#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Feb 13 11:18:50 2020

@author: dberke

A script to read in data on transition offsets vs. several stellar parameters
from a database, and perform multi-component fitting to it.
"""

import argparse
import csv
from inspect import signature
import os
from pathlib import Path
import pickle
from pprint import pprint
from time import sleep, time
import sys

from adjustText import adjust_text
import h5py
import hickle
from matplotlib.gridspec import GridSpec
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import numpy.ma as ma
from scipy.optimize import curve_fit
from tqdm import tqdm
import unyt as u

import varconlib as vcl
import varconlib.fitting.fitting as fit

# Define style parameters to use for stellar parameter plots.
style_pre = {'color': 'Chocolate',
             'ecolor_thick': 'DarkOrange',
             'ecolor_thin': 'BurlyWood'}
style_post = {'color': 'DodgerBlue',
              'ecolor_thick': 'CornFlowerBlue',
              'ecolor_thin': 'LightSkyBlue'}
style_ref = {'color': 'DarkGreen',
             'ecolor_thick': 'ForestGreen',
             'ecolor_thin': 'DarkSeaGreen'}
style_markers = {'markeredgecolor': 'Black',
                 'markeredgewidth': 1,
                 'alpha': 0.7,
                 'markersize': 4}
style_caps = {'capsize_thin': 4,
              'capsize_thick': 7,
              'linewidth_thin': 2,
              'linewidth_thick': 3,
              'cap_thin': 1.5,
              'cap_thick': 2.5}

def create_comparison_figure(ylims=None, fit_target='transitions',
                             temp_lims=(5300 * u.K, 6200 * u.K),
                             mtl_lims=(-0.75, 0.4),
                             mag_lims=(4, 5.8),
                             logg_lims=(4.1, 4.6)):
    """Create and returns a figure with pre-set subplots.

    This function creates the background figure and subplots for use with the
    --compare-stellar-parameter-* flags.

    Optional
    ----------
    ylims : 2-tuple of floats or ints
        A tuple of length 2 containing the upper and lower limits of the
        subplots in the figure.
    fit_target : str, ['transitions', 'pairs']
        A string denoting whether these plots are for transitions or pairs.
    temp_lims : 2-tuple of floats or ints (optional dimensions of temperature)
        A tuple of length containing upper and lower limits for the x-axis of
        the temperature subplot.
    mtl_lims : 2-tuple of floats or ints
        A tuple of length containing upper and lower limits for the x-axis of
        the metallicity subplot.
    mag_lims : 2-tuple of floats or ints
        A tuple of length containing upper and lower limits for the x-axis of
        the absolute magnitude subplot.
    logg_lims : 2-tuple of floats or ints
        A tuple of length containing upper and lower limits for the x-axis of
        the log(g) subplot.

    Returns
    -------
    tuple
        A tuple containing the figure itself and the various axes of the
        subplots within it.

    """

    comp_fig = plt.figure(figsize=(12, 8), tight_layout=True)
    gs = GridSpec(ncols=4, nrows=2, figure=comp_fig,
                  width_ratios=(5, 5, 5, 3))

    temp_ax_pre = comp_fig.add_subplot(gs[0, 0])
    temp_ax_post = comp_fig.add_subplot(gs[1, 0],
                                        sharex=temp_ax_pre,
                                        sharey=temp_ax_pre)
    mtl_ax_pre = comp_fig.add_subplot(gs[0, 1],
                                      sharey=temp_ax_pre)
    mtl_ax_post = comp_fig.add_subplot(gs[1, 1],
                                       sharex=mtl_ax_pre,
                                       sharey=mtl_ax_pre)
    logg_ax_pre = comp_fig.add_subplot(gs[0, 2],
                                       sharey=temp_ax_pre)
    logg_ax_post = comp_fig.add_subplot(gs[1, 2],
                                        sharex=logg_ax_pre,
                                        sharey=logg_ax_pre)
    hist_ax_pre = comp_fig.add_subplot(gs[0, 3],
                                       sharey=temp_ax_pre)
    hist_ax_post = comp_fig.add_subplot(gs[1, 3],
                                        sharex=hist_ax_pre,
                                        sharey=hist_ax_pre)

    all_axes = (temp_ax_pre, temp_ax_post, mtl_ax_pre, mtl_ax_post,
                logg_ax_pre, logg_ax_post, hist_ax_pre, hist_ax_post)
    # Set the plot limits here. The y-limits for temp_ax1 are
    # used for all subplots.
    if ylims is not None:
        temp_ax_pre.set_ylim(bottom=ylims[0],
                             top=ylims[1])
    temp_ax_pre.set_xlim(left=temp_lims[0],
                         right=temp_lims[1])
    mtl_ax_pre.set_xlim(left=mtl_lims[0],
                        right=mtl_lims[1])
    logg_ax_pre.set_xlim(left=logg_lims[0],
                         right=logg_lims[1])

    # Axis styles for all subplots.
    for ax in all_axes:
        if not args.full_range:
            ax.yaxis.set_major_locator(ticker.MultipleLocator(
                                      base=100))
            ax.yaxis.set_minor_locator(ticker.MultipleLocator(
                                      base=50))
        else:
            ax.yaxis.set_major_locator(ticker.AutoLocator())
            ax.yaxis.set_minor_locator(ticker.AutoMinorLocator())
        ax.axhline(y=0, color='Black', linestyle='--')
        ax.yaxis.grid(which='major', color='Gray',
                      linestyle='--', alpha=0.5)
        ax.yaxis.grid(which='minor', color='Gray',
                      linestyle=':', alpha=0.5)
        if ax not in (hist_ax_pre, hist_ax_post):
            ax.xaxis.grid(which='major', color='Gray',
                          linestyle='--', alpha=0.65)
        ax.tick_params(labelsize=14)

    for ax in (temp_ax_pre, temp_ax_post):
        ax.set_xlabel('Temperature (K)', size=15)
        ax.xaxis.set_major_locator(ticker.MultipleLocator(base=200))
        ax.xaxis.set_minor_locator(ticker.MultipleLocator(base=100))
    for ax in (mtl_ax_pre, mtl_ax_post):
        ax.set_xlabel('Metallicity [Fe/H]', size=15)
        ax.xaxis.set_major_locator(ticker.MultipleLocator(base=0.2))
        ax.xaxis.set_minor_locator(ticker.MultipleLocator(base=0.1))
    for ax in (logg_ax_pre, logg_ax_post):
        ax.set_xlabel(r'log $g$ $(\mathrm{cm}/\mathrm{s}^2)$', size=15)
        ax.xaxis.set_major_locator(ticker.MultipleLocator(base=0.1))
        ax.xaxis.set_minor_locator(ticker.MultipleLocator(base=0.05))

    # Just label the left-most two subplots' y-axes.
    for ax, era in zip((temp_ax_pre, temp_ax_post),
                       ('Pre', 'Post')):
        if fit_target == 'transitions':
            ax.set_ylabel(f'{era}-fiber change offset (m/s)')
        elif fit_target == 'pairs':
            ax.set_ylabel(f'{era}-fiber change separation (m/s)')
        else:
            raise RuntimeError(f'Unallowed value for fit_target: {fit_target}')

    axes_dict = {'temp_pre': temp_ax_pre, 'temp_post': temp_ax_post,
                 'mtl_pre': mtl_ax_pre, 'mtl_post': mtl_ax_post,
                 'logg_pre': logg_ax_pre, 'logg_post': logg_ax_post,
                 'hist_pre': hist_ax_pre, 'hist_post': hist_ax_post}

    return comp_fig, axes_dict

def plot_data_points(axis, x_pos, y_pos, thick_err, thin_err, era=None,
                     ref=False):
    """Plot a data point for a star.

    Parameters
    ----------
    axis : `matplotlib.axes.Axes`
        An axes to plot the data on.
    x_pos : iterable of floats or `unyt.unyt_quantity`
        The x-positions of the points to plot.
    y_pos : iterable of floats or `unyt.unyt_quantity`
        The y-positions of the points to plot. The length must match the length
        of `x_pos`.
    thick_err : iterable of floats or `unyt.unyt_quantity` or None
        The values of the thick error bars to plot. The length must match the
        length of `x_pos`.
    thin_err : iterable of floats or `unyt.unyt_quantity` or None
        The values of the thin error bars to plot. The length must match the
        length of `x_pos`.
    era : string, ['pre', 'post'], or None, Default : None
        Whether the time period of the plot is pre- or post-fiber
        change. The only allowed string values are 'pre' and 'post'. Controls
        color of the points. If `ref` is *True*, the value of `era` is
        ignored, and can be left unspecified, otherwise it needs a
        value to be given.
    ref : bool, Default : False
        Whether this data point is for the reference star. If *True*,
        will use a special separate color scheme.

    Returns
    -------
    None.

    """

    if ref:
        params = style_ref
    elif era == 'pre':
        params = style_pre
    elif era == 'post':
        params = style_post
    else:
        raise ValueError("Keyword 'era' received an unknown value"
                         f" (valid values are 'pre' & 'post'): {era}")

    if thin_err is not None:
        axis.errorbar(x=x_pos, y=y_pos,
                      yerr=thin_err, linestyle='',
                      marker='', capsize=style_caps['capsize_thin'],
                      color=params['color'],
                      ecolor=params['ecolor_thin'],
                      elinewidth=style_caps['linewidth_thin'],
                      capthick=style_caps['cap_thin'])
    if thick_err is not None:
        axis.errorbar(x=x_pos, y=y_pos,
                      yerr=thick_err, linestyle='',
                      marker='o', markersize=style_markers['markersize'],
                      markeredgewidth=style_markers['markeredgewidth'],
                      markeredgecolor=style_markers['markeredgecolor'],
                      alpha=style_markers['alpha'],
                      capsize=style_caps['capsize_thick'],
                      color=params['color'],
                      ecolor=params['ecolor_thick'],
                      elinewidth=style_caps['linewidth_thick'],
                      capthick=style_caps['cap_thick'])

def find_star(star_params, stellar_params, star_names):
    """Return the name of the star which matches the given parameters.

    Parameters
    ----------
    star_params : iterable
        An iterable (in practice, an array slice of length 3) which contains
        values of stellar parameters, temperature in position 0, metallicity in
        position 1, and magnitude in position 2.
    stellar_params : array-like
        A 3xN array of stellar parameters, with the same order as `star_params`.
    star_names : `bidict.bidict`
        A `bidict` containing star names in strings as keys, and index numbers
        as ints for the values.

    Returns
    -------
    str
        The name of the star associated with the given parameters.

    """

    for i in range(len(stellar_params[0])):
        if np.all(np.isclose(star_params, stellar_params[:, i])):
            return star_names[i]

def main():
    """Run the main routine of the script."""

    # Define the limits to plot in the various stellar parameters.
    temp_lims = (5400, 6300) * u.K
    mtl_lims = (-0.75, 0.45)
    # mag_lims = (4, 5.8)
    logg_lims = (4.1, 4.6)

    # Define the model to use:
    if args.constant:
        model_func = fit.constant_model
    elif args.linear:
        model_func = fit.linear_model
    elif args.quadratic:
        model_func = fit.quadratic_model
    elif args.cubic:
        model_func = fit.cubic_model
    elif args.cross_term:
        model_func = fit.cross_term_model
    elif args.quadratic_cross_term:
        model_func = fit.quad_cross_term_model
    elif args.quad_cross_term:
        model_func = fit.quad_cross_term_model

    model_name = '_'.join(model_func.__name__.split('_')[:-1])

    if args.transitions:
        tqdm.write('Unpickling transitions list.')
        with open(vcl.final_selection_file, 'r+b') as f:
            transitions_list = pickle.load(f)
        vprint(f'Found {len(transitions_list)} transitions.')
    elif args.pairs:
        tqdm.write('Unpickling pairs list.')
        with open(vcl.final_pair_selection_file, 'r+b') as f:
            pairs_list = pickle.load(f)
        vprint(f'Found {len(pairs_list)} pairs in the list.')

    db_file = vcl.databases_dir / 'stellar_db_uncorrected.hdf5'
    if not db_file.exists():
        raise FileNotFoundError('The given stellar database does not exist:'
                                f' {db_file}')

    # Load data from HDF5 database file.
    tqdm.write('Reading data from stellar database file...')
    if args.transitions:
        star_transition_offsets = u.unyt_array.from_hdf5(
                db_file, dataset_name='star_transition_offsets')
        star_transition_offsets_EotWM = u.unyt_array.from_hdf5(
                db_file, dataset_name='star_transition_offsets_EotWM')
        star_transition_offsets_EotM = u.unyt_array.from_hdf5(
                db_file, dataset_name='star_transition_offsets_EotM')
        # star_transition_offsets_stds = u.unyt_array.from_hdf5(
        #         db_file, dataset_name='star_standard_deviations')
    elif args.pairs:
        star_pair_separations = u.unyt_array.from_hdf5(
               db_file, dataset_name='star_pair_separations')
        star_pair_separations_EotWM = u.unyt_array.from_hdf5(
                db_file, dataset_name='star_pair_separations_EotWM')
        star_pair_separations_EotM = u.unyt_array.from_hdf5(
                db_file, dataset_name='star_pair_separations_EotM')
    star_temperatures = u.unyt_array.from_hdf5(
            db_file, dataset_name='star_temperatures')

    with h5py.File(db_file, mode='r') as f:

        star_metallicities = hickle.load(f, path='/star_metallicities')
        # star_magnitudes = hickle.load(f, path='/star_magnitudes')
        star_gravities = hickle.load(f, path='/star_gravities')
        transition_column_dict = hickle.load(f, path='/transition_column_index')
        pair_column_dict = hickle.load(f, path='/pair_column_index')

        star_names = hickle.load(f, path='/star_row_index')

    # Handle various fitting and plotting setup:
    eras = {'pre': 0, 'post': 1}
    param_dict = {'temp': 0, 'mtl': 1, 'logg': 2}

    # Create lists to store information about each fit in:
    index_nums = []
    chi_squareds_pre, sigmas_pre, sigma_sys_pre = [], [], []
    chi_squareds_post, sigmas_post, sigma_sys_post = [], [], []
    index_num = 0

    # Figure out how many parameters the model function takes, so we know how
    # many to dynamically give it later. Subtract 1 for the parameter which
    # takes the stellar parameters.
    params_list = [0 for i in range(len(signature(model_func).parameters)-1)]

    # Define the folder to put plots in.
    output_dir = vcl.output_dir
    if args.transitions:
        fit_target = 'transitions'
    elif args.pairs:
        fit_target = 'pairs'
    plots_folder = output_dir /\
        f'stellar_parameter_fits_{fit_target}_{args.sigma}sigma/{model_name}'
    vprint(f'Creating plots in {plots_folder}')
    if not plots_folder.exists():
        os.makedirs(plots_folder)

    # Create a dictionary of fit coefficients assigned to each transition's
    # label
    coefficients_dict = {}
    covariance_dict = {}
    sigmas_dict = {}
    sigma_sys_dict = {}

    if args.transitions:
        tqdm.write('Creating plots for each transition...')
        for transition in tqdm(transitions_list):
            for order_num in transition.ordersToFitIn:
                index_nums.append(index_num)
                index_num += 1
                label = '_'.join([transition.label, str(order_num)])
                vprint(20 * '-')
                vprint(f'Analyzing {label}...')

                # The column number to use for this transition:
                col = transition_column_dict[label]
                ylimits = (-300 * u.m / u.s,
                           300 * u.m / u.s) if not args.full_range else None

                comp_fig, axes_dict = create_comparison_figure(
                                ylims=ylimits,
                                fit_target='transitions',
                                temp_lims=temp_lims,
                                mtl_lims=mtl_lims,
                                logg_lims=logg_lims)

                for time in eras.keys():

                    vprint(20 * '=')
                    vprint(f'Working on {time}-change era.')
                    mean = np.nanmean(star_transition_offsets[eras[time],
                                      :, col])

                    # First, create a masked version to catch any missing
                    # entries:
                    m_offsets = ma.masked_invalid(star_transition_offsets[
                                eras[time], :, col])
                    total_stars = ma.count(m_offsets)
                    vprint(f'Found {total_stars} stars with data.')
                    m_offsets = m_offsets.reshape([len(m_offsets), 1])
                    # Then create a new array from the non-masked data:
                    offsets = u.unyt_array(m_offsets[~m_offsets.mask],
                                           units=u.m/u.s)
                    vprint(f'Median of offsets is {np.nanmedian(offsets)}')

    #                m_stds = ma.masked_invalid(star_transition_offsets_stds[
    #                            eras[time], :, col])
    #                m_stds = m_stds.reshape([len(m_stds), 1])
    #                stds = u.unyt_array(m_stds[~m_stds.mask],
    #                                    units=u.m/u.s)

                    m_eotwms = ma.masked_invalid(star_transition_offsets_EotWM[
                            eras[time], :, col])
                    m_eotwms = m_eotwms.reshape([len(m_eotwms), 1])
                    eotwms = u.unyt_array(m_eotwms[~m_offsets.mask],
                                          units=u.m/u.s)

                    m_eotms = ma.masked_invalid(star_transition_offsets_EotM[
                            eras[time], :, col])
                    m_eotms = m_eotms.reshape([len(m_eotms), 1])
                    # Use the same mask as for the offsets.
                    eotms = u.unyt_array(m_eotms[~m_offsets.mask],
                                         units=u.m/u.s)
                    # Create an error array which uses the greater of the error
                    # on the mean or the error on the weighted mean.
                    err_array = ma.array(np.maximum(eotwms, eotms).value)

                    vprint(f'Mean is {np.mean(offsets)}')
                    weighted_mean = np.average(offsets, weights=err_array**-2)
                    vprint(f'Weighted mean is {weighted_mean}')

                    # Mask the various stellar parameter arrays with the same
                    # mask so that everything stays in sync.
                    temperatures = ma.masked_array(star_temperatures)
                    temps = temperatures[~m_offsets.mask]
                    metallicities = ma.masked_array(star_metallicities)
                    metals = metallicities[~m_offsets.mask]
                    # magnitudes = ma.masked_array(star_magnitudes)
                    # mags = magnitudes[~m_offsets.mask]
                    gravities = ma.masked_array(star_gravities)
                    loggs = gravities[~m_offsets.mask]

                    stars = ma.masked_array([key for key in
                                             star_names.keys()]).reshape(
                                                 len(star_names.keys()), 1)
                    names = stars[~m_offsets.mask]

                    # Stack the stellar parameters into vertical slices
                    # for passing to model functions.
                    x_data = ma.array(np.stack((temps, metals, loggs), axis=0))

                    # Create the parameter list for this run of fitting.
                    params_list[0] = float(mean)

                    beta0 = tuple(params_list)
                    vprint(beta0)

                    results = fit.find_sys_scatter(model_func,
                                                   x_data,
                                                   ma.array(offsets.value),
                                                   err_array, beta0,
                                                   n_sigma=args.sigma,
                                                   tolerance=0.001,
                                                   verbose=args.verbose)

                    mask = results['mask_list'][-1]
                    residuals = ma.array(results['residuals'], mask=mask)
                    x_data.mask = mask
                    err_array.mask = mask

                    # for item1, item2 in zip(residuals, ma.getdata(residuals)):
                    #     print(f'{item1:10.3f}   {item2:10.3f}')

                    chi_squared_nu = results['chi_squared_list'][-1]
                    sys_err = results['sys_err_list'][-1] * u.m / u.s

                    vprint(f'Terminated with sys_err = {sys_err}')
                    vprint(f'Finished {label}_{time} in'
                           f' {len(results["sys_err_list"])} steps.')
                    # Add the optimized parameters and covariances to the
                    # dictionary. Make sure we separate them by time period.
                    coefficients_dict[label + '_' + time] = results['popt']
                    covariance_dict[label + '_' + time] = results['pcov']

                    sigma = np.nanstd(residuals) * u.m/u.s

                    sigmas_dict[label + '_' + time] = sigma
                    sigma_sys_dict[label + '_' + time] = sys_err

                    if time == 'pre':
                        chi_squareds_pre.append(chi_squared_nu)
                        sigmas_pre.append(sigma.value)
                        sigma_sys_pre.append(sys_err.value)
                    else:
                        chi_squareds_post.append(chi_squared_nu)
                        sigmas_post.append(sigma.value)
                        sigma_sys_post.append(sys_err.value)

                    for plot_type, lims in zip(('temp', 'mtl', 'logg'),
                                               (temp_lims, mtl_lims,
                                                logg_lims)):
                        ax = axes_dict[f'{plot_type}_{time}']
                        plot_data_points(ax,
                                         ma.compressed(x_data[
                                             param_dict[plot_type]]),
                                         ma.compressed(residuals),
                                         thick_err=ma.compressed(err_array),
                                         # thin_err=iter_err_array,
                                         thin_err=None,
                                         era=time)
                        if args.label_outliers:
                            # Find outliers more than 3 sigma away from zero so
                            # we can label them.
                            labels = []
                            for x, y, e in zip(range(len(
                                    x_data[param_dict[plot_type]])), residuals,
                                    err_array):
                                sig_lim = args.sigma * e
                                if abs(y) > sig_lim:
                                    star_name = find_star(x_data[:, x],
                                                          x_data, names)

                                    labels.append(ax.text(
                                        x_data[param_dict[plot_type], x],
                                        y, star_name,
                                        horizontalalignment='left',
                                        verticalalignment='top',
                                        size=8, weight='bold', color='Red'))
                            # print(labels)
                            adjust_text(labels,
                                        ax=ax,
                                        only_move={'points': 'y',
                                                   'text': 'xy',
                                                   'objects': 'xy'},
                                        arrowprops=dict(arrowstyle='-',
                                                        color='OliveDrab'),
                                        autoalign=True,
                                        lim=1000, fontsize=9)

                        points = residuals.count()
                        outliers = total_stars - points
                        ax.annotate(f'Blendedness: {transition.blendedness}\n'
                                    f'Stars: {points}\n'
                                    f'Outliers: {outliers}',
                                    (0.01, 0.99),
                                    xycoords='axes fraction',
                                    verticalalignment='top')
                        ax.annotate(fr'$\chi^2_\nu$: {chi_squared_nu:.4f}'
                                    '\n'
                                    fr'$\sigma$: {sigma:.2f}'
                                    '\n'
                                    r'$\sigma_\mathrm{sys}$:'
                                    f' {sys_err:.2f}',
                                    (0.99, 0.99),
                                    xycoords='axes fraction',
                                    horizontalalignment='right',
                                    verticalalignment='top')
                        data = np.array(ma.masked_invalid(
                            residuals).compressed())
                        axes_dict[f'hist_{time}'].hist(data,
                                                       bins='fd',
                                                       color='Black',
                                                       histtype='step',
                                                       orientation='horizontal')

                file_name = plots_folder / f'{label}_{model_name}.png'
                vprint(f'Saving file {label}.png')
                vprint('\n')

                comp_fig.savefig(str(file_name))
                plt.close('all')

    elif args.pairs:
        tqdm.write('Creating plots for each pair...')
        for pair in tqdm(pairs_list):
            for order_num in pair.ordersToMeasureIn:
                index_nums.append(index_num)
                index_num += 1
                label = '_'.join([pair.label, str(order_num)])
                vprint(20 * '-')
                vprint(f'Analyzing {label}...')

                # The column number to use for this transition:
                col = pair_column_dict[label]
                ylimits = (-300 * u.m / u.s,
                           300 * u.m / u.s) if not args.full_range else None

                comp_fig, axes_dict = create_comparison_figure(
                                ylims=ylimits,
                                fit_target='pairs',
                                temp_lims=temp_lims,
                                mtl_lims=mtl_lims,
                                logg_lims=logg_lims)

                for time in eras.keys():

                    vprint(20 * '=')
                    vprint(f'Working on {time}-change era.')
                    mean = np.nanmean(star_pair_separations[eras[time],
                                      :, col])

                    # First, create a masked version to catch any missing
                    # entries:
                    m_seps = ma.masked_invalid(star_pair_separations[
                                eras[time], :, col])
                    total_stars = ma.count(m_seps)
                    vprint(f'Found {total_stars} stars with data.')
                    m_seps = m_seps.reshape([len(m_seps), 1])
                    # Then create a new array from the non-masked data:
                    separations = u.unyt_array(m_seps[~m_seps.mask],
                                               units=u.m/u.s)
                    vprint('Median of separations is'
                           f' {np.nanmedian(separations)}')

                    m_eotwms = ma.masked_invalid(star_pair_separations_EotWM[
                            eras[time], :, col])
                    m_eotwms = m_eotwms.reshape([len(m_eotwms), 1])
                    eotwms = u.unyt_array(m_eotwms[~m_seps.mask],
                                          units=u.m/u.s)

                    m_eotms = ma.masked_invalid(star_pair_separations_EotM[
                            eras[time], :, col])
                    m_eotms = m_eotms.reshape([len(m_eotms), 1])
                    # Use the same mask as for the offsets.
                    eotms = u.unyt_array(m_eotms[~m_seps.mask],
                                         units=u.m/u.s)
                    # Create an error array which uses the greater of the error
                    # on the mean or the error on the weighted mean.
                    err_array = ma.array(np.maximum(eotwms, eotms).value)

                    vprint(f'Mean is {np.mean(separations)}')
                    weighted_mean = np.average(separations,
                                               weights=err_array**-2)
                    vprint(f'Weighted mean is {weighted_mean}')

                    # Mask the various stellar parameter arrays with the same
                    # mask so that everything stays in sync.
                    temperatures = ma.masked_array(star_temperatures)
                    temps = temperatures[~m_seps.mask]
                    metallicities = ma.masked_array(star_metallicities)
                    metals = metallicities[~m_seps.mask]
                    gravities = ma.masked_array(star_gravities)
                    loggs = gravities[~m_seps.mask]

                    stars = ma.masked_array([key for key in
                                             star_names.keys()]).reshape(
                                                 len(star_names.keys()), 1)
                    names = stars[~m_seps.mask]

                    # Stack the stellar parameters into vertical slices
                    # for passing to model functions.
                    x_data = ma.array(np.stack((temps, metals, loggs), axis=0))

                    # Create the parameter list for this run of fitting.
                    params_list[0] = float(mean)

                    beta0 = tuple(params_list)
                    vprint(beta0)

                    results = fit.find_sys_scatter(model_func,
                                                   x_data,
                                                   ma.array(separations.value),
                                                   err_array, beta0,
                                                   n_sigma=args.sigma,
                                                   tolerance=0.001,
                                                   verbose=args.verbose)

                    mask = results['mask_list'][-1]
                    residuals = ma.array(results['residuals'], mask=mask)
                    x_data.mask = mask
                    err_array.mask = mask

                    chi_squared_nu = results['chi_squared_list'][-1]
                    sys_err = results['sys_err_list'][-1] * u.m / u.s

                    vprint(f'Terminated with sys_err = {sys_err}')
                    vprint(f'Finished {label}_{time} in'
                           f' {len(results["sys_err_list"])} steps.')
                    # Add the optimized parameters and covariances to the
                    # dictionary. Make sure we separate them by time period.
                    coefficients_dict[label + '_' + time] = results['popt']
                    covariance_dict[label + '_' + time] = results['pcov']

                    sigma = np.nanstd(residuals) * u.m/u.s

                    sigmas_dict[label + '_' + time] = sigma
                    sigma_sys_dict[label + '_' + time] = sys_err

                    if time == 'pre':
                        chi_squareds_pre.append(chi_squared_nu)
                        sigmas_pre.append(sigma.value)
                        sigma_sys_pre.append(sys_err.value)
                    else:
                        chi_squareds_post.append(chi_squared_nu)
                        sigmas_post.append(sigma.value)
                        sigma_sys_post.append(sys_err.value)

                    for plot_type, lims in zip(('temp', 'mtl', 'logg'),
                                               (temp_lims, mtl_lims,
                                                logg_lims)):
                        ax = axes_dict[f'{plot_type}_{time}']
                        ax.tick_params(labelsize=14)
                        plot_data_points(
                            ax,
                            ma.compressed(x_data[param_dict[plot_type]]),
                            ma.compressed(residuals),
                            thick_err=ma.compressed(err_array),
                            # thin_err=None,
                            thin_err=np.sqrt(ma.compressed(err_array) ** 2 +
                                             sys_err.value ** 2),
                            era=time)
                        if args.label_outliers:
                            # Find outliers more than 3 sigma away from zero so
                            # we can label them.
                            labels = []
                            for x, y, e in zip(range(len(
                                    x_data[param_dict[plot_type]])), residuals,
                                    err_array):
                                sig_lim = args.sigma * e
                                if abs(y) > sig_lim:
                                    star_name = find_star(x_data[:, x],
                                                          x_data, names)

                                    labels.append(ax.text(
                                        x_data[param_dict[plot_type], x],
                                        y, star_name,
                                        horizontalalignment='left',
                                        verticalalignment='top',
                                        size=8, weight='bold', color='Red'))
                            # print(labels)
                            adjust_text(labels,
                                        ax=ax,
                                        only_move={'points': 'y',
                                                   'text': 'xy',
                                                   'objects': 'xy'},
                                        arrowprops=dict(arrowstyle='-',
                                                        color='OliveDrab'),
                                        autoalign=True,
                                        lim=1000, fontsize=9)

                        points = residuals.count()
                        outliers = total_stars - points
                        ax.annotate(f'Blend tuple: {pair.blendTuple}\n'
                                    f'Stars: {points}\n'
                                    f'Outliers: {outliers}',
                                    (0.01, 0.99),
                                    xycoords='axes fraction',
                                    verticalalignment='top')
                        ax.annotate(fr'$\chi^2_\nu$: {chi_squared_nu:.4f}'
                                    '\n'
                                    fr'$\sigma$: {sigma:.2f}'
                                    '\n'
                                    r'$\sigma_\mathrm{sys}$:'
                                    f' {sys_err:.2f}',
                                    (0.99, 0.99),
                                    xycoords='axes fraction',
                                    horizontalalignment='right',
                                    verticalalignment='top')
                        data = np.array(ma.masked_invalid(
                            residuals).compressed())
                        axes_dict[f'hist_{time}'].hist(data,
                                                       bins='fd',
                                                       color='Black',
                                                       histtype='step',
                                                       orientation='horizontal')

                file_name = plots_folder / f'{label}_{model_name}.png'
                vprint(f'Saving file {label}.png')
                vprint('\n')

                comp_fig.savefig(str(file_name))
                plt.close('all')

    # Save metadata from this run's fits to CSV:
    csv_file = plots_folder / f'{model_name}_{fit_target}_fit_results.csv'

    with open(csv_file, 'w', newline='') as f:
        datawriter = csv.writer(f)
        header = ('#index', 'chi_squared_pre', 'sigma_pre', 'sigma_sys_pre',
                  'chi_squared_post', 'sigma_post', 'sigma_sys_post')
        datawriter.writerow(header)
        for row in zip(index_nums, chi_squareds_pre, sigmas_pre, sigma_sys_pre,
                       chi_squareds_post, sigmas_post, sigma_sys_post):
            datawriter.writerow(row)

    # Save the function used and the parameters found for each transition/pair
    # to an HDF5 file for use in other scripts.
    output_dir = output_dir / 'fit_params'
    hdf5_file = output_dir /\
        f'{model_name}_{fit_target}_{args.sigma:.1f}sigma_params.hdf5'
    if not hdf5_file.parent.exists():
        os.mkdir(hdf5_file.parent)

    vprint(f'Writing HDF5 file with fit parameters at {hdf5_file}')
    if hdf5_file.exists():
        os.unlink(hdf5_file)
    with h5py.File(hdf5_file, mode='a') as f:
        f.attrs['type'] = 'A file containing a fitting function and the' +\
                          ' parameters for it for each transition or pair' +\
                          'in /coeffs_dict'
        hickle.dump(model_func, f, path='/fitting_function')
        hickle.dump(coefficients_dict, f, path='/coeffs_dict')
        hickle.dump(covariance_dict, f, path='/covariance_dict')
        hickle.dump(sigmas_dict, f, path='/sigmas_dict')
        hickle.dump(sigma_sys_dict, f, path='/sigma_sys_dict')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Use stored data from stars'
                                     ' to fit transition offsets to stellar'
                                     ' parameters.')
    parser.add_argument('--full-range', action='store_true',
                        help='Plot the full range of values instead of'
                        ' restricting to a  fixed range.')
    parser.add_argument('--label-outliers', action='store_true',
                        help='Label the points which are more than'
                        ' 3 sigma away from the mean.')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Print out more information about the script.')

    parser.add_argument('--sigma', action='store', type=float, default=2.5,
                        help='The number to use (in standard deviations)'
                        ' beyond which to consider a data point an outlier.')

    func = parser.add_mutually_exclusive_group(required=True)
    func.add_argument('--constant', action='store_true',
                      help='Use a constant function.')
    func.add_argument('--linear', action='store_true',
                      help='Use a function linear in all three variables.')
    func.add_argument('--quadratic', action='store_true',
                      help='Use a function quadratic in all three variables.')
    func.add_argument('--cubic', action='store_true',
                      help='Use a cubic function for all three variables.')
    func.add_argument('--quartic', action='store_true',
                      help='Use a quartic function for all three variables.')
    func.add_argument('--quintic', action='store_true',
                      help='Use a quintic function for all three variables.')
    func.add_argument('--cross-term', action='store_true',
                      help='Use a linear model with cross term ([Fe/H]/Teff).')
    func.add_argument('--quadratic-cross-term', action='store_true',
                      help='Use a quadratic model with cross terms between'
                      ' metallicity and temperature.')
    func.add_argument('--quadratic-magnitude', action='store_true',
                      help='Use a cross term with quadratic magnitude.')
    func.add_argument('--quad-cross-terms', action='store_true',
                      help='Use a quadratic model with full cross terms.')

    fit_target = parser.add_mutually_exclusive_group(required=True)
    fit_target.add_argument('-T', '--transitions', action='store_true',
                            help='Fit individual transitions.')
    fit_target.add_argument('-P', '--pairs', action='store_true',
                            help='Fit pairs.')

    args = parser.parse_args()

    vprint = vcl.verbose_print(args.verbose)

    start_time = time()

    main()

    duration = time() - start_time
    print(f'Finished in {duration:.2f} seconds.')


"""
##################################################################################################
# Copyright Info :    Copyright (c) Davar Lab @ Hikvision Research Institute. All rights reserved.
# Filename       :    affine_transformation.py
# Abstract       :    Implementations of the affine transformation

# Current Version:    1.0.0
# Date           :    2021-03-07
##################################################################################################
"""
import logging

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from mmcv.runner import load_checkpoint

from ..builder import TRANSFORMATIONS

@TRANSFORMATIONS.register_module()
class Affine_SpatialTransformer(nn.Module):
    """ Rectification Network of RARE, namely Affine-based STN [1]

    Ref: [1] Spatial Transformer Network. NIPS-2016

    Usage:
    transformation=dict(
         type='Affine_SpatialTransformer',
         I_size=(32, 100),
         I_r_size=(32, 100),
         I_channel_num=1),

    """

    def __init__(self,
                 I_size,
                 I_r_size,
                 I_channel_num=1,
                 fix_transformation=False):
        """
        Spatial Transformation Network based on Affine

        Args:
            I_size (tuple): size of input images
            I_r_size (tuple): size of rectified images
            I_channel_num (int): the number of channels of the input image I
            fix_transformation (bool): if fix the parameters of the transformation during training

        """
        super(Affine_SpatialTransformer, self).__init__()
        self.I_size = I_size
        self.I_r_size = I_r_size
        self.I_channel_num = I_channel_num
        self.LocalizationNetwork = AffineLocalizationNetwork(self.I_channel_num, self.I_size)

        # if set STN parameters fixed, the weights and bias of all layers are set to False
        if fix_transformation:
            for p in self.LocalizationNetwork.parameters():
                p.requires_grad = False

    def init_weights(self, pretrained=None):
        """

        Args:
            pretrained (str): save path of pretrained model

        """
        if isinstance(pretrained, str):
            logger = logging.getLogger()
            logger.info("Affine_SpatialTransformer:")
            load_checkpoint(self, pretrained,
                            strict=False, logger=logger)

    def forward(self, batch_I):
        """

        Args:
            batch_I (tensor): batch of input images[batch_size x I_channel_num x I_r_height x I_r_width]

        Returns:
            torch.Tensor: rectified image [batch_size x I_channel_num x I_r_height x I_r_width]

        """
        xy_batch_C_prime = self.LocalizationNetwork(batch_I).view(-1, 2, 3)

        row_rev = torch.cat([xy_batch_C_prime[:, 1], xy_batch_C_prime[:, 0]], 1).view(-1, 2, 3).permute(0, 2, 1)
        colomn_rev = torch.cat([row_rev[:, 1], row_rev[:, 0]], 1)
        colomn_rev = torch.cat([colomn_rev, row_rev[:, 2]], 1).view(-1, 3, 2).permute(0, 2, 1)
        batch_C_prime = colomn_rev

        build_P_prime = F.affine_grid(batch_C_prime, torch.Size((batch_C_prime.size(0),
                                                                 self.I_channel_num,
                                                                 self.I_r_size[0],
                                                                 self.I_r_size[1])))

        batch_I_r = F.grid_sample(batch_I, build_P_prime, padding_mode='border')

        return batch_I_r

class AffineLocalizationNetwork(nn.Module):
    """
    Localization Network of STN,
    which predicts 6 params of Affine Transformation

    """

    def __init__(self, I_channel_num, I_size):
        """
        Args:
            I_channel_num (int): channel number of input
            I_size (tuple): input image size
        """
        super(AffineLocalizationNetwork, self).__init__()
        self.I_channel_num = I_channel_num
        self.I_size = I_size
        self.loc_conv1_channel = nn.Conv2d(in_channels=self.I_channel_num,
                                           out_channels=48, kernel_size=3,
                                           stride=1, padding=1)
        self.loc_relu1 = nn.ReLU(inplace=True)
        self.loc_pool1 = nn.MaxPool2d(2, 2)
        self.loc_conv2_channel = nn.Conv2d(48, 48,
                                           3, 1, 1)
        self.loc_relu2_new = nn.ReLU(inplace=True)
        self.loc_pool2_new = nn.MaxPool2d(2, 2)
        self.loc_conv3_channel = nn.Conv2d(48, 64,
                                           3, 1, 1)
        self.loc_relu3_new = nn.ReLU(inplace=True)
        self.loc_pool3_new = nn.MaxPool2d(2, 2)
        self.loc_conv4_channel_64 = nn.Conv2d(64, 64,
                                              3, 1, 1)
        self.loc_relu4 = nn.ReLU(inplace=True)
        self.loc_pool4 = nn.MaxPool2d(2, 2)

        self.loc_relu5 = nn.ReLU(inplace=True)
        self.loc_reg = nn.Linear(768, 6)

        # Init fc2 in LocalizationNetwork
        self.loc_reg.weight.data.fill_(0)
        initial_bias = np.array([1, 0, 0, 0, 1, 0])
        self.loc_reg.bias.data = torch.from_numpy(initial_bias).float().view(-1)

        # for item in self.no_update:
        #    item.weight.requires_grad = False
        #    item.bias.requires_grad = False

    def forward(self, x):
        """
        Args:
            x (tensor): input image feature maps [batch_size x I_channel_num x I_height x I_width]

        Returns:
            torch.Tensor: Predicted coordinates of fiducial points for input batch [batch_size x F x 2]
        """
        batch_size = x.size(0)
        x = self.loc_conv1_channel(x)
        x = self.loc_relu1(x)
        x = self.loc_pool1(x)
        x = self.loc_conv2_channel(x)
        x = self.loc_relu2_new(x)
        x = self.loc_pool2_new(x)
        x = self.loc_conv3_channel(x)
        x = self.loc_relu3_new(x)
        x = self.loc_pool3_new(x)
        x = self.loc_conv4_channel_64(x)
        x = self.loc_relu4(x)
        x = self.loc_pool4(x)

        y = x.view(batch_size, -1)
        y = self.loc_relu5(y)
        batch_C_prime = self.loc_reg(y)

        return batch_C_prime


"""
Calculate driving_force due to ZPF tielines.

The general approach is similar to the PanOptimizer rough search method.

1. With all phases active, calculate the chemical potentials of the tieline
   endpoints via ``equilibrium`` calls. Done in ``estimate_hyperplane``.
2. Calculate the target chemical potentials, which are the average chemical
   potentials of all of the current chemical potentials at the tieline endpoints.
3. Calculate the current chemical potentials of the desired single phases
4. The error is the difference between these chemical potentials

There's some special handling for tieline endpoints where we do not know the
composition conditions to calculate chemical potentials at.
"""

import logging
from dataclasses import dataclass
from collections import OrderedDict
from typing import Sequence, Dict, Any, Union, List, Tuple, Type, Optional

import numpy as np
from numpy.typing import ArrayLike
from scipy.stats import norm
import tinydb

from pycalphad import Database, Model, variables as v
from pycalphad.codegen.callables import build_phase_records
from pycalphad.core.utils import instantiate_models, filter_phases, unpack_components
from pycalphad.core.phase_rec import PhaseRecord
from espei.utils import PickleableTinyDB
from espei.shadow_functions import equilibrium_, calculate_, no_op_equilibrium_, update_phase_record_parameters, constrained_equilibrium
from pycalphad.core.calculate import _sample_phase_constitution
from pycalphad.core.utils import point_sample

_log = logging.getLogger(__name__)

@dataclass
class RegionVertex:
    phase_name: str
    composition: ArrayLike  # 1D of size (number nonvacant pure elements)
    comp_conds: Dict[v.X, float]
    points: ArrayLike
    phase_records: Dict[str, PhaseRecord]
    is_disordered: bool
    has_missing_comp_cond: bool

@dataclass
class PhaseRegion:
    vertices: Sequence[RegionVertex]
    potential_conds: Dict[v.StateVariable, float]
    species: Sequence[v.Species]
    phases: Sequence[str]
    models: Dict[str, Model]

    def eq_str(self):
        phase_compositions = ', '.join(f'{vtx.phase_name}: {vtx.comp_conds}' for vtx in self.vertices)
        return f"conds: ({self.potential_conds}), comps: ({phase_compositions})"

def _extract_pot_conds(all_conditions: Dict[v.StateVariable, np.ndarray], idx: int) -> Dict[v.StateVariable, float]:
    """Conditions are either scalar or 1d arrays for the conditions in the entire dataset.
    This function extracts the condition corresponding to the current region,
    based on the index in the 1d condition array.
    """
    pot_conds = {}  # e.g. v.P, v.T
    for cond_key, cond_val in all_conditions.items():
        cond_val = np.atleast_1d(np.asarray(cond_val))
        # If the conditions is an array, we want the corresponding value
        # Otherwise treat it as a scalar
        if len(cond_val) > 1:
            cond_val = cond_val[idx]
        pot_conds[getattr(v, cond_key)] = float(cond_val)
    return pot_conds

def _extract_phases_comps(vertex):
    """Extract the phase name, phase compositions and disordered flag from a vertex
    """
    if len(vertex) == 4:  # phase_flag within
        phase_name, components, compositions, flag = vertex
        if flag == "disordered":
            disordered_flag = True
        else:
            disordered_flag = False
    elif len(vertex) == 3:  # no phase_flag within
        phase_name, components, compositions = vertex
        disordered_flag = False
    else:
        raise ValueError("Wrong number of data in tie-line point")
    comp_conds = dict(zip(map(v.X, map(str.upper, components)), compositions))
    return phase_name, comp_conds, disordered_flag

def _phase_is_stoichiometric(mod):
    return all(len(subl) == 1 for subl in mod.constituents)

def _compute_vertex_composition(comps: Sequence[str], comp_conds: Dict[str, float]):
    """Compute the overall composition in a vertex assuming an N=1 normalization condition"""
    pure_elements = sorted(c for c in comps if c != 'VA')
    vertex_composition = np.empty(len(pure_elements), dtype=np.float_)
    unknown_indices = []
    for idx, el in enumerate(pure_elements):
        amt = comp_conds.get(v.X(el), None)
        if amt is None:
            unknown_indices.append(idx)
            vertex_composition[idx] = np.nan
        else:
            vertex_composition[idx] = amt
    if len(unknown_indices) == 1:
        # Determine the dependent component by mass balance
        vertex_composition[unknown_indices[0]] = 1 - np.nansum(vertex_composition)
    return vertex_composition

def _subsample_phase_points(phase_record, phase_points, target_composition, avg_mass_residual_tol=0.02):
    # Compute the mole fractions of each point
    phase_compositions = np.zeros((phase_points.shape[0], target_composition.size), order='F')
    # TODO: potential bug here if the composition has dependence (even piecewise
    #   dependence) in the state variables. The compositions may be nan in this case.
    statevar_placeholder = np.ones((phase_points.shape[0], phase_record.num_statevars))
    dof = np.hstack((statevar_placeholder, phase_points))
    for el_idx in range(target_composition.size):
        phase_record.mass_obj_2d(phase_compositions[:, el_idx], dof, el_idx)

    # Find the points indicdes where the mass is within the average mass residual tolerance
    idxs = np.nonzero(np.mean(np.abs(phase_compositions - target_composition), axis=1) < avg_mass_residual_tol)[0]

    # Return the sub-space of points where this condition holds valid
    return phase_points[idxs]

def get_zpf_data(dbf: Database, comps: Sequence[str], phases: Sequence[str], datasets: PickleableTinyDB, parameters: Dict[str, float], model: Optional[Dict[str, Type[Model]]] = None):
    """
    Return the ZPF data used in the calculation of ZPF error

    Parameters
    ----------
    comps : list
        List of active component names
    phases : list
        List of phases to consider
    datasets : espei.utils.PickleableTinyDB
        Datasets that contain single phase data
    parameters : dict
        Dictionary mapping symbols to optimize to their initial values
    model : Optional[Dict[str, Type[Model]]]
        Dictionary phase names to pycalphad Model classes.

    Returns
    -------
    list
        List of data dictionaries with keys ``weight``, ``phase_regions`` and ``dataset_references``.
    """
    desired_data = datasets.search((tinydb.where('output') == 'ZPF') &
                                   (tinydb.where('components').test(lambda x: set(x).issubset(comps))) &
                                   (tinydb.where('phases').test(lambda x: len(set(phases).intersection(x)) > 0)))

    zpf_data = []  # 1:1 correspondence with each dataset
    for data in desired_data:
        data_comps = list(set(data['components']).union({'VA'}))
        species = sorted(unpack_components(dbf, data_comps), key=str)
        data_phases = filter_phases(dbf, species, candidate_phases=phases)
        models = instantiate_models(dbf, species, data_phases, model=model, parameters=parameters)
        all_phase_points = {phase_name: _sample_phase_constitution(models[phase_name], point_sample, True, 50) for phase_name in data_phases}
        all_regions = data['values']
        conditions = data['conditions']
        phase_regions = []
        # Each phase_region is one set of phases in equilibrium (on a tie-line),
        # e.g. [["ALPHA", ["B"], [0.25]], ["BETA", ["B"], [0.5]]]
        for idx, phase_region in enumerate(all_regions):
            # Extract the conditions for entire phase region
            pot_conds = _extract_pot_conds(conditions, idx)
            pot_conds.setdefault(v.N, 1.0) # Add v.N condition, if missing
            # Extract all the phases and compositions from the tie-line points
            vertices = []
            for vertex in phase_region:
                phase_name, comp_conds, disordered_flag = _extract_phases_comps(vertex)
                phase_recs = build_phase_records(dbf, species, data_phases, {**pot_conds, **comp_conds}, models, parameters=parameters, build_gradients=True, build_hessians=True)
                # Construct single-phase points satisfying the conditions for each phase in the region
                mod = models[phase_name]
                composition = _compute_vertex_composition(data_comps, comp_conds)
                if np.any(np.isnan(composition)):
                    # We can't construct points because we don't have a known composition
                    has_missing_comp_cond = True
                    phase_points = None
                elif _phase_is_stoichiometric(mod):
                    has_missing_comp_cond = False
                    phase_points = None
                else:
                    has_missing_comp_cond = False
                    # Only sample points that have an average mass residual within tol
                    tol = 0.05
                    phase_points = _subsample_phase_points(phase_recs[phase_name], all_phase_points[phase_name], composition, tol)
                    assert phase_points.shape[0] > 0, "at least one set of points is within the target tolerance"
                vtx = RegionVertex(phase_name, composition, comp_conds, phase_points, phase_recs, disordered_flag, has_missing_comp_cond)
                vertices.append(vtx)
            region = PhaseRegion(vertices, pot_conds, species, data_phases, models)
            phase_regions.append(region)

        data_dict = {
            'weight': data.get('weight', 1.0),
            'phase_regions': phase_regions,
            'dataset_reference': data['reference']
        }
        zpf_data.append(data_dict)
    return zpf_data

def estimate_hyperplane(phase_region: PhaseRegion, parameters: np.ndarray, approximate_equilibrium: bool = False) -> np.ndarray:
    """
    Calculate the chemical potentials for the target hyperplane, one vertex at a time

    Notes
    -----
    This takes just *one* set of phase equilibria, a phase region, e.g. a dataset point of
    [['FCC_A1', ['CU'], [0.1]], ['LAVES_C15', ['CU'], [0.3]]]
    and calculates the chemical potentials given all the phases possible at the
    given compositions. Then the average chemical potentials of each end point
    are taken as the target hyperplane for the given equilibria.

    """
    if approximate_equilibrium:
        _equilibrium = no_op_equilibrium_
    else:
        _equilibrium = equilibrium_
    target_hyperplane_chempots = []
    target_hyperplane_phases = []
    species = phase_region.species
    phases = phase_region.phases
    models = phase_region.models
    for vertex in phase_region.vertices:
        phase_records = vertex.phase_records
        update_phase_record_parameters(phase_records, parameters)
        cond_dict = {**vertex.comp_conds, **phase_region.potential_conds}
        if vertex.has_missing_comp_cond:
            # This composition is unknown -- it doesn't contribute to hyperplane estimation
            pass
        else:
            # Extract chemical potential hyperplane from multi-phase calculation
            # Note that we consider all phases in the system, not just ones in this tie region
            str_statevar_dict = OrderedDict([(str(key), cond_dict[key]) for key in sorted(phase_region.potential_conds.keys(), key=str)])
            grid = calculate_(species, phases, str_statevar_dict, models, phase_records, pdens=50, fake_points=True)
            multi_eqdata = _equilibrium(phase_records, cond_dict, grid)
            target_hyperplane_phases.append(multi_eqdata.Phase.squeeze())
            # Does there exist only a single phase in the result with zero internal degrees of freedom?
            # We should exclude those chemical potentials from the average because they are meaningless.
            num_phases = np.sum(multi_eqdata.Phase.squeeze() != '')
            Y_values = multi_eqdata.Y.squeeze()
            no_internal_dof = np.all((np.isclose(Y_values, 1.)) | np.isnan(Y_values))
            MU_values = multi_eqdata.MU.squeeze()
            if (num_phases == 1) and no_internal_dof:
                target_hyperplane_chempots.append(np.full_like(MU_values, np.nan))
            else:
                target_hyperplane_chempots.append(MU_values)
    target_hyperplane_mean_chempots = np.nanmean(target_hyperplane_chempots, axis=0, dtype=np.float_)
    return target_hyperplane_mean_chempots

def driving_force_to_hyperplane(target_hyperplane_chempots: np.ndarray,
                                phase_region: PhaseRegion, vertex: RegionVertex,
                                parameters: np.ndarray, approximate_equilibrium: bool = False) -> float:
    """Calculate the integrated driving force between the current hyperplane and target hyperplane.
    """
    species = phase_region.species
    models = phase_region.models
    current_phase = vertex.phase_name
    cond_dict = {**phase_region.potential_conds, **vertex.comp_conds}
    str_statevar_dict = OrderedDict([(str(key),cond_dict[key]) for key in sorted(phase_region.potential_conds.keys(), key=str)])
    phase_points = vertex.points
    phase_records = vertex.phase_records
    update_phase_record_parameters(phase_records, parameters)
    if phase_points is None:
        # We don't have the phase composition here, so we estimate the driving force.
        # Can happen if one of the composition conditions is unknown or if the phase is
        # stoichiometric and the user did not specify a valid phase composition.
        single_eqdata = calculate_(species, [current_phase], str_statevar_dict, models, phase_records, pdens=50)
        df = np.multiply(target_hyperplane_chempots, single_eqdata.X).sum(axis=-1) - single_eqdata.GM
        driving_force = float(df.max())
    elif vertex.is_disordered:
        # Construct disordered sublattice configuration from composition dict
        # Compute energy
        # Compute residual driving force
        # TODO: Check that it actually makes sense to declare this phase 'disordered'
        num_dof = sum([len(subl) for subl in models[current_phase].constituents])
        desired_sitefracs = np.ones(num_dof, dtype=np.float_)
        dof_idx = 0
        for subl in models[current_phase].constituents:
            dof = sorted(subl, key=str)
            num_subl_dof = len(subl)
            if v.Species("VA") in dof:
                if num_subl_dof == 1:
                    _log.debug('Cannot predict the site fraction of vacancies in the disordered configuration %s of %s. Returning driving force of zero.', subl, current_phase)
                    return 0
                else:
                    sitefracs_to_add = [1.0]
            else:
                sitefracs_to_add = np.array([cond_dict.get(v.X(d)) for d in dof], dtype=np.float_)
                # Fix composition of dependent component
                sitefracs_to_add[np.isnan(sitefracs_to_add)] = 1 - np.nansum(sitefracs_to_add)
            desired_sitefracs[dof_idx:dof_idx + num_subl_dof] = sitefracs_to_add
            dof_idx += num_subl_dof
        single_eqdata = calculate_(species, [current_phase], str_statevar_dict, models, phase_records, points=np.asarray([desired_sitefracs]))
        driving_force = np.multiply(target_hyperplane_chempots, single_eqdata.X).sum(axis=-1) - single_eqdata.GM
        driving_force = float(np.squeeze(driving_force))
    else:
        # Extract energies from single-phase calculations
        grid = calculate_(species, [current_phase], str_statevar_dict, models, phase_records, points=phase_points, pdens=50, fake_points=True)
        # TODO: consider enabling approximate for this?
        converged, energy = constrained_equilibrium(phase_records, cond_dict, grid)
        if not converged:
            _log.debug('Calculation failure: constrained equilibrium not converged for %s, conditions: %s, parameters %s', current_phase, cond_dict, parameters)
            return np.inf
        driving_force = float(np.dot(target_hyperplane_chempots, vertex.composition) - float(energy))
    return driving_force

def calculate_zpf_driving_forces(zpf_data: Sequence[Dict[str, Any]],
                                 parameters: ArrayLike = None,
                                 approximate_equilibrium: bool = False,
                                 short_circuit: bool = False
                                 ) -> Tuple[List[List[float]], List[List[float]]]:
    """
    Calculate error due to phase equilibria data

    zpf_data : Sequence[Dict[str, Any]]
        Datasets that contain single phase data
    parameters : ArrayLike
        Array of parameters to calculate the error with.
    approximate_equilibrium : bool
        Whether or not to use an approximate version of equilibrium that does
        not refine the solution and uses ``starting_point`` instead.
    short_circuit: bool
        If True, immediately return a size 1 array with a driving force of
        ``np.nan`` (failed hyperplane) or ``np.inf`` (failed driving force).
        Can save computational time if the caller will aggregate driving forces.

    Returns
    -------
    Tuple[List[List[float]], List[List[float]]]
        Driving forces and weights as ragged 2D arrays with shape
        ``(len(zpf_data), len(vertices in each zpf_data))``

    Notes
    -----
    The physical picture of the standard deviation is that we've measured a ZPF
    line. That line corresponds to some equilibrium chemical potentials. The
    standard deviation is the standard deviation of those 'measured' chemical
    potentials.

    """
    if parameters is None:
        parameters = np.array([])
    driving_forces = []
    weights = []
    for data in zpf_data:
        data_driving_forces = []
        data_weights = []
        weight = data['weight']
        dataset_ref = data['dataset_reference']
        # for the set of phases and corresponding tie-line verticies in equilibrium
        for phase_region in data['phase_regions']:
            # 1. Calculate the average multiphase hyperplane
            eq_str = phase_region.eq_str()
            target_hyperplane = estimate_hyperplane(phase_region, parameters, approximate_equilibrium=approximate_equilibrium)
            if np.any(np.isnan(target_hyperplane)):
                _log.debug('NaN target hyperplane. Equilibria: (%s), driving force: 0.0, reference: %s.', eq_str, dataset_ref)
                data_driving_forces.extend([0]*len(phase_region.vertices))
                data_weights.extend([weight]*len(phase_region.vertices))
                continue
            # 2. Calculate the driving force to that hyperplane for each vertex
            for vertex in phase_region.vertices:
                driving_force = driving_force_to_hyperplane(target_hyperplane, phase_region, vertex, parameters,
                                                            approximate_equilibrium=approximate_equilibrium,
                                                            )
                if np.isinf(driving_force) and short_circuit:
                    _log.debug('Equilibria: (%s), current phase: %s, hyperplane: %s, driving force: %s, reference: %s. Short circuiting.', eq_str, vertex.phase_name, target_hyperplane, driving_force, dataset_ref)
                    return [[np.inf]], [[np.inf]]
                data_driving_forces.append(driving_force)
                data_weights.append(weight)
                _log.debug('Equilibria: (%s), current phase: %s, hyperplane: %s, driving force: %s, reference: %s', eq_str, vertex.phase_name, target_hyperplane, driving_force, dataset_ref)
        driving_forces.append(data_driving_forces)
        weights.append(data_weights)
    return driving_forces, weights

def calculate_zpf_error(zpf_data: Sequence[Dict[str, Any]],
                        parameters: np.ndarray = None,
                        data_weight: int = 1.0,
                        approximate_equilibrium: bool = False) -> float:
    """
    Calculate the likelihood due to phase equilibria data.

    For detailed documentation, see ``calculate_zpf_driving_forces``

    Returns
    -------
    float
        Log probability of ZPF driving forces

    """
    if len(zpf_data) == 0:
        return 0.0
    driving_forces, weights = calculate_zpf_driving_forces(zpf_data, parameters, approximate_equilibrium, short_circuit=True)
    # Driving forces and weights are 2D ragged arrays with the shape (len(zpf_data), len(zpf_data['values']))
    driving_forces = np.concatenate(driving_forces)
    weights = np.concatenate(weights)
    if np.any(np.logical_or(np.isinf(driving_forces), np.isnan(driving_forces))):
        return -np.inf
    log_probabilites = norm.logpdf(driving_forces, loc=0, scale=1000/data_weight/weights)
    _log.debug('Data weight: %s, driving forces: %s, weights: %s, probabilities: %s', data_weight, driving_forces, weights, log_probabilites)
    return np.sum(log_probabilites)