"""
Utilities for rock chemistry and mineral abundance classification.

Todo
-------

* Petrological classifiers: QAPF (aphanitic/phaneritic),
  gabbroic Pyroxene-Olivine-Plagioclase,
  ultramafic Olivine-Orthopyroxene-Clinopyroxene
"""
import os
import json
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.text
import matplotlib.pyplot as plt
import matplotlib.patches
import matplotlib.lines
from .plot.style import patchkwargs
from .plot.axes import init_axes
from .plot.helpers import get_centroid
from .meta import (
    pyrolite_datafolder,
    subkwargs,
    sphinx_doi_link,
    update_docstring_references,
)
from .log import Handle

logger = Handle(__name__)

class PolygonClassifier(object):
    """
    A classifier model built form a series of polygons defining specific classes.

    Parameters
    -----------
    name : :class:`str`
        A name for the classifier model.
    axes : :class:`list` | :class:`tuple`
        Names of the axes corresponding to the polygon coordinates.
    fields : :class:`dict`
        Dictionary describing indiviudal polygons, with identifiers as keys and
        dictionaries containing 'name' and 'fields' items.
    scale : :class:`float`
        Default maximum scale for the axes. Typically 100 (wt%) or 1 (fractional).
    xlim : :class:`tuple`
        Default x-limits for this classifier for plotting.
    ylim : :class:`tuple`
        Default y-limits for this classifier for plotting.
    """

    def __init__(
        self, name=None, axes=None, fields=None, scale=1.0, xlim=None, ylim=None,
    ):
        self.default_scale = scale
        self._scale = self.default_scale
        self.xlim = xlim
        self.ylim = ylim

        self.name = name
        self.axes = axes or []
        # check axes for ratios, adition/subtraction etc
        self.fields = fields or []
        self.classes = list(self.fields.keys())

    def predict(self, X, data_scale=None):
        """
        Predict the classification of samples using the polygon-based classifier.

        Parameters
        -----------
        X : :class:`numpy.ndarray` | :class:`pandas.DataFrame`
            Data to classify.
        data_scale : :class:`float`
            Maximum scale for the data. Typically 100 (wt%) or 1 (fractional).

        Returns
        -------
        :class:`pandas.Series`
            Series containing classifer predictions. If a dataframe was input,
            it inherit the index.
        """
        classes = [k for (k, cfg) in self.fields.items() if cfg["poly"]]
        polys = [
            matplotlib.patches.Polygon(self.fields[k]["poly"], closed=True)
            for k in classes
        ]
        if isinstance(X, pd.DataFrame):
            # check whether the axes names are in the columns
            axes = self.axis_components
            idx = X.index
            X = X.loc[:, axes].values
        else:
            idx = np.arange(X.shape[0])
        out = pd.Series(index=idx, dtype="object")

        rescale_by = 1.0  # rescaling the data to fit the classifier scale
        if data_scale is not None:
            if not np.isclose(self.default_scale, data_scale):
                rescale_by = self.default_scale / data_scale
        X = X * rescale_by

        indexes = np.array([p.contains_points(X) for p in polys]).T
        notfound = np.logical_not(indexes.sum(axis=-1))

        outlist = list(map(lambda ix: classes[ix], np.argmax(indexes, axis=-1)))
        out.loc[:] = outlist
        out.loc[(notfound)] = "none"
        return out

    @property
    def axis_components(self):
        """
        Get the axis components used by the classifier.

        Returns
        -------
        :class:`tuple`
            Names of the x and y axes for the classifier.
        """
        return self.axes.get("x"), self.axes.get("y")

    def _add_polygons_to_axes(
        self, ax=None, fill=False, axes_scale=100.0, labels=None, **kwargs
    ):
        """
        Add the polygonal fields from the classifier to an axis.

        Parameters
        ----------
        ax : :class:`matplotlib.axes.Axes`
            Axis to add the polygons to.
        fill : :class:`bool`
            Whether to fill the polygons.
        axes_scale : :class:`float`
            Maximum scale for the axes. Typically 100 (for wt%) or 1 (fractional).
        labels : :class:`str`
            Which labels to add to the polygons (e.g. for TAS, 'volcanic', 'intrusive'
            or the field 'ID').

        Returns
        --------
        ax : :class:`matplotlib.axes.Axes`
        """
        if ax is None:
            ax = init_axes(**kwargs)

        rescale_by = 1.0
        if axes_scale is not None:  # rescale polygons to fit ax
            if not np.isclose(self.default_scale, axes_scale):
                rescale_by = axes_scale / self.default_scale
        pgns = []
        for k, cfg in self.fields.items():
            if cfg["poly"]:
                if not fill:
                    kwargs["facecolor"] = "none"
                verts = np.array(cfg["poly"]) * rescale_by
                pg = matplotlib.patches.Polygon(
                    verts, closed=True, edgecolor="k", **patchkwargs(kwargs)
                )
                pgns.append(pg)
                ax.add_patch(pg)

        # if the axis has the default scaling, there's a good chance that it hasn't
        # been rescaled/rendered. We need to rescale to show the polygons.
        if np.allclose(ax.get_xlim(), [0, 1]) & np.allclose(ax.get_ylim(), [0, 1]):
            ax.set_xlim(np.array(self.xlim) * rescale_by)
            ax.set_ylim(np.array(self.ylim) * rescale_by)

        return ax

    def add_to_axes(self, ax=None, fill=False, axes_scale=1, **kwargs):
        """
        Add the polygonal fields from the classifier to an axis.

        Parameters
        ----------
        ax : :class:`matplotlib.axes.Axes`
            Axis to add the polygons to.
        fill : :class:`bool`
            Whether to fill the polygons.
        axes_scale : :class:`float`
            Maximum scale for the axes. Typically 100 (for wt%) or 1 (fractional).

        Returns
        --------
        ax : :class:`matplotlib.axes.Axes`
        """
        ax = self._add_polygons_to_axes(
            ax=ax, fill=fill, axes_scale=axes_scale, **kwargs
        )
        if self.axes:  # may be none?
            ax.set_ylabel(self.axes[0])
            ax.set_xlabel(self.axes[1])
        return ax

class TAS(PolygonClassifier):
    """
    Total-alkali Silica Diagram classifier from Le Bas (1992) [#ref_1]_.

    Parameters
    -----------
    name : :class:`str`
        A name for the classifier model.
    axes : :class:`list` | :class:`tuple`
        Names of the axes corresponding to the polygon coordinates.
    fields : :class:`dict`
        Dictionary describing indiviudal polygons, with identifiers as keys and
        dictionaries containing 'name' and 'fields' items.
    scale : :class:`float`
        Default maximum scale for the axes. Typically 100 (wt%) or 1 (fractional).
    xlim : :class:`tuple`
        Default x-limits for this classifier for plotting.
    ylim : :class:`tuple`
        Default y-limits for this classifier for plotting.

    References
    -----------
    .. [#ref_1] Le Bas, M.J., Le Maitre, R.W., Woolley, A.R., 1992.
                The construction of the Total Alkali-Silica chemical
                classification of volcanic rocks.
                Mineralogy and Petrology 46, 1–22.
                doi: {LeBas1992}
    """

    @update_docstring_references
    def __init__(self, **kwargs):
        src = pyrolite_datafolder(subfolder="models") / "TAS" / "config.json"

        with open(src, "r") as f:
            config = json.load(f)
        kw = dict(scale=100.0, xlim=[35, 85], ylim=[0, 20])
        kw.update(kwargs)
        poly_config = {**config, **kw}
        super().__init__(**poly_config)

    def add_to_axes(self, ax=None, fill=False, axes_scale=100.0, labels=None, **kwargs):
        """
        Add the TAS fields from the classifier to an axis.

        Parameters
        ----------
        ax : :class:`matplotlib.axes.Axes`
            Axis to add the polygons to.
        fill : :class:`bool`
            Whether to fill the polygons.
        axes_scale : :class:`float`
            Maximum scale for the axes. Typically 100 (for wt%) or 1 (fractional).
        labels : :class:`str`
            Which labels to add to the polygons (e.g. for TAS, 'volcanic', 'intrusive'
            or the field 'ID').

        Returns
        --------
        ax : :class:`matplotlib.axes.Axes`
        """
        # use and override the default add_to_axes
        ax = self._add_polygons_to_axes(
            ax=ax, fill=fill, axes_scale=axes_scale, **kwargs
        )
        rescale_by = 1.0
        if axes_scale is not None:  # rescale polygons to fit ax
            if not np.isclose(self.default_scale, axes_scale):
                rescale_by = axes_scale / self.default_scale
        if labels is not None:
            for k, cfg in self.fields.items():
                if cfg["poly"]:
                    verts = np.array(cfg["poly"]) * rescale_by
                    x, y = get_centroid(matplotlib.patches.Polygon(verts))
                    if "volc" in labels:  # use the volcanic name
                        label = cfg["name"][0]
                    elif "intr" in labels:  # use the intrusive name
                        label = cfg["name"][-1]
                    else:  # use the field identifier
                        label = k
                    ax.annotate(
                        "\n".join(label.split()),
                        xy=(x, y),
                        ha="center",
                        va="center",
                        **subkwargs(kwargs, ax.annotate, matplotlib.text.Text)
                    )

        ax.set_ylabel("$Na_2O + K_2O$")
        ax.set_xlabel("$SiO_2$")
        return ax

class PeralkalinityClassifier(object):
    def __init__(self):
        self.fields = None

    def predict(self, df: pd.DataFrame):
        TotalAlkali = df.Na2O + df.K2O
        perkalkaline_where = (df.Al2O3 < (TotalAlkali + df.CaO)) & (
            TotalAlkali > df.Al2O3
        )
        metaluminous_where = (df.Al2O3 > (TotalAlkali + df.CaO)) & (
            TotalAlkali < df.Al2O3
        )
        peraluminous_where = (df.Al2O3 < (TotalAlkali + df.CaO)) & (
            TotalAlkali < df.Al2O3
        )
        out = pd.Series(index=df.index, dtype="object")
        out.loc[peraluminous_where] = "Peraluminous"
        out.loc[metaluminous_where] = "Metaluminous"
        out.loc[perkalkaline_where] = "Peralkaline"
        return out

TAS.__init__.__doc__ = TAS.__init__.__doc__.format(
    LeBas1992=sphinx_doi_link("10.1007/BF01160698")
)


# -*- coding: utf-8 -*-
# Copyright (C) 2017
# Full license can be found in LICENSE.txt
#---------------------------------------------------------------------------
""" Perform OCB gridding for SuperDARN vorticity data

Functions
----------------------------------------------------------------------------
vort2ascii_ocb(vortfile, outfile, kwargs)
    Write and ASCII file with SuperDARN data and the OCB coordinates for each
    data point
load_vorticity_ascii_data(filename, save_all=False)
    Load vorticity block ASCII data files

Data
----------------------------------------------------------------------------
Specialised SuperDARN data product, available from: gchi@bas.ac.uk
"""
import logbook as logging
import numpy as np

def vort2ascii_ocb(vortfile, outfile, ocb=None, ocbfile=None, max_sdiff=600,
                   save_all=False, min_sectors=7, rcent_dev=8.0, max_r=23.0,
                   min_r=10.0, min_j=0.15):
    """ Coverts the location of vorticity data in AACGM coordinates into a frame
    that is relative to the open-closed field-line boundary (OCB) as determined
    from a circle fit to the poleward boundary of the auroral oval

    Parameters
    ----------
    vortfile : (str)
        file containing the required vorticity file sorted by time
    outfile : (str)
        filename for the output data
    ocb : (ocbpy.ocboundary.OCBoundary or NoneType)
        Object containing open closed boundary data or None to load from file
    ocbfile : (str or NoneType)
        file containing the required OC boundary data sorted by time, or None
        to use ocb object or IMAGE WIC file (default=None)
    max_sdiff : (int)
        maximum seconds between OCB and data record in sec (default=600)
    save_all : (bool)
        Save all data (True), or only that needed to calcuate OCB and vorticity
        (False). (default=False)
    min_sectors : (int)
        Minimum number of MLT sectors required for good OCB. (default=7)
    rcent_dev : (float)
        Maximum number of degrees between the new centre and the AACGM pole
        (default=8.0).
    max_r : (float)
        Maximum radius for open-closed field line boundary in degrees.
        (default=23.0)
    min_r : (float)
        Minimum radius for open-closed field line boundary in degrees
        (default=10.0)
    min_j : (float)
        Minimum unitless current magnitude scale difference (default=0.15)

    Returns
    ---------
    Void

    Notes
    --------
    Input header or col_names must include the names in the default string.
    """
    import ocbpy
    import ocbpy.ocb_scaling as ocbscal
    import datetime as dt

    assert ocbpy.instruments.test_file(vortfile), \
        logging.error("vorticity file cannot be opened[{:s}]".format(vortfile))
    assert isinstance(outfile, str), \
        logging.error("output filename is not a string [{:}]".format(outfile))

    # Read the vorticity data
    vdata = load_vorticity_ascii_data(vortfile, save_all=save_all)
    need_keys = ["VORTICITY", "CENTRE_MLAT", "DATETIME", "MLT"]

    if vdata is None or not all([kk in vdata.keys() for kk in need_keys]):
        estr = "unable to load necessary data from [{:s}]".format(vortfile)
        logging.error(estr)
        return

    # Load the OCB data
    if ocb is None or not isinstance(ocb, ocbpy.ocboundary.OCBoundary):
        vstart = vdata['DATETIME'][0] - dt.timedelta(seconds=max_sdiff+1)
        vend = vdata['DATETIME'][-1] + dt.timedelta(seconds=max_sdiff+1)
        ocb = ocbpy.ocboundary.OCBoundary(ocbfile, stime=vstart, etime=vend)

    if ocb.filename is None or ocb.records == 0:
        try:
            logging.error("no data in OCB file {:s} ".format(ocb.filename))
        except:
            logging.error("bad OCB file specified")
        return

    # Set the reference radius
    ref_r = 90.0 - abs(ocb.boundary_lat)

    # Open and test the file to ensure it can be written
    try:
        fout = open(outfile, 'w')
    except:
        logging.error("unable to create output file [{:}]".format(outfile))
        return

    # Write header line
    outline = "#DATE TIME "

    if save_all:
        vkeys = vdata.keys()
        vkeys.pop(vkeys.index("DATETIME"))
        outline = "{:s}{:s} ".format(outline, " ".join(vkeys))

    outline = "{:s}OCB_LAT OCB_MLT NORM_VORT\n".format(outline)

    try:
        fout.write(outline)
    except:
        estr = "unable to write [{:s}] because of error ".format(outline)
        estr = "{:s}[{:}]".format(estr, e)
        logging.error(estr)
        return

    # Initialise the ocb and vorticity indices
    ivort = 0
    num_vort = vdata['DATETIME'].shape[0]

    # Cycle through the data, matching vorticity and OCB records
    while ivort < num_vort and ocb.rec_ind < ocb.records:
        ivort = ocbpy.match_data_ocb(ocb, vdata['DATETIME'], idat=ivort,
                                     max_tol=max_sdiff, min_sectors=min_sectors,
                                     rcent_dev=rcent_dev, max_r=max_r,
                                     min_r=min_r, min_j=min_j)

        if ivort < num_vort and ocb.rec_ind < ocb.records:
            # Use the indexed OCB to convert the AACGM grid coordinate to one
            # related to the OCB
            nlat, nmlt = ocb.normal_coord(vdata['CENTRE_MLAT'][ivort],
                                          vdata['MLT'][ivort])
            nvort = ocbscal.normal_curl_evar(vdata['VORTICITY'][ivort],
                                             ocb.r[ocb.rec_ind], ref_r)

            # Format the output line
            #    DATE TIME (SAVE_ALL) OCB_LAT OCB_MLT NORM_VORT
            outline = "{:} ".format(vdata['DATETIME'][ivort])

            if save_all:
                for k in vkeys:
                    outline = "{:s}{:} ".format(outline, vdata[k][ivort])

            outline = "{:s}{:.2f} {:.6f} {:.6f}\n".format(outline, nlat, nmlt,
                                                          nvort)

            try:
                fout.write(outline)
            except e:
                estr = "unable to write [{:s}] ".format(outline)
                estr = "{:s}because of error [{:}]".format(estr, e)
                logging.error(estr)
                return

            # Move to next line
            ivort += 1

    # Close output file
    fout.close()

    return

def load_vorticity_ascii_data(vortfile, save_all=False):
    """Load SuperDARN vorticity data files.

    Parameters
    -----------
    vortfile : (str)
        SuperDARN vorticity file in block format
    save_all : (bool)
        Save all data from the file (True), or only data needed to calculate
        the OCB coordinates and normalised vorticity (False). (default=False)

    Returns
    ---------
    vdata : (dict)
        Dictionary of numpy arrays
    """
    from ocbpy.instruments import test_file
    import datetime as dt

    if not test_file(vortfile):
        return None

    # Open the data file
    try:
        fvort = open(vortfile, "r")
    except:
        logging.error("unable to open vorticity file [{:s}]".format(vortfile))
        return None

    # Initialise the output dictionary
    vkeys = ["YEAR", "MONTH", "DAY", "UTH", "VORTICITY", "MLT", "CENTRE_MLAT",
             "DATETIME"]
    if save_all:
        vkeys.extend(["R1BM1", "R1BM2", "R2BM1", "R2BM2", "AREA", "CENTRE_GLAT",
                      "CENTRE_GLON", "C1_GLAT", "C1_GLON", "C2_GLAT", "C2_GLON",
                      "C3_GLAT", "C3_GLON", "C4_GLAT", "C4_GLON", "CENTRE_MLON",
                      "C1_MLAT", "C1_MLON", "C2_MLAT", "C2_MLON", "C3_MLAT",
                      "C3_MLON", "C4_MLAT", "C4_MLON"])
    vdata = {k:list() for k in vkeys}
    vkeys = set(vkeys)

    # Set the data block keys
    bkeys = [["R1BM1", "R1BM2", "R2BM1", "R2BM2", "AREA", "VORTICITY", "MLT"],
             ["GFLG", "CENTRE_GLAT", "CENTRE_GLON", "C1_GLAT", "C1_GLON",
              "C2_GLAT", "C2_GLON", "C3_GLAT", "C3_GLON", "C4_GLAT", "C4_GLON"],
             ["MFLG", "CENTRE_MLAT", "CENTRE_MLON", "C1_MLAT", "C1_MLON",
              "C2_MLAT", "C2_MLON", "C3_MLAT", "C3_MLON", "C4_MLAT", "C4_MLON"]]

    # Read the lines and assign data.  Recall that blank lines in file are
    # returned as '\n'
    vline = fvort.readline()
    vsplit = vline.split()
    vinc = 0

    while len(vline) > 0:
        if vinc == 0:
            # This is a date line
            if len(vsplit) != 4:
                estr = "unexpected line encountered when date line "
                estr = "{:s}expected [{:s}]".format(estr, vline)
                logging.error(estr)
                fvort.close()
                return None

            # Save the data in the format desired for the output dict
            yy = int(vsplit[0])
            mm = int(vsplit[1])
            dd = int(vsplit[2])
            hh = float(vsplit.pop())

            # Calculate and save the datetime
            stime = " ".join(vsplit)
            dtime = (dt.datetime.strptime(stime, "%Y %m %d") +
                     dt.timedelta(seconds=np.floor(hh * 3600.0)))
            vinc += 1
        elif vinc == 1:
            # This is a number of entries line
            if len(vsplit) != 1:
                estr = "unexpected line encountered when number of entries "
                estr = "{:s}line expected [{:s}]".format(estr, vline)
                logging.error(estr)
                fvort.close()
                return None

            # Save the number of entries
            nentries = int(vsplit[0])
            vinc += 1
        else:
            # This is an entry.  For each entry there are three lines
            ninc = 0
            while ninc < nentries:
                # Save the time data
                vdata['YEAR'].append(yy)
                vdata['MONTH'].append(mm)
                vdata['DAY'].append(dd)
                vdata['UTH'].append(hh)
                vdata['DATETIME'].append(dtime)

                for bklist in bkeys:
                    # Test to see that this line has the right number of col
                    if len(vsplit) != len(bklist):
                        estr = "unexpected line encountered for a data block "
                        estr = "{:s}[{:s}]".format(estr, vline)
                        logging.error(estr)
                        fvort.close()
                        return None

                    # Save all desired keys
                    gkeys = list(vkeys.intersection(bklist))

                    for gk in gkeys:
                        ik = bklist.index(gk)
                        vdata[gk].append(float(vsplit[ik]))

                    # Move to next line
                    vline = fvort.readline()
                    vsplit = vline.split()

                # All data lines for this entry have been processed, incriment
                ninc += 1

            # All entries in block have been processed, reset incriment
            vinc = 0

        # Move to next line
        vline = fvort.readline()
        vsplit = vline.split()

    # Close file handle
    fvort.close()

    # Recast lists as numpy arrays
    for k in vdata.keys():
        vdata[k] = np.array(vdata[k])

    return vdata


"""Position reconstruction for Xenon-nT"""

import os
import tempfile
import tarfile
import numpy as np
import strax
import straxen
from warnings import warn
export, __all__ = strax.exporter()

DEFAULT_POSREC_ALGO_OPTION = tuple([strax.Option("default_reconstruction_algorithm",
                 help="default reconstruction algorithm that provides (x,y)",
                 default="mlp", infer_type=False,
                 )])

@export
@strax.takes_config(
    strax.Option('min_reconstruction_area',
                 help='Skip reconstruction if area (PE) is less than this',
                 default=10, infer_type=False,),
    strax.Option('n_top_pmts', default=straxen.n_top_pmts, infer_type=False,
                 help="Number of top PMTs")
)
class PeakPositionsBaseNT(strax.Plugin):
    """
    Base class for reconstructions.
    This class should only be used when subclassed for the different
    algorithms. Provides x_algorithm, y_algorithm for all peaks > than
    min-reconstruction area based on the top array.
    """
    depends_on = ('peaks',)
    algorithm = None
    compressor = 'zstd'
    # Using parallel = 'process' is not allowed as we cannot Pickle
    # self.model during multiprocessing (to fix?)
    parallel = True
    __version__ = '0.0.0'

    def infer_dtype(self):
        if self.algorithm is None:
            raise NotImplementedError(f'Base class should not be used without '
                                      f'algorithm as done in {__class__.__name__}')
        dtype = [('x_' + self.algorithm, np.float32,
                  f'Reconstructed {self.algorithm} S2 X position (cm), uncorrected'),
                 ('y_' + self.algorithm, np.float32,
                  f'Reconstructed {self.algorithm} S2 Y position (cm), uncorrected')]
        dtype += strax.time_fields
        return dtype

    def setup(self):
        self.model_file = self._get_model_file_name()
        if self.model_file is None:
            warn(f'No file provided for {self.algorithm}. Setting all values '
                 f'for {self.provides} to None.')
            # No further setup required
            return

        # Load the tensorflow model
        import tensorflow as tf
        if os.path.exists(self.model_file):
            print(f"Path is local. Loading {self.algorithm} TF model locally "
                  f"from disk.")
        else:
            downloader = straxen.MongoDownloader()
            try:
                self.model_file = downloader.download_single(self.model_file)
            except straxen.mongo_storage.CouldNotLoadError as e:
                raise RuntimeError(f'Model files {self.model_file} is not found') from e
        with tempfile.TemporaryDirectory() as tmpdirname:
            tar = tarfile.open(self.model_file, mode="r:gz")
            tar.extractall(path=tmpdirname)
            self.model = tf.keras.models.load_model(tmpdirname)

    def compute(self, peaks):
        result = np.ones(len(peaks), dtype=self.dtype)
        result['time'], result['endtime'] = peaks['time'], strax.endtime(peaks)

        result['x_' + self.algorithm] *= float('nan')
        result['y_' + self.algorithm] *= float('nan')

        if self.model_file is None:
            # This plugin is disabled since no model is provided
            return result

        # Keep large peaks only
        peak_mask = peaks['area'] > self.config['min_reconstruction_area']
        if not np.sum(peak_mask):
            # Nothing to do, and .predict crashes on empty arrays
            return result

        # Getting actual position reconstruction
        _in = peaks['area_per_channel'][peak_mask, 0:self.config['n_top_pmts']]
        with np.errstate(divide='ignore', invalid='ignore'):
            _in = _in / np.max(_in, axis=1).reshape(-1, 1)
        _in = _in.reshape(-1, self.config['n_top_pmts'])
        _out = self.model.predict(_in)

        # writing output to the result
        result['x_' + self.algorithm][peak_mask] = _out[:, 0]
        result['y_' + self.algorithm][peak_mask] = _out[:, 1]
        return result

    def _get_model_file_name(self):

        config_file = f'{self.algorithm}_model'
        model_from_config = self.config.get(config_file, 'No file')
        if model_from_config == 'No file':
            raise ValueError(f'{__class__.__name__} should have {config_file} '
                             f'provided as an option.')
        if isinstance(model_from_config, str) and os.path.exists(model_from_config):
            # Allow direct path specification
            return model_from_config
        if model_from_config is None:
            # Allow None to be specified (disables processing for given posrec)
            return model_from_config

        # Use CMT
        model_file = straxen.get_correction_from_cmt(self.run_id, model_from_config)
        return model_file

@export
@strax.takes_config(
    strax.Option('mlp_model',
                 help='Neural network model.'
                      'If CMT, specify as (mlp_model, ONLINE, True)'
                      'Set to None to skip the computation of this plugin.',
                 default=('mlp_model', "ONLINE", True), infer_type=False,
                )
)
class PeakPositionsMLP(PeakPositionsBaseNT):
    """Multilayer Perceptron (MLP) neural net for position reconstruction"""
    provides = "peak_positions_mlp"
    algorithm = "mlp"

@export
@strax.takes_config(
    strax.Option('gcn_model',
                 help='Neural network model.'
                      'If CMT, specify as  (gcn_model, ONLINE, True)'
                      'Set to None to skip the computation of this plugin.',
                 default=('gcn_model', "ONLINE", True), infer_type=False,
                )
)
class PeakPositionsGCN(PeakPositionsBaseNT):
    """Graph Convolutional Network (GCN) neural net for position reconstruction"""
    provides = "peak_positions_gcn"
    algorithm = "gcn"
    __version__ = '0.0.1'

@export
@strax.takes_config(
    strax.Option('cnn_model',
                 help='Neural network model.'
                      'If CMT, specify as (cnn_model, ONLINE, True)'
                      'Set to None to skip the computation of this plugin.',
                 default=('cnn_model', "ONLINE", True), infer_type=False,
                )
)
class PeakPositionsCNN(PeakPositionsBaseNT):
    """Convolutional Neural Network (CNN) neural net for position reconstruction"""
    provides = "peak_positions_cnn"
    algorithm = "cnn"
    __version__ = '0.0.1'

@export
@strax.takes_config(
    *DEFAULT_POSREC_ALGO_OPTION
)
class PeakPositionsNT(strax.MergeOnlyPlugin):
    """
    Merge the reconstructed algorithms of the different algorithms
    into a single one that can be used in Event Basics.

    Select one of the plugins to provide the 'x' and 'y' to be used
    further down the chain. Since we already have the information
    needed here, there is no need to wait until events to make the
    decision.

    Since the computation is trivial as it only combined the three
    input plugins, don't save this plugins output.
    """
    provides = "peak_positions"
    depends_on = ("peak_positions_cnn", "peak_positions_mlp", "peak_positions_gcn")
    save_when = strax.SaveWhen.NEVER
    __version__ = '0.0.0'

    def infer_dtype(self):
        dtype = strax.merged_dtype([self.deps[d].dtype_for(d) for d in self.depends_on])
        dtype += [('x', np.float32, 'Reconstructed S2 X position (cm), uncorrected'),
                  ('y', np.float32, 'Reconstructed S2 Y position (cm), uncorrected')]
        return dtype

    def compute(self, peaks):
        result = {dtype: peaks[dtype] for dtype in peaks.dtype.names}
        algorithm = self.config['default_reconstruction_algorithm']
        if not 'x_' + algorithm in peaks.dtype.names:
            raise ValueError
        for xy in ('x', 'y'):
            result[xy] = peaks[f'{xy}_{algorithm}']
        return result

@export
@strax.takes_config(
    strax.Option('recon_alg_included', help = 'The list of all reconstruction algorithm considered.',
                 default = ('_mlp', '_gcn', '_cnn'), infer_type=False,
                )
)
class S2ReconPosDiff(strax.Plugin):
    '''
    Plugin that provides position reconstruction difference for S2s in events, see note:
    https://xe1t-wiki.lngs.infn.it/doku.php?id=xenon:shengchao:sr0:reconstruction_quality
    '''

    __version__ = '0.0.3'
    parallel = True
    depends_on = 'event_basics'
    provides = 's2_recon_pos_diff'
    save_when = strax.SaveWhen.EXPLICIT

    def infer_dtype(self):
        dtype = [
        ('s2_recon_avg_x', np.float32,
         'Mean value of x for main S2'),
        ('alt_s2_recon_avg_x', np.float32,
         'Mean value of x for alternatice S2'),
        ('s2_recon_avg_y', np.float32,
         'Mean value of y for main S2'),
        ('alt_s2_recon_avg_y', np.float32,
         'Mean value of y for alternatice S2'),
        ('s2_recon_pos_diff', np.float32,
         'Reconstructed position difference for main S2'),
        ('alt_s2_recon_pos_diff', np.float32,
         'Reconstructed position difference for alternative S2'),
    ]
        dtype += strax.time_fields
        return dtype

    def compute(self, events):

        result = np.zeros(len(events), dtype = self.dtype)
        result['time'] = events['time']
        result['endtime'] = strax.endtime(events)
        # Computing position difference
        self.compute_pos_diff(events, result)
        return result

    def cal_avg_and_std(self, values, axis = 1):
        average = np.mean(values, axis = axis)
        std = np.std(values, axis = axis)
        return average, std

    def eval_recon(self, data, name_x_list, name_y_list):
        """
        This function reads the name list based on s2/alt_s2 and all recon algorithm registered
        Each row consists the reconstructed x/y and their average and standard deviation is calculated
        """
        x_avg, x_std = self.cal_avg_and_std(np.array(data[name_x_list].tolist())) #lazy fix to delete field name in array, otherwise np.mean will complain
        y_avg, y_std = self.cal_avg_and_std(np.array(data[name_y_list].tolist()))
        r_std = np.sqrt(x_std**2 + y_std**2)
        res = x_avg, y_avg, r_std
        return res

    def compute_pos_diff(self, events, result):

        alg_list = self.config['recon_alg_included']
        for peak_type in ['s2', 'alt_s2']:
            # Selecting S2s for pos diff
            # - must exist (index != -1)
            # - must have positive AFT
            # - must contain all alg info
            cur_s2_bool = (events[peak_type + '_index'] !=- 1)
            cur_s2_bool &= (events[peak_type + '_area_fraction_top'] > 0)
            for name in self.config['recon_alg_included']:
                cur_s2_bool &= ~np.isnan(events[peak_type+'_x'+name])
                cur_s2_bool &= ~np.isnan(events[peak_type+'_y'+name])

            # default value is nan, it will be ovewrite if the event satisfy the requirments
            result[peak_type + '_recon_pos_diff'][:] = np.nan
            result[peak_type + '_recon_avg_x'][:] = np.nan
            result[peak_type + '_recon_avg_y'][:] = np.nan

            if np.any(cur_s2_bool):
                name_x_list = []
                name_y_list = []
                for alg in alg_list:
                    name_x_list.append(peak_type + '_x' + alg)
                    name_y_list.append(peak_type + '_y' + alg)

                # Calculating average x,y, and position difference
                x_avg, y_avg, r_std = self.eval_recon(events[cur_s2_bool], name_x_list, name_y_list)
                result[peak_type + '_recon_pos_diff'][cur_s2_bool] = r_std
                result[peak_type + '_recon_avg_x'][cur_s2_bool] = x_avg
                result[peak_type + '_recon_avg_y'][cur_s2_bool] = y_avg


# Generate angular tensors using symbolic expressions for the GTO's and derivatives

# There are two steps to the code generation, and only the second is automated.
#  1. Use read_order.py to generate:
#      A. A python version of get_ijk, which is pasted into this file.
#      B. A C++ version of get_ABC, which is pasted into CartesianTensor.h.in
#      C. Repeat B for a version of get_ABC to be pasted into SoaCartesianTensor.h.in
#  2. Run create_cartesian_tensor.py
#      A. Copy CartesianTensor.h up on directory (to src/Numerics)
#      B. Copy SoaCartesianTensor.h up to src/QMCWaveFunctions/LCAO)
#      C. Apply clang-format on modified source files.

from collections import namedtuple, defaultdict
from sympy import *

# See the GaussianOrbitals notebook in the qmc_algorithms repo for more explanation,
#  especially about the normalization.

def create_gto_symbolic():
    pre = sympify('x**i * y**j * z**k * exp(-alpha *r**2)')
    return pre

# this function generated from 'read_order.py', should match GAMESS order
def get_ijk():
  ijk = []
  # S
  ijk.append( (0,0,0,"S") )
  # P
  ijk.append( (1,0,0,"X") )
  ijk.append( (0,1,0,"Y") )
  ijk.append( (0,0,1,"Z") )
  # D
  ijk.append( (2,0,0,"XX") )
  ijk.append( (0,2,0,"YY") )
  ijk.append( (0,0,2,"ZZ") )
  ijk.append( (1,1,0,"XY") )
  ijk.append( (1,0,1,"XZ") )
  ijk.append( (0,1,1,"YZ") )
  # F
  ijk.append( (3,0,0,"XXX") )
  ijk.append( (0,3,0,"YYY") )
  ijk.append( (0,0,3,"ZZZ") )
  ijk.append( (2,1,0,"XXY") )
  ijk.append( (2,0,1,"XXZ") )
  ijk.append( (1,2,0,"YYX") )
  ijk.append( (0,2,1,"YYZ") )
  ijk.append( (1,0,2,"ZZX") )
  ijk.append( (0,1,2,"ZZY") )
  ijk.append( (1,1,1,"XYZ") )
  # G
  ijk.append( (4,0,0,"XXXX") )
  ijk.append( (0,4,0,"YYYY") )
  ijk.append( (0,0,4,"ZZZZ") )
  ijk.append( (3,1,0,"XXXY") )
  ijk.append( (3,0,1,"XXXZ") )
  ijk.append( (1,3,0,"YYYX") )
  ijk.append( (0,3,1,"YYYZ") )
  ijk.append( (1,0,3,"ZZZX") )
  ijk.append( (0,1,3,"ZZZY") )
  ijk.append( (2,2,0,"XXYY") )
  ijk.append( (2,0,2,"XXZZ") )
  ijk.append( (0,2,2,"YYZZ") )
  ijk.append( (2,1,1,"XXYZ") )
  ijk.append( (1,2,1,"YYXZ") )
  ijk.append( (1,1,2,"ZZXY") )
  # H
  ijk.append( (5,0,0,"XXXXX") )
  ijk.append( (0,5,0,"YYYYY") )
  ijk.append( (0,0,5,"ZZZZZ") )
  ijk.append( (4,1,0,"XXXXY") )
  ijk.append( (4,0,1,"XXXXZ") )
  ijk.append( (1,4,0,"YYYYX") )
  ijk.append( (0,4,1,"YYYYZ") )
  ijk.append( (1,0,4,"ZZZZX") )
  ijk.append( (0,1,4,"ZZZZY") )
  ijk.append( (3,2,0,"XXXYY") )
  ijk.append( (3,0,2,"XXXZZ") )
  ijk.append( (2,3,0,"YYYXX") )
  ijk.append( (0,3,2,"YYYZZ") )
  ijk.append( (2,0,3,"ZZZXX") )
  ijk.append( (0,2,3,"ZZZYY") )
  ijk.append( (3,1,1,"XXXYZ") )
  ijk.append( (1,3,1,"YYYXZ") )
  ijk.append( (1,1,3,"ZZZXY") )
  ijk.append( (2,2,1,"XXYYZ") )
  ijk.append( (2,1,2,"XXZZY") )
  ijk.append( (1,2,2,"YYZZX") )
  # I
  ijk.append( (6,0,0,"X6") )
  ijk.append( (0,6,0,"Y6") )
  ijk.append( (0,0,6,"Z6") )
  ijk.append( (5,1,0,"X5Y") )
  ijk.append( (5,0,1,"X5Z") )
  ijk.append( (1,5,0,"Y5X") )
  ijk.append( (0,5,1,"Y5Z") )
  ijk.append( (1,0,5,"Z5X") )
  ijk.append( (0,1,5,"Z5Y") )
  ijk.append( (4,2,0,"X4Y2") )
  ijk.append( (4,0,2,"X4Z2") )
  ijk.append( (2,4,0,"Y4X2") )
  ijk.append( (0,4,2,"Y4Z2") )
  ijk.append( (2,0,4,"Z4X2") )
  ijk.append( (0,2,4,"Z4Y2") )
  ijk.append( (4,1,1,"X4YZ") )
  ijk.append( (1,4,1,"Y4XZ") )
  ijk.append( (1,1,4,"Z4XY") )
  ijk.append( (3,3,0,"X3Y3") )
  ijk.append( (3,0,3,"X3Z3") )
  ijk.append( (0,3,3,"Y3Z3") )
  ijk.append( (3,2,1,"X3Y2Z") )
  ijk.append( (3,1,2,"X3Z2Y") )
  ijk.append( (2,3,1,"Y3X2Z") )
  ijk.append( (1,3,2,"Y3Z2X") )
  ijk.append( (2,1,3,"Z3X2Y") )
  ijk.append( (1,2,3,"Z3Y2X") )
  ijk.append( (2,2,2,"X2Y2Z2") )

  return ijk

#  Input is a list of i,j,k,s ( descriptive string )
#  Output is a list that adds the maximum L (sum of i,j,k) up to that point.
def gen_lmax(ijk_list):
  list_with_lmax = []
  lmax = -1
  for i,j,k,s in ijk_list:
    current_l = i+j+k
    if current_l > lmax:
      lmax = current_l
    list_with_lmax.append( (i,j,k,s,lmax) )
  return list_with_lmax

# Replace powers with pre-computed values.
# One reason for this replacement is the current code prints powers as '**' and so is not even valid C++.
#  Using the C printer and printing as 'pow' could lead to poor performance.
# This does a manual common subexpression elimination.  Could imagine using the cse module to do more
# automatically.  The tricky part might be only computing pieces needed for lmax.
def replace_common_subexpressions(expr, slist=None):
  x,y,z = symbols('x y z')
  x2,y2,z2 = symbols('x2 y2 z2')
  x3,y3,z3 = symbols('x3 y3 z3')
  x4,y4,z4 = symbols('x4 y4 z4')
  x5,y5,z5 = symbols('x5 y5 z5')

  rlist0 = {x**6:x*x5, y**6:y*y5, z**6:z*z5}
  rlist1 = {x**5:x*x4, y**5:y*y4, z**5:z*z4}
  rlist2 = {x**4:x4, y**4:y4, z**4: z4}
  rlist3 = {x**3:x3, y**3:y3, z**3:z3}
  rlist4 = {x**2:x2, y**2:y2, z**2:z2}

  if slist:
    expr = expr.subs(slist)

  return expr.subs(rlist0).subs(rlist1).subs(rlist2).subs(rlist3).subs(rlist4)

def gen_evaluate():
  out_str = """
template<class T, class Point_t, class Tensor_t, class GGG_t>
void CartesianTensor<T,Point_t, Tensor_t, GGG_t>::evaluate(const Point_t& p)
{
  value_type x=p[0], y=p[1], z=p[2];
  value_type x2=x*x, y2=y*y, z2=z*z;
  value_type x3=x2*x, y3=y2*y, z3=z2*z;
  value_type x4=x3*x, y4=y3*y, z4=z3*z;
  value_type x5=x4*x, y5=y4*y, z5=z4*z;
  switch(Lmax)
  {
%s
  }
  for (int i=0; i<XYZ.size(); i++)
    XYZ[i]*= NormFactor[i];
}
"""
  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  # put index and values in a list so it can be reversed
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}
    val = replace_common_subexpressions(gto_s, slist)
    body_str += "    XYZ[%d] = %s; // %s\n"%(idx,val,s)

  return out_str%body_str

def gen_evaluate_all():
  out_str = """
template<class T, class Point_t, class Tensor_t, class GGG_t>
void CartesianTensor<T,Point_t, Tensor_t, GGG_t>::evaluateAll(const Point_t& p)
{
  value_type x=p[0], y=p[1], z=p[2];
  value_type x2=x*x, y2=y*y, z2=z*z;
  value_type x3=x2*x, y3=y2*y, z3=z2*z;
  value_type x4=x3*x, y4=y3*y, z4=z3*z;
  value_type x5=x4*x, y5=y4*y, z5=z4*z;
  int ntot=XYZ.size();
  for (int i=0; i<ntot; i++)
    gradXYZ[i]=0.0;
  for (int i=0; i<ntot; i++)
    laplXYZ[i]=0.0;

  switch(Lmax)
  {
%s
  }
  for (int i=0; i<ntot; i++)
    XYZ[i]*= NormFactor[i];
  for (int i=0; i<ntot; i++)
    gradXYZ[i]*= NormFactor[i];
  for (int i=0; i<ntot; i++)
    laplXYZ[i]*= NormFactor[i];

}
"""
  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  # put index and values in a list so it can be reversed
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  x,y,z = symbols('x y z')

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    # Compute derivatives symbolically
    dx = diff(gto_s, x)
    dy = diff(gto_s, y)
    dz = diff(gto_s, z)
    lap = diff(gto_s, x, 2) + diff(gto_s, y, 2) + diff(gto_s, z, 2)

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}
    val = replace_common_subexpressions(gto_s, slist)
    gx = replace_common_subexpressions(dx, slist)
    gy = replace_common_subexpressions(dy, slist)
    gz = replace_common_subexpressions(dz, slist)
    lap_val = replace_common_subexpressions(lap, slist)

    body_str += "    XYZ[%d] = %s;     // %s\n"%(idx,val,s)

    if gx != 0:
      body_str += "    gradXYZ[%d][0] = %s;\n"%(idx,gx)
    if gy != 0:
      body_str += "    gradXYZ[%d][1] = %s;\n"%(idx,gy)
    if gz != 0:
      body_str += "    gradXYZ[%d][2] = %s;\n"%(idx,gz)

    if lap_val != 0:
      body_str += "    laplXYZ[%d] = %s;\n"%(idx,lap_val)

  return out_str%body_str

def gen_evaluate_with_hessian():
  out_str = """
template<class T, class Point_t, class Tensor_t, class GGG_t>
void CartesianTensor<T,Point_t, Tensor_t, GGG_t>::evaluateWithHessian(const Point_t& p)
{
  value_type x=p[0], y=p[1], z=p[2];
  value_type x2=x*x, y2=y*y, z2=z*z;
  value_type x3=x2*x, y3=y2*y, z3=z2*z;
  value_type x4=x3*x, y4=y3*y, z4=z3*z;
  value_type x5=x4*x, y5=y4*y, z5=z4*z;
  int ntot=XYZ.size();
  for (int i=0; i<ntot; i++)
    gradXYZ[i]=0.0;
  for (int i=0; i<ntot; i++)
    hessXYZ[i]=0.0;

  switch(Lmax)
  {
%s
  }
  for (int i=0; i<ntot; i++)
    XYZ[i]*= NormFactor[i];
  for (int i=0; i<ntot; i++)
    gradXYZ[i]*= NormFactor[i];
  for (int i=0; i<ntot; i++)
    hessXYZ[i]*= NormFactor[i];
}
"""

  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  x,y,z = symbols('x y z')

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}

    # Compute derivatives symbolically
    dx = diff(gto_s, x)
    dy = diff(gto_s, y)
    dz = diff(gto_s, z)

    val = replace_common_subexpressions(gto_s, slist)
    body_str += "    XYZ[%d] = %s;     // %s\n"%(idx,val,s)

    gx = replace_common_subexpressions(dx, slist)
    gy = replace_common_subexpressions(dy, slist)
    gz = replace_common_subexpressions(dz, slist)
    if gx != 0:
      body_str += "    gradXYZ[%d][0] = %s;\n"%(idx,gx)
    if gy != 0:
      body_str += "    gradXYZ[%d][1] = %s;\n"%(idx,gy)
    if gz != 0:
      body_str += "    gradXYZ[%d][2] = %s;\n"%(idx,gz)

    axis_syms = [Symbol('x'), Symbol('y'), Symbol('z')]
    for ii,si in enumerate(axis_syms):
      for jj,sj in enumerate(axis_syms):
        # Compute Hessian elements symbolically
        h_s = diff(diff(gto_s, si), sj)
        hess_val = replace_common_subexpressions(h_s, slist)
        if hess_val != 0:
          body_str += "    hessXYZ[%d](%d,%d) = %s;\n"%(idx,ii,jj,hess_val)

  return out_str%body_str

def gen_evaluate_with_third_deriv():
  out_str = """
template<class T, class Point_t, class Tensor_t, class GGG_t>
void CartesianTensor<T,Point_t, Tensor_t, GGG_t>::evaluateWithThirdDeriv(const Point_t& p)
{
  value_type x=p[0], y=p[1], z=p[2];
  value_type x2=x*x, y2=y*y, z2=z*z;
  value_type x3=x2*x, y3=y2*y, z3=z2*z;
  value_type x4=x3*x, y4=y3*y, z4=z3*z;
  value_type x5=x4*x, y5=y4*y, z5=z4*z;

  int ntot=XYZ.size();
  for (int i=0; i<ntot; i++)
    gradXYZ[i]=0.0;
  for (int i=0; i<ntot; i++)
    hessXYZ[i]=0.0;
  for (int i=0; i<ntot; i++)
  {
    gggXYZ[i][0]=0.0;
    gggXYZ[i][1]=0.0;
    gggXYZ[i][2]=0.0;
  }

  switch(Lmax)
  {
%s
  }

  for (int i=0; i<ntot; i++)
    XYZ[i]*= NormFactor[i];
  for (int i=0; i<ntot; i++)
    gradXYZ[i]*= NormFactor[i];
  for (int i=0; i<ntot; i++)
    hessXYZ[i]*= NormFactor[i];
  for (int i=0; i<ntot; i++)
  {
    gggXYZ[i][0] *= NormFactor[i];
    gggXYZ[i][1] *= NormFactor[i];
    gggXYZ[i][2] *= NormFactor[i];
  }
}

"""

  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  # put index and values in a list so it can be reversed
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  x,y,z = symbols('x y z')

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}

    dx = diff(gto_s, x)
    dy = diff(gto_s, y)
    dz = diff(gto_s, z)

    val = replace_common_subexpressions(gto_s, slist)
    body_str += "    XYZ[%d] = %s;     // %s\n"%(idx,val,s)

    gx = replace_common_subexpressions(dx, slist)
    gy = replace_common_subexpressions(dy, slist)
    gz = replace_common_subexpressions(dz, slist)
    if gx != 0:
      body_str += "    gradXYZ[%d][0] = %s;\n"%(idx,gx)
    if gy != 0:
      body_str += "    gradXYZ[%d][1] = %s;\n"%(idx,gy)
    if gz != 0:
      body_str += "    gradXYZ[%d][2] = %s;\n"%(idx,gz)

    axis_syms = [Symbol('x'), Symbol('y'), Symbol('z')]
    for ii,si in enumerate(axis_syms):
      for jj,sj in enumerate(axis_syms):
        h_s = diff(diff(gto_s, si), sj)
        hess_val = replace_common_subexpressions(h_s, slist)
        if hess_val != 0:
          body_str += "    hessXYZ[%d](%d,%d) = %s;\n"%(idx,ii,jj,hess_val)

    for ii,si in enumerate(axis_syms):
      for jj,sj in enumerate(axis_syms):
        for kk,sk in enumerate(axis_syms):
          ggg_s = diff(diff(diff(gto_s, si), sj), sk)
          ggg_val = replace_common_subexpressions(ggg_s, slist)

          if ggg_val != 0:
            body_str += "    gggXYZ[%d][%d](%d,%d) = %s;\n"%(idx,ii,jj,kk,ggg_val)

  return out_str%body_str

def gen_evaluate_third_deriv_only():
  out_str = """
template<class T, class Point_t, class Tensor_t, class GGG_t>
void CartesianTensor<T,Point_t, Tensor_t, GGG_t>::evaluateThirdDerivOnly(const Point_t& p)
{
  value_type x=p[0], y=p[1], z=p[2];
  value_type x2=x*x, y2=y*y, z2=z*z;
  value_type x3=x2*x, y3=y2*y, z3=z2*z;
  int ntot=XYZ.size();
  for (int i=0; i<ntot; i++)
  {
    gggXYZ[i][0]=0.0;
    gggXYZ[i][1]=0.0;
    gggXYZ[i][2]=0.0;
  }

  switch(Lmax)
  {
%s
  }

  for (int i=0; i<ntot; i++)
  {
    gggXYZ[i][0] *= NormFactor[i];
    gggXYZ[i][1] *= NormFactor[i];
    gggXYZ[i][2] *= NormFactor[i];
  }
}

"""

  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  # put index and values in a list so it can be reversed
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  x,y,z = symbols('x y z')

  case_has_content = False
  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      body_str += "  case %d:\n"%lmax
      if not case_has_content and curr_lmax != -1:
        body_str += "        ; // empty statement\n"
      curr_lmax = lmax
      case_has_content = False

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}

    axis_syms = [Symbol('x'), Symbol('y'), Symbol('z')]
    for ii,si in enumerate(axis_syms):
      for jj,sj in enumerate(axis_syms):
        for kk,sk in enumerate(axis_syms):
          ggg_s = diff(diff(diff(gto_s, si), sj), sk)
          ggg_val = replace_common_subexpressions(ggg_s, slist)

          if ggg_val != 0:
            body_str += "    gggXYZ[%d][%d](%d,%d) = %s;\n"%(idx,ii,jj,kk,ggg_val)
            case_has_content = True

  return out_str%body_str

def gen_soa_evaluate_bare():
  out_str = """
template<class T>
void SoaCartesianTensor<T>::evaluate_bare(T x, T y, T z, T* restrict XYZ) const
{
  const T x2=x*x, y2=y*y, z2=z*z;
  const T x3=x2*x, y3=y2*y, z3=z2*z;
  const T x4=x3*x, y4=y3*y, z4=z3*z;
  const T x5=x4*x, y5=y4*y, z5=z4*z;
  switch(Lmax)
  {
%s
  }
}
"""
  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  # put index and values in a list so it can be reversed
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}
    val = replace_common_subexpressions(gto_s, slist)
    body_str += "    XYZ[%d] = %s; // %s\n"%(idx,val,s)

  return out_str%body_str

def gen_soa_evaluate_vgl():
  out_str = """
template<class T>
void SoaCartesianTensor<T>::evaluateVGL(T x, T y, T z)
{

  constexpr T czero(0);
  cXYZ=czero;

  const T x2=x*x, y2=y*y, z2=z*z;
  const T x3=x2*x, y3=y2*y, z3=z2*z;
  const T x4=x3*x, y4=y3*y, z4=z3*z;
  const T x5=x4*x, y5=y4*y, z5=z4*z;
  T* restrict XYZ=cXYZ.data(0);
  T* restrict gr0=cXYZ.data(1);
  T* restrict gr1=cXYZ.data(2);
  T* restrict gr2=cXYZ.data(3);
  T* restrict lap=cXYZ.data(4);

  switch(Lmax)
  {
%s
  }

  const size_t ntot=NormFactor.size();
  for (size_t i=0; i<ntot; i++)
  {
    XYZ[i]*= NormFactor[i];
    gr0[i]*= NormFactor[i];
    gr1[i]*= NormFactor[i];
    gr2[i]*= NormFactor[i];
    lap[i]*= NormFactor[i];
  }
}
"""
  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  # put index and values in a list so it can be reversed
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  x,y,z = symbols('x y z')

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    # Compute derivatives symbolically
    dx = diff(gto_s, x)
    dy = diff(gto_s, y)
    dz = diff(gto_s, z)
    lap = diff(gto_s, x, 2) + diff(gto_s, y, 2) + diff(gto_s, z, 2)

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}
    val = replace_common_subexpressions(gto_s, slist)
    gx = replace_common_subexpressions(dx, slist)
    gy = replace_common_subexpressions(dy, slist)
    gz = replace_common_subexpressions(dz, slist)
    lap_val = replace_common_subexpressions(lap, slist)

    body_str += "    XYZ[%d] = %s;     // %s\n"%(idx,val,s)

    if gx != 0:
      body_str += "    gr0[%d] = %s;\n"%(idx,gx)
    if gy != 0:
      body_str += "    gr1[%d] = %s;\n"%(idx,gy)
    if gz != 0:
      body_str += "    gr2[%d] = %s;\n"%(idx,gz)

    if lap_val != 0:
      body_str += "    lap[%d] = %s;\n"%(idx,lap_val)

  return out_str%body_str

def gen_soa_evaluate_vgh():
  out_str = """
template<class T>
void SoaCartesianTensor<T>::evaluateVGH(T x, T y, T z)
{
  constexpr T czero(0);
  cXYZ=czero;

  const T x2=x*x, y2=y*y, z2=z*z;
  const T x3=x2*x, y3=y2*y, z3=z2*z;
  const T x4=x3*x, y4=y3*y, z4=z3*z;
  const T x5=x4*x, y5=y4*y, z5=z4*z;

  T* restrict XYZ=cXYZ.data(0);
  T* restrict gr0=cXYZ.data(1);
  T* restrict gr1=cXYZ.data(2);
  T* restrict gr2=cXYZ.data(3);
  T* restrict h00=cXYZ.data(4);
  T* restrict h01=cXYZ.data(5);
  T* restrict h02=cXYZ.data(6);
  T* restrict h11=cXYZ.data(7);
  T* restrict h12=cXYZ.data(8);
  T* restrict h22=cXYZ.data(9);

  switch(Lmax)
  {
%s
  }

  const size_t ntot=cXYZ.size();
  for(size_t i=0; i<ntot; ++i)
  {
    XYZ[i]*= NormFactor[i];
    gr0[i]*= NormFactor[i];
    gr1[i]*= NormFactor[i];
    gr2[i]*= NormFactor[i];
    h00[i]*= NormFactor[i];
    h01[i]*= NormFactor[i];
    h02[i]*= NormFactor[i];
    h11[i]*= NormFactor[i];
    h12[i]*= NormFactor[i];
    h22[i]*= NormFactor[i];
  }

}
"""

  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  x,y,z = symbols('x y z')

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}

    # Compute derivatives symbolically
    dx = diff(gto_s, x)
    dy = diff(gto_s, y)
    dz = diff(gto_s, z)

    val = replace_common_subexpressions(gto_s, slist)
    body_str += "    XYZ[%d] = %s;     // %s\n"%(idx,val,s)

    gx = replace_common_subexpressions(dx, slist)
    gy = replace_common_subexpressions(dy, slist)
    gz = replace_common_subexpressions(dz, slist)
    if gx != 0:
      body_str += "    gr0[%d] = %s;\n"%(idx,gx)
    if gy != 0:
      body_str += "    gr1[%d] = %s;\n"%(idx,gy)
    if gz != 0:
      body_str += "    gr2[%d] = %s;\n"%(idx,gz)

    axis_syms = [Symbol('x'), Symbol('y'), Symbol('z')]
    for ii,si in enumerate(axis_syms):
      for jj,sj in enumerate(axis_syms):
        if ii <= jj:
          # Compute Hessian elements symbolically
          h_s = diff(diff(gto_s, si), sj)
          hess_val = replace_common_subexpressions(h_s, slist)
          if hess_val != 0 :
            body_str += "    h%d%d[%d] = %s;\n"%(ii,jj,idx,hess_val)

  return out_str%body_str

def gen_soa_evaluate_vghgh():
  out_str = """
template<class T>
void SoaCartesianTensor<T>::evaluateVGHGH(T x, T y, T z)
{
  constexpr T czero(0);
  cXYZ=czero;

  const T x2=x*x, y2=y*y, z2=z*z;
  const T x3=x2*x, y3=y2*y, z3=z2*z;
  const T x4=x3*x, y4=y3*y, z4=z3*z;
  const T x5=x4*x, y5=y4*y, z5=z4*z;

  T* restrict XYZ   = cXYZ.data(0);
  T* restrict gr0   = cXYZ.data(1);
  T* restrict gr1   = cXYZ.data(2);
  T* restrict gr2   = cXYZ.data(3);
  T* restrict h00   = cXYZ.data(4);
  T* restrict h01   = cXYZ.data(5);
  T* restrict h02   = cXYZ.data(6);
  T* restrict h11   = cXYZ.data(7);
  T* restrict h12   = cXYZ.data(8);
  T* restrict h22   = cXYZ.data(9);
  T* restrict gh000 = cXYZ.data(10);
  T* restrict gh001 = cXYZ.data(11);
  T* restrict gh002 = cXYZ.data(12);
  T* restrict gh011 = cXYZ.data(13);
  T* restrict gh012 = cXYZ.data(14);
  T* restrict gh022 = cXYZ.data(15);
  T* restrict gh111 = cXYZ.data(16);
  T* restrict gh112 = cXYZ.data(17);
  T* restrict gh122 = cXYZ.data(18);
  T* restrict gh222 = cXYZ.data(19);

  switch(Lmax)
  {
%s
  }

  const size_t ntot=cXYZ.size();
  for(size_t i=0; i<ntot; ++i)
  {
    XYZ[i]   *= NormFactor[i];
    gr0[i]   *= NormFactor[i];
    gr1[i]   *= NormFactor[i];
    gr2[i]   *= NormFactor[i];
    h00[i]   *= NormFactor[i];
    h01[i]   *= NormFactor[i];
    h02[i]   *= NormFactor[i];
    h11[i]   *= NormFactor[i];
    h12[i]   *= NormFactor[i];
    h22[i]   *= NormFactor[i];
    gh000[i] *= NormFactor[i];
    gh001[i] *= NormFactor[i];
    gh002[i] *= NormFactor[i];
    gh011[i] *= NormFactor[i];
    gh012[i] *= NormFactor[i];
    gh022[i] *= NormFactor[i];
    gh111[i] *= NormFactor[i];
    gh112[i] *= NormFactor[i];
    gh122[i] *= NormFactor[i];
    gh222[i] *= NormFactor[i];
  }

}
"""

  gto_s = create_gto_symbolic()
  # just the 'angular' part
  gto_s = gto_s.subs(Symbol('alpha'),0)

  ijk_with_lmax = gen_lmax(get_ijk())
  body_str = ''
  curr_lmax = -1
  ijk_l = [(idx,c) for idx,c in enumerate(ijk_with_lmax)]

  x,y,z = symbols('x y z')

  for idx, (i,j,k,s,lmax) in reversed(ijk_l):
    if lmax != curr_lmax:
      curr_lmax = lmax
      body_str += "  case %d:\n"%curr_lmax

    slist = {Symbol('i'):i, Symbol('j'):j, Symbol('k'):k}

    # Compute derivatives symbolically
    dx = diff(gto_s, x)
    dy = diff(gto_s, y)
    dz = diff(gto_s, z)

    val = replace_common_subexpressions(gto_s, slist)
    body_str += "    XYZ[%d] = %s;     // %s\n"%(idx,val,s)

    gx = replace_common_subexpressions(dx, slist)
    gy = replace_common_subexpressions(dy, slist)
    gz = replace_common_subexpressions(dz, slist)
    if gx != 0:
      body_str += "    gr0[%d] = %s;\n"%(idx,gx)
    if gy != 0:
      body_str += "    gr1[%d] = %s;\n"%(idx,gy)
    if gz != 0:
      body_str += "    gr2[%d] = %s;\n"%(idx,gz)

    axis_syms = [Symbol('x'), Symbol('y'), Symbol('z')]
    for ii,si in enumerate(axis_syms):
      for jj,sj in enumerate(axis_syms):
        if ii <= jj:
          # Compute Hessian elements symbolically
          h_s = diff(diff(gto_s, si), sj)
          hess_val = replace_common_subexpressions(h_s, slist)
          if hess_val != 0 :
            body_str += "    h%d%d[%d] = %s;\n"%(ii,jj,idx,hess_val)

    for ii,si in enumerate(axis_syms):
      for jj,sj in enumerate(axis_syms):
         for kk,sk in enumerate(axis_syms):
           if ii <= jj and jj <= kk:
           # Compute Grad Hessian elements symbolically
             ghess_s = diff(diff(diff(gto_s, si), sj),sk)
             ghess_val = replace_common_subexpressions(ghess_s, slist)
             if ghess_val != 0 :
               body_str += "    gh%d%d%d[%d] = %s;\n"%(ii,jj,kk,idx,ghess_val)

  return out_str%body_str

# A simple template replacement engine.
# Template items to be replaced start on a line with '%'.
def run_template(fname_in, fname_out, bodies):
  out = ''
  with open(fname_in, 'r') as f:
    for line in f:
      if line.startswith('%'):
        key = line.strip()[1:]
        if key in bodies:
          line = bodies[key]
        else:
          print 'Error, template item not found, key:',key, ' line = ',line
      out += line

  with open(fname_out, 'w') as f:
    f.write(out)

dire_codegen_text = """
/*
 DO NOT MAKE PERMANENT EDITS IN THIS FILE
 This file is generated from src/Numerics/codegen/%(script_name)s and %(template_file_name)s

 Edit %(template_file_name)s, rerun %(script_name)s, and copy the generated file here.
*/
"""

def create_cartesian_tensor_h():
  bodies = dict()
  bodies['evaluate'] = gen_evaluate()
  bodies['evaluate_all'] = gen_evaluate_all()
  bodies['evaluate_with_hessian'] = gen_evaluate_with_hessian()
  bodies['evaluate_with_third_deriv'] = gen_evaluate_with_third_deriv()
  bodies['evaluate_third_deriv_only'] = gen_evaluate_third_deriv_only()
  fname_out = 'CartesianTensor.h'
  fname_in= 'CartesianTensor.h.in'

  bodies['dire_codegen_warning'] = dire_codegen_text%({'script_name':'gen_cartesian_tensor.py', 'template_file_name':fname_in})

  run_template(fname_in, fname_out, bodies)

def create_soa_cartesian_tensor_h():
  bodies = dict()
  bodies['evaluate_bare'] = gen_soa_evaluate_bare()
  bodies['evaluate_vgl'] = gen_soa_evaluate_vgl()
  bodies['evaluate_vgh'] = gen_soa_evaluate_vgh()
  bodies['evaluate_vghgh'] = gen_soa_evaluate_vghgh()
  fname_in= 'SoaCartesianTensor.h.in'
  fname_out = 'SoaCartesianTensor.h'

  bodies['dire_codegen_warning'] = dire_codegen_text%({'script_name':'gen_cartesian_tensor.py', 'template_file_name':fname_in})

  run_template(fname_in, fname_out, bodies)

if __name__ == '__main__':
    #print gen_evaluate()
    #print gen_evaluate_all()
    #print gen_evaluate_with_hessian()
    #print gen_evaluate_with_third_deriv()
    #print gen_evaluate_third_deriv_only()

    # Create CartesianTensor.h from CartesianTensor.h.in
    create_cartesian_tensor_h()

    # Create SoaCartesianTensor.h from SoaCartesianTensor.h.in
    create_soa_cartesian_tensor_h()


import os
import warnings
import numpy as np
from pyrates.utility.genetic_algorithm import CGSGeneticAlgorithm
from pandas import DataFrame, read_hdf
from copy import deepcopy

class CustomGOA(CGSGeneticAlgorithm):

    def eval_fitness(self, target: list, **kwargs):

        # define simulation conditions
        worker_file = self.cgs_config['worker_file'] if 'worker_file' in self.cgs_config else None
        param_grid = self.pop.drop(['fitness', 'sigma', 'results'], axis=1)
        result_vars = ['r_e', 'r_p', 'r_a', 'r_m', 'r_f']
        freq_targets = [0.0, np.nan, np.nan, np.nan, np.nan]
        #param_grid, invalid_params = eval_params(param_grid)
        conditions = [{},  # healthy control
                      {'k_pe': 0.2, 'k_ae': 0.2},  # AMPA blockade in GPe
                      {'k_pe': 0.2, 'k_ae': 0.2, 'k_pp': 0.2, 'k_pa': 0.2, 'k_pm': 0.2, 'k_aa': 0.2, 'k_ap': 0.2,
                       'k_am': 0.2},  # AMPA blockade and GABAA blockade in GPe
                      {'k_pp': 0.2, 'k_pa': 0.2, 'k_pm': 0.2, 'k_aa': 0.2, 'k_ap': 0.2,
                       'k_am': 0.2},  # GABAA blockade in GPe
                      {'k_pe': 0.0, 'k_ae': 0.0},  # STN blockade
                      {'k_ep': 0.2},  # GABAA blocker in STN
                      ]
        param_scalings = [
            ('delta_e', 'tau_e', 2.0),
            ('delta_p', 'tau_p', 2.0),
            ('delta_a', 'tau_a', 2.0),
            ('delta_m', 'tau_m', 2.0),
            ('delta_f', 'tau_f', 2.0),
            ('k_ee', 'delta_e', 0.5),
            ('k_ep', 'delta_e', 0.5),
            ('k_pe', 'delta_p', 0.5),
            ('k_pp', 'delta_p', 0.5),
            ('k_pa', 'tau_p', 0.5),
            ('k_pm', 'tau_p', 0.5),
            ('k_ae', 'tau_a', 0.5),
            ('k_ap', 'tau_a', 0.5),
            ('k_aa', 'tau_a', 0.5),
            ('k_am', 'tau_a', 0.5),
            ('k_mf', 'delta_m', 0.5),
            ('k_mm', 'delta_m', 0.5),
            ('k_fa', 'delta_f', 0.5),
            ('k_ff', 'delta_f', 0.5),
            ('eta_e', 'delta_e', 1.0),
            ('eta_p', 'delta_p', 1.0),
            ('eta_a', 'delta_a', 1.0),
            ('eta_m', 'delta_m', 1.0),
            ('eta_f', 'delta_f', 1.0),
        ]
        chunk_size = [
            60,  # carpenters
            100,  # osttimor
            60,  # spanien
            100,  # animals
            60,  # kongo
            60,  # tschad
            #100,  # uganda
            # 50,  # tiber
            #50,  # giraffe
            40,  # lech
            20,  # rilke
            12,  # dinkel
            #10,  # rosmarin
            #10,  # mosambik
            # 50,  # compute servers
            # 40,
            # 30,
            # 20,
            # 10,
            # 50,
            # 40,
            # 30,
            # 20,
            # 10,
            # 50,
            # 40,
            # 30,
            # 20,
            # 10,
            # 50,
            # 40,
        ]

        # perform simulations
        if len(param_grid) > 0:
            self.gs_config['init_kwargs'].update(kwargs)
            res_file = self.cgs.run(
                circuit_template=self.gs_config['circuit_template'],
                param_grid=deepcopy(param_grid),
                param_map=self.gs_config['param_map'],
                simulation_time=self.gs_config['simulation_time'],
                dt=self.gs_config['step_size'],
                inputs=self.gs_config['inputs'],
                outputs=self.gs_config['outputs'],
                sampling_step_size=self.gs_config['sampling_step_size'],
                permute=False,
                chunk_size=chunk_size,
                worker_file=worker_file,
                worker_env=self.cgs_config['worker_env'],
                gs_kwargs={'init_kwargs': self.gs_config['init_kwargs'], 'conditions': conditions,
                           'param_scalings': param_scalings},
                worker_kwargs={'y': target, 'time_lim': 7200.0, 'freq_targets': freq_targets},
                result_concat_axis=0)
            results_tmp = read_hdf(res_file, key=f'Results/results')

            # calculate fitness
            for gene_id in param_grid.index:
                self.pop.at[gene_id, 'fitness'] = 1.0 / results_tmp.at[gene_id, 'fitness']
                self.pop.at[gene_id, 'results'] = [results_tmp.at[gene_id, v] for v in result_vars]

        # set fitness of invalid parametrizations
        #for gene_id in invalid_params.index:
        #    self.pop.at[gene_id, 'fitness'] = 0.0
        #    self.pop.at[gene_id, 'results'] = [0. for _ in result_vars]

def fitness(y, t):
    y = np.asarray(y).flatten()
    t = np.asarray(t).flatten()
    diff = np.asarray([0.0 if np.isnan(t_tmp) else y_tmp - t_tmp for y_tmp, t_tmp in zip(y, t)]).flatten()
    t[np.isnan(t)] = 1.0
    t[t == 0] = 1.0
    weights = 1 / np.abs(t)
    return weights @ np.abs(diff)

if __name__ == "__main__":
    warnings.filterwarnings("ignore")

    pop_size = 1024
    pop_genes = {
        'k_ee': {'min': 0, 'max': 15, 'size': pop_size, 'sigma': 0.1, 'loc': 1.0, 'scale': 0.5},
        'k_ae': {'min': 0, 'max': 150, 'size': pop_size, 'sigma': 0.5, 'loc': 20.0, 'scale': 2.0},
        'k_pe': {'min': 0, 'max': 150, 'size': pop_size, 'sigma': 0.5, 'loc': 20.0, 'scale': 2.0},
        'k_pp': {'min': 0, 'max': 100, 'size': pop_size, 'sigma': 0.5, 'loc': 10.0, 'scale': 1.0},
        'k_ep': {'min': 0, 'max': 150, 'size': pop_size, 'sigma': 0.5, 'loc': 20.0, 'scale': 2.0},
        'k_ap': {'min': 0, 'max': 100, 'size': pop_size, 'sigma': 0.5, 'loc': 10.0, 'scale': 1.0},
        'k_aa': {'min': 0, 'max': 50, 'size': pop_size, 'sigma': 0.5, 'loc': 10.0, 'scale': 1.0},
        'k_pa': {'min': 0, 'max': 50, 'size': pop_size, 'sigma': 0.5, 'loc': 10.0, 'scale': 1.0},
        'k_fa': {'min': 0, 'max': 100, 'size': pop_size, 'sigma': 0.5, 'loc': 20.0, 'scale': 2.0},
        'k_mm': {'min': 0, 'max': 50, 'size': pop_size, 'sigma': 0.5, 'loc': 10.0, 'scale': 1.0},
        'k_am': {'min': 0, 'max': 200, 'size': pop_size, 'sigma': 0.8, 'loc': 40.0, 'scale': 4.0},
        'k_pm': {'min': 0, 'max': 200, 'size': pop_size, 'sigma': 0.5, 'loc': 5.0, 'scale': 1.0},
        'k_mf': {'min': 0, 'max': 150, 'size': pop_size, 'sigma': 0.5, 'loc': 20.0, 'scale': 2.0},
        'k_ff': {'min': 0, 'max': 100, 'size': pop_size, 'sigma': 0.5, 'loc': 10.0, 'scale': 1.0},
        'eta_e': {'min': -5, 'max': 5, 'size': pop_size, 'sigma': 0.2, 'loc': 0.0, 'scale': 0.5},
        'eta_p': {'min': -5, 'max': 5, 'size': pop_size, 'sigma': 0.2, 'loc': 0.0, 'scale': 0.5},
        'eta_a': {'min': -5, 'max': 5, 'size': pop_size, 'sigma': 0.2, 'loc': 0.0, 'scale': 0.5},
        'eta_m': {'min': -10, 'max': 0, 'size': pop_size, 'sigma': 0.2, 'loc': -3.0, 'scale': 0.5},
        'eta_f': {'min': -5, 'max': 5, 'size': pop_size, 'sigma': 0.2, 'loc': 0.0, 'scale': 0.5},
        'delta_e': {'min': 0.01, 'max': 1.0, 'size': pop_size, 'sigma': 0.05, 'loc': 0.1, 'scale': 0.1},
        'delta_p': {'min': 0.01, 'max': 1.0, 'size': pop_size, 'sigma': 0.05, 'loc': 0.2, 'scale': 0.1},
        'delta_a': {'min': 0.01, 'max': 1.5, 'size': pop_size, 'sigma': 0.05, 'loc': 0.4, 'scale': 0.1},
        'delta_m': {'min': 0.01, 'max': 1.5, 'size': pop_size, 'sigma': 0.05, 'loc': 0.2, 'scale': 0.1},
        'delta_f': {'min': 0.01, 'max': 1.5, 'size': pop_size, 'sigma': 0.05, 'loc': 0.2, 'scale': 0.1},
        'tau_e': {'min': 12, 'max': 12, 'size': pop_size, 'sigma': 0.0, 'loc': 12.0, 'scale': 0.0},
        'tau_p': {'min': 24, 'max': 24, 'size': pop_size, 'sigma': 0.0, 'loc': 24.0, 'scale': 0.0},
        'tau_a': {'min': 20, 'max': 20, 'size': pop_size, 'sigma': 0.0, 'loc': 20.0, 'scale': 0.0},
        'tau_m': {'min': 20, 'max': 20, 'size': pop_size, 'sigma': 0.0, 'loc': 20.0, 'scale': 0.0},
        'tau_f': {'min': 20, 'max': 20, 'size': pop_size, 'sigma': 0.0, 'loc': 20.0, 'scale': 0.0},
        #'tau_ee_v': {'min': 0.5, 'max': 1.0, 'size': 2, 'sigma': 0.1, 'loc': 0.5, 'scale': 0.1},
        # 'tau_ei': {'min': 3.0, 'max': 5.0, 'size': 1, 'sigma': 0.1, 'loc': 4.0, 'scale': 0.1},
        #'tau_ei_v': {'min': 0.5, 'max': 1.0, 'size': 2, 'sigma': 0.1, 'loc': 1.0, 'scale': 0.2},
        # 'tau_ie': {'min': 2.0, 'max': 4.0, 'size': 1, 'sigma': 0.1, 'loc': 3.0, 'scale': 0.1},
        #'tau_ie_v': {'min': 0.8, 'max': 1.6, 'size': 2, 'sigma': 0.1, 'loc': 0.7, 'scale': 0.1},
        #'tau_ii_v': {'min': 0.5, 'max': 1.0, 'size': 2, 'sigma': 0.1, 'loc': 0.5, 'scale': 0.1},
    }

    param_map = {
        'k_ee': {'vars': ['weight'], 'edges': [('stn', 'stn')]},
        'k_ae': {'vars': ['weight'], 'edges': [('stn', 'gpe_a')]},
        'k_pe': {'vars': ['weight'], 'edges': [('stn', 'gpe_p')]},
        'k_pp': {'vars': ['weight'], 'edges': [('gpe_p', 'gpe_p')]},
        'k_ep': {'vars': ['weight'], 'edges': [('gpe_p', 'stn')]},
        'k_ap': {'vars': ['weight'], 'edges': [('gpe_p', 'gpe_a')]},
        'k_aa': {'vars': ['weight'], 'edges': [('gpe_a', 'gpe_a')]},
        'k_pa': {'vars': ['weight'], 'edges': [('gpe_a', 'gpe_p')]},
        'k_fa': {'vars': ['weight'], 'edges': [('gpe_a', 'fsi')]},
        'k_mm': {'vars': ['weight'], 'edges': [('msn', 'msn')]},
        'k_am': {'vars': ['weight'], 'edges': [('msn', 'gpe_a')]},
        'k_pm': {'vars': ['weight'], 'edges': [('msn', 'gpe_p')]},
        'k_ff': {'vars': ['weight'], 'edges': [('fsi', 'fsi')]},
        'k_mf': {'vars': ['weight'], 'edges': [('fsi', 'msn')]},
        'eta_e': {'vars': ['stn_op/eta_e'], 'nodes': ['stn']},
        'eta_p': {'vars': ['gpe_proto_op/eta_i'], 'nodes': ['gpe_p']},
        'eta_a': {'vars': ['gpe_arky_op/eta_a'], 'nodes': ['gpe_a']},
        'eta_m': {'vars': ['str_msn_op/eta_s'], 'nodes': ['msn']},
        'eta_f': {'vars': ['str_fsi_op/eta_f'], 'nodes': ['fsi']},
        'delta_e': {'vars': ['stn_op/delta_e'], 'nodes': ['stn']},
        'delta_p': {'vars': ['gpe_proto_op/delta_i'], 'nodes': ['gpe_p']},
        'delta_a': {'vars': ['gpe_arky_op/delta_a'], 'nodes': ['gpe_a']},
        'delta_m': {'vars': ['str_msn_op/delta_s'], 'nodes': ['msn']},
        'delta_f': {'vars': ['str_fsi_op/delta_f'], 'nodes': ['fsi']},
        'tau_e': {'vars': ['stn_op/tau_e'], 'nodes': ['stn']},
        'tau_p': {'vars': ['gpe_proto_op/tau_i'], 'nodes': ['gpe_p']},
        'tau_a': {'vars': ['gpe_arky_op/tau_a'], 'nodes': ['gpe_a']},
        'tau_m': {'vars': ['str_msn_op/tau_s'], 'nodes': ['msn']},
        'tau_f': {'vars': ['str_fsi_op/tau_f'], 'nodes': ['fsi']},
    }

    T = 2000.
    dt = 1e-2
    dts = 1e-1
    compute_dir = f"{os.getcwd()}/stn_gpe_str_opt"

    # perform genetic optimization
    ga = CustomGOA(fitness_measure=fitness,
                   gs_config={
                       'circuit_template': f"{os.getcwd()}/config/stn_gpe/stn_gpe_str",
                       'permute_grid': True,
                       'param_map': param_map,
                       'simulation_time': T,
                       'step_size': dt,
                       'sampling_step_size': dts,
                       'inputs': {},
                       'outputs': {'r_e': "stn/stn_op/R_e", 'r_p': 'gpe_p/gpe_proto_op/R_i',
                                   'r_a': 'gpe_a/gpe_arky_op/R_a', 'r_m': 'msn/str_msn_op/R_s',
                                   'r_f': 'fsi/str_fsi_op/R_f'},
                       'init_kwargs': {'backend': 'numpy', 'solver': 'scipy', 'step_size': dt},
                   },
                   cgs_config={'nodes': [
                       'carpenters',
                       'osttimor',
                       'spanien',
                       'animals',
                       'kongo',
                       'tschad',
                       #'uganda',
                       # 'tiber',
                       #'giraffe',
                       'lech',
                       'rilke',
                       'dinkel',
                       #'rosmarin',
                       #'mosambik',
                       # 'comps06h01',
                       # 'comps06h02',
                       # 'comps06h03',
                       # 'comps06h04',
                       # 'comps06h05',
                       # 'comps06h06',
                       # 'comps06h07',
                       # 'comps06h08',
                       # 'comps06h09',
                       # 'comps06h10',
                       # 'comps06h11',
                       # 'comps06h12',
                       # 'comps06h13',
                       # 'comps06h14',
                       # 'scorpions',
                       # 'spliff',
                       # 'supertramp',
                       # 'ufo'
                   ],
                       'compute_dir': compute_dir,
                       'worker_file': f'{os.getcwd()}/stn_gpe_str_worker.py',
                       'worker_env': "/data/u_rgast_software/anaconda3/envs/pyrates/bin/python3",
                   })

    drop_save_dir = f'{compute_dir}/PopulationDrops/'
    os.makedirs(drop_save_dir, exist_ok=True)

    winner = ga.run(
        initial_gene_pool=pop_genes,
        gene_sampling_func=np.random.normal,
        new_member_sampling_func=np.random.normal,
        target=[[20, 60, 20, 2, 20],  # healthy control
                [np.nan, 2/3, np.nan, np.nan, np.nan],  # ampa blockade in GPe
                [np.nan, 1, np.nan, np.nan, np.nan],  # ampa and gabaa blockade in GPe
                [np.nan, 2, np.nan, np.nan, np.nan],  # GABAA blockade in GPe
                [np.nan, 1/2, np.nan, np.nan, np.nan],  # STN blockade
                [2, 2, np.nan, np.nan, np.nan],  # GABAA blockade in STN
                ],
        max_iter=100,
        enforce_max_iter=True,
        min_fit=1.0,
        n_winners=10,
        n_parent_pairs=40,
        n_new=62,
        sigma_adapt=0.05,
        candidate_save=f'{compute_dir}/GeneticCGSCandidatestn.h5',
        drop_save=drop_save_dir,
        new_pop_on_drop=True,
        pop_save=f'{drop_save_dir}/pop_summary',
        permute=False
    )

    # winner.to_hdf(f'{drop_save_dir}/winner.h5', key='data')


#!/usr/bin/env python
from __future__ import print_function
import numpy as np
from cycler import cycler
import matplotlib.pyplot as plt
import StarKillerMicrophysics as SKM

class BurnerDriver(object):
    def __init__(self, probin_file):
        skinit = SKM.Starkiller_Initialization_Module.starkiller_initialize
        skinit(probin_file)

        self.nspec = SKM.Actual_Network().nspec
        self.short_species_names = [SKM.Network().get_network_short_species_name(i+1).decode("ASCII").strip().lower() for i in range(self.nspec)]
        self.species_names = [SKM.Network().get_network_species_name(i+1).decode("ASCII").strip().lower() for i in range(self.nspec)]

        self.history = History()
        self.plotting = BurnPlotting()

        self.burn_module = SKM.Actual_Burner_Module()
        self.burn_type_module = SKM.Burn_Type_Module()

        self.eos_module = SKM.Eos_Module()
        self.eos_type_module = SKM.Eos_Type_Module()

        self.rhs_module = SKM.actual_rhs_module

        self.initial_burn_state = self.burn_type_module.burn_t()
        self.initial_burn_state.rho = 0.0
        self.initial_burn_state.t = 0.0
        self.initial_burn_state.xn = np.zeros(self.nspec, dtype=np.float64)

        self.end_time = 0.0
        self.num_steps = 0
        self.dt = 0.0

    def list_species(self):
        print("Species in network:\n")
        for sname, short_sname in zip(self.species_names, self.short_species_names):
            print("{} ({})".format(sname, short_sname))

    def set_initial_density(self, dens):
        self.initial_burn_state.rho = dens

    def set_initial_temperature(self, temp):
        self.initial_burn_state.t = temp

    def set_initial_massfractions(self, xn):
        self.initial_burn_state.xn = xn[:]

    def set_initial_species(self, species_name, xspec):
        sname = species_name.lower()
        idx = -1
        try:
            idx = self.short_species_names.index(sname)
        except ValueError:
            try:
                idx = self.species_names.index(sname)
            except ValueError:
                print("ERROR: species {} is not in this network.".format(species_name))
                self.list_species()
                return
        self.initial_burn_state.xn[idx] = xspec

    def get_initial_state(self):
        return self.initial_burn_state

    def burn(self, end_time, num_steps):
        self.history = History()
        self.end_time = end_time
        self.num_steps = num_steps
        self.dt = self.end_time/self.num_steps

        current_time = 0.0

        state_in = self.burn_type_module.burn_t()
        state_out = self.burn_type_module.burn_t()

        self.burn_type_module.copy_burn_t(state_in, self.initial_burn_state)
        self.burn_type_module.copy_burn_t(state_out, self.initial_burn_state)

        for istep in range(self.num_steps):
            self.burn_module.actual_burner(state_in, state_out, self.dt, 0.0)
            current_time = current_time + self.dt
            self.history.store(state_out, current_time, self.dt, istep+1)
            self.burn_type_module.copy_burn_t(state_in, state_out)

        self.plotting.plot_burn_history(self.history)

    def eos(self, input, burn_state):
        eos_state = self.eos_type_module.eos_t()
        self.burn_type_module.burn_to_eos(burn_state, eos_state)
        self.eos_module.eos(input, eos_state)
        self.burn_type_module.eos_to_burn(eos_state, burn_state)

    def rhs(self, burn_state):
        # Call the EOS in (r,t) mode and then evaluate the rhs
        self.eos(self.eos_type_module.eos_input_rt, burn_state)
        self.rhs_module.actual_rhs(burn_state)

    def jac(self, burn_state):
        # Call the EOS in (r,t) mode and then evaluate the jacobian
        self.eos(self.eos_type_module.eos_input_rt, burn_state)
        self.rhs_module.actual_jac(burn_state)

    def save(self, file_name):
        self.history.save(self.species_names, self.initial_burn_state, file_name)

    def get_temp_dot(self, burn_state):
        return burn_state.ydot[-2]

    def get_enuc_dot(self, burn_state):
        return burn_state.ydot[-1]

class History(object):
    def __init__(self):
        self.nspec = SKM.Actual_Network().nspec
        self.xn = [[] for i in range(self.nspec)]
        self.t = []
        self.edot = []
        self.time = []
        self.step = []

    def append_xn(self, xn):
        for i, xi in enumerate(xn):
            self.xn[i].append(xn[i])

    def store(self, state, time, dt, step):
        self.append_xn(state.xn)
        self.t.append(state.t)
        self.edot.append(state.e/dt)
        self.time.append(time)
        self.step.append(step)

    def get_save_string(self, step, time, temp, enuc, xn):
        saveline = "   ".join([str(step), str(time), str(temp), str(enuc)] + [str(xi) for xi in xn])
        saveline = saveline + "\n"
        return saveline

    def get_species_vector(self, ixn):
        xvec = np.array([xi[ixn] for xi in self.xn])
        return xvec

    def save(self, spec_names, initial_state, filename):
        f = open("{}.dat".format(filename), "w")
        spec_names_string = "   ".join(spec_names)
        f.write("step   time   temperature   enucdot   " + spec_names_string + "\n")
        f.write(self.get_save_string(0, 0.0, initial_state.t, 0.0, initial_state.xn))
        for ii, (istep, itime, itemp, ienuc) in enumerate(zip(self.step, self.time, self.t, self.edot)):
            xvec = self.get_species_vector(ii)
            f.write(self.get_save_string(istep, itime, itemp, ienuc, xvec))
        f.close()

class BurnPlotting(object):
    def __init__(self):
        self.nspec = SKM.Actual_Network().nspec
        self.short_species_names = [SKM.Network().get_network_short_species_name(i+1).decode("ASCII").strip() for i in range(self.nspec)]

    def rgba_to_hex(self, rgba):
        r = int(rgba[0]*255.0)
        g = int(rgba[1]*255.0)
        b = int(rgba[2]*255.0)
        return '#{:02X}{:02X}{:02X}'.format(r,g,b)

    def plot_burn_history(self, history, logtime=True):
        plt.clf()

        if logtime:
            xlabel = '$\mathrm{Log_{10}~Time~(s)}$'
            xvec = np.log10(history.time)
            xlim = [np.log10(history.time[0]), np.log10(history.time[-1])]
        else:
            xlabel = '$\mathrm{Time~(s)}$'
            xvec = history.time
            xlim = [history.time[0], history.time[-1]]

        # Get set of colors to use for abundances
        cm = plt.get_cmap('nipy_spectral')
        clist = [cm(1.0*i/self.nspec) for i in range(self.nspec)]
        hexclist = [self.rgba_to_hex(ci) for ci in clist]

        # Get the figure
        fig = plt.figure()
        fig.set_figheight(10.0)
        fig.set_figwidth(5.0)

        # Plot X vs. time
        ax = fig.add_subplot(211)
        ax.set_prop_cycle(cycler('color', hexclist))
        for i in range(self.nspec):
            ax.plot(xvec, np.log10(history.xn[i]), label=self.short_species_names[i])
        lgd = ax.legend(bbox_to_anchor=(1.15, 1.0), loc=2, borderaxespad=0.0)
        ax.set_xlim(xlim)
        plt.setp(ax.get_xticklabels(), visible=False)
        # plt.xlabel(xlabel)
        plt.ylabel('$\mathrm{Log_{10}~X}$')

        # Plot T, edot vs. time
        # Find where edot = 0
        def y_where_x_zero(y, x):
            yzero = []
            xiszero = False
            ylo = 0.0
            for yi, xi in zip(y, x):
                if xi == 0.0 and not xiszero:
                    xiszero = True
                    ylo = yi
                if xi != 0.0 and xiszero:
                    xiszero = False
                    yzero.append([ylo, yi])
            if xiszero:
                yzero.append([ylo, y[-1]])
            return yzero

        edotzero = y_where_x_zero(history.time, history.edot)
        axe = fig.add_subplot(212)
        axe.plot(xvec, np.log10(history.t), label='Temperature', color='red')
        lgd1 = axe.legend(bbox_to_anchor=(0.3, 1.02, 0.7, 1.02), loc='lower left',
                          ncol=1, mode='expand', borderaxespad=0.0)
        axe.set_xlabel(xlabel)
        axe.set_ylabel('$\mathrm{Log_{10}~T~(K)}$')
        axe.set_xlim(xlim)
        ax2 = axe.twinx()
        ax2.plot(xvec, np.log10(history.edot), label='E Gen Rate', color='blue')
        ax2.set_ylabel('$\mathrm{Log_{10}~\\dot{e}~(erg/g/s)}$')
        ax2.set_xlim(xlim)
        lgd2 = ax2.legend(bbox_to_anchor=(0.7, 1.02, 1.0, 1.02), loc='lower left',
                          ncol=1, borderaxespad=0.0)
        # hatch where edot=0
        for edz in edotzero:
            plt.axvspan(np.log10(edz[0]), np.log10(edz[1]), color='blue', fill=False,
                        linewidth=0, hatch='/', alpha=0.2)

        plt.savefig('burn.eps',
                    bbox_extra_artists=(lgd, lgd1, lgd2,), bbox_inches='tight')
        plt.savefig('burn.png', dpi=300,
                    bbox_extra_artists=(lgd, lgd1, lgd2,), bbox_inches='tight')

        plt.show()