"""
Substitution of molecular geometries with functional groups.

Given a cartesian geometry and element labels, a substituent can
be added knowing (a) the substituent identity (b) the desired
position to substitute and (c) the bond axis for substitution.

This requires some default information, such as the default structure
and orientation of the substituent (relative to an axis) and the
bond length of the substituent. For now, only single bonds are treated.
"""
import numpy as np
import gimbal.displace as displace
import gimbal.fileio as fileio
import gimbal.constants as con

class SubLib(object):
    """
    Object containing a library of substituent geometries.

    Attributes
    ----------
    syn : dict
        A dictionary of synonyms used to find the appropriate
        substituents.
    elem : dict
        A dictionary of atomic symbols for each substituent.
    xyz : dict
        A dictionary of atomic cartesian coordinates for each
        substituent.
    """
    def __init__(self):
        self.syn = dict()
        self.elem = dict()
        self.xyz = dict()
        self._populate_syn()
        self._populate_elem()
        self._populate_xyz()
        self._add_comb()

    def _populate_syn(self):
        """Adds a dictionary of synonyms for labels."""
        synlist = [['h'],
                   ['d'],
                   ['me', 'ch3', 'h3c'],
                   ['et', 'ch2ch3', 'c2h5', 'ch3ch2'],
                   ['npr', 'ch2ch2ch3', 'c3h7', 'ch3ch2ch2'],
                   ['ipr', 'chch3ch3', 'ch(ch3)2', '(ch3)2ch', 'ch3chch3'],
                   ['nbu', 'ch2ch2ch2ch3', 'c4h9', 'ch3ch2ch2ch2'],
                   ['ibu', 'chch3ch2ch3', 'ch3chch2ch3', 'ch3ch3chch3'],
                   ['tbu', 'cch3ch3ch3', 'c(ch3)3', 'ch3ch3ch3c', '(ch3)3c'],
                   ['vi', 'chch2', 'c2h3', 'h2chc', 'h3c2'],
                   ['ey', 'cch', 'c2h', 'hcc', 'hc2'],
                   ['ph', 'c6h5', 'h5c6'],
                   ['am', 'nh2', 'h2n'],
                   ['im', 'chnh', 'cnh2', 'nhch'],
                   ['cn', 'nc'],
                   ['oh', 'ho'],
                   ['ome', 'meo', 'och3', 'ch3o'],
                   ['al', 'cho', 'coh', 'och', 'ohc'],
                   ['ac', 'coch3', 'cch3o'],
                   ['ca', 'cooh', 'co2h', 'hooc', 'ho2c'],
                   ['nt', 'no2', 'o2n'],
                   ['f'],
                   ['tfm', 'cf3', 'f3c'],
                   ['sh', 'hs'],
                   ['sf', 'so2h', 'sooh', 'sho2', 'ho2s', 'hso2'],
                   ['ms', 'sfme', 'mesf', 'sfch3', 'so2me', 'so2ch3'],
                   ['cl']]
        for subl in synlist:
            for item in subl:
                self.syn[item] = subl[0]

    def _populate_elem(self):
        """Adds element labels to self.elem."""
        self.elem['h'] = np.array(['H'])
        self.elem['d'] = np.array(['D'])
        self.elem['me'] = np.array(['C', 'H', 'H', 'H'])
        self.elem['vi'] = np.array(['C', 'C', 'H', 'H', 'H'])
        self.elem['ey'] = np.array(['C', 'C', 'H'])
        self.elem['ph'] = np.array(['C', 'C', 'C', 'C', 'C', 'C',
                                    'H', 'H', 'H', 'H', 'H'])
        self.elem['am'] = np.array(['N', 'H', 'H'])
        self.elem['im'] = np.array(['C', 'N', 'H', 'H'])
        self.elem['cn'] = np.array(['C', 'N'])
        self.elem['oh'] = np.array(['O', 'H'])
        self.elem['al'] = np.array(['C', 'O', 'H'])
        self.elem['nt'] = np.array(['N', 'O', 'O'])
        self.elem['f'] = np.array(['F'])
        self.elem['sh'] = np.array(['S', 'H'])
        self.elem['sf'] = np.array(['S', 'O', 'O', 'H'])
        self.elem['cl'] = np.array(['Cl'])

    def _populate_xyz(self):
        """Adds cartesian geometries to self.xyz.

        For all substituents, the bonding atom is at the origin,
        the bonding axis is the z-axis and the plane axis
        is the y-axis.
        """
        self.xyz['h'] = np.array([[ 0.000,  0.000,  0.000]])
        self.xyz['d'] = np.array([[ 0.000,  0.000,  0.000]])
        self.xyz['me'] = np.array([[ 0.000,  0.000,  0.000],
                                   [ 0.511, -0.886,  0.377],
                                   [ 0.511,  0.886,  0.377],
                                   [-1.023,  0.000,  0.377]])
        self.xyz['vi'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-1.124,  0.000,  0.730],
                                   [ 0.971,  0.000,  0.495],
                                   [-1.067,  0.000,  1.818],
                                   [-2.095,  0.000,  0.235]])
        self.xyz['ey'] = np.array([[ 0.000,  0.000,  0.000],
                                   [ 0.000,  0.000,  1.210],
                                   [ 0.000,  0.000,  2.280]])
        self.xyz['ph'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-1.212,  0.000,  0.700],
                                   [ 1.212,  0.000,  0.700],
                                   [-1.212,  0.000,  2.100],
                                   [ 1.212,  0.000,  2.100],
                                   [ 0.000,  0.000,  2.800],
                                   [-2.156,  0.000,  0.155],
                                   [ 2.156,  0.000,  0.155],
                                   [-2.156,  0.000,  2.645],
                                   [ 2.156,  0.000,  2.645],
                                   [ 0.000,  0.000,  3.890]])
        self.xyz['am'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-0.577, -0.771,  0.332],
                                   [-0.577,  0.771,  0.332]])
        self.xyz['im'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-1.082,  0.000,  0.703],
                                   [ 0.980,  0.000,  0.499],
                                   [-0.869,  0.000,  1.710]])
        self.xyz['cn'] = np.array([[ 0.000,  0.000,  0.000],
                                   [ 0.000,  0.000,  1.136]])
        self.xyz['oh'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-0.913,  0.000,  0.297]])
        self.xyz['al'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-1.011,  0.000,  0.700],
                                   [ 0.998,  0.000,  0.463]])
        self.xyz['nt'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-1.105,  0.000,  0.563],
                                   [ 1.105,  0.000,  0.563]])
        self.xyz['f'] = np.array([[ 0.000,  0.000,  0.000]])
        self.xyz['sh'] = np.array([[ 0.000,  0.000,  0.000],
                                   [-1.331,  0.000,  0.156]])
        self.xyz['sf'] = np.array([[ 0.000,  0.000,  0.000],
                                   [ 0.548, -1.266,  0.448],
                                   [ 0.548,  1.266,  0.448],
                                   [-1.311,  0.000,  0.279]])
        self.xyz['cl'] = np.array([[ 0.000,  0.000,  0.000]])

    def _add_comb(self):
        """Adds substituents made by combining multiple substituents."""
        self.elem['et'], self.xyz['et'] = self.add_subs('me', 'me')
        self.elem['npr'], self.xyz['npr'] = self.add_subs('me', 'me', 'me')
        self.elem['ipr'], self.xyz['ipr'] = self.add_subs('me', 'me', 'me',
                                                          inds=1)
        self.elem['nbu'], self.xyz['nbu'] = self.add_subs('me', 'me', 'me',
                                                          'me')
        self.elem['ibu'], self.xyz['ibu'] = self.add_subs('me', 'me', 'me',
                                                          'me', inds=[2, 1, -1])
        self.elem['tbu'], self.xyz['tbu'] = self.add_subs('me', 'me', 'me',
                                                          'me', inds=1)
        self.elem['ome'], self.xyz['ome'] = self.add_subs('oh', 'me')
        self.elem['ac'], self.xyz['ac'] = self.add_subs('al', 'me')
        self.elem['ca'], self.xyz['ca'] = self.add_subs('al', 'oh')
        self.elem['tfm'], self.xyz['tfm'] = self.add_subs('me', 'f', 'f', 'f',
                                                          inds=1)
        self.elem['ms'], self.xyz['ms'] = self.add_subs('sf', 'me')

    def get_sub(self, label):
        """Returns the element list and cartesian geometry of a
        substituent.

        Parameters
        ----------
        label : str
            The substituent label of the desired substituent.

        Returns
        -------
        elem : (N,) ndarray
            The atomic symbols of the substituent.
        xyz : (N, 3) ndarray
            The atomic cartesian coordinates of the substituent.
        """
        lbl = self.syn[label.lower()]
        return self.elem[lbl], self.xyz[lbl]

    def add_subs(self, *lbls, inds=-1):
        """Returns the element list and cartesian geometry from a
        combination of substituents.

        Parameters
        ----------
        lbls : list
            A list of substituent labels to be combined.
        inds : int or array_like, optional
            The indices for substitution between substituents. Setting
            inds=-1 (default) makes the last atom the subtituted atom.
            Otherwise a list of indices can be given for the first of
            each pair of substituents.

        Returns
        -------
        elem : (N,) ndarray
            The atomic symbols of the combined substituent.
        xyz : (N, 3) ndarray
            The atomic cartesian coordinates of the combined substituent.
        """
        if isinstance(inds, int):
            inds = (len(lbls) - 1) * [inds]
        elif len(inds) != len(lbls) - 1:
            raise ValueError('Number of inds != number of labels - 1')

        rot = 0
        lbl0 = self.syn[lbls[0].lower()]
        elem = self.elem[lbl0]
        xyz = self.xyz[lbl0]
        for i, label in zip(inds, lbls[1:]):
            dist = np.linalg.norm(xyz - xyz[i], axis=1)
            dist[i] += np.max(dist)
            ibond = np.argmin(dist)
            rot = (rot + 1) % 2
            ax = con.unit_vec(xyz[i] - xyz[ibond])
            lbl = self.syn[label.lower()]
            new_elem = self.elem[lbl]
            new_xyz = displace.rotate(self.xyz[lbl], rot*np.pi, 'Z')
            new_xyz = displace.align_axis(new_xyz, 'Z', ax)
            blen = con.get_covrad(elem[ibond]) + con.get_covrad(new_elem[0])
            new_xyz += xyz[ibond] + blen * ax
            elem = np.hstack((np.delete(elem, i), new_elem))
            xyz = np.vstack((np.delete(xyz, i, axis=0), new_xyz))

        return elem, xyz

def import_sub(label):
    """Returns the element list and cartesian geometry of a substituent
    given its label.

    Parameters
    ----------
    label : str
        The substituent label.

    Returns
    -------
    elem : (N,) ndarray
        The atomic symbols of the substituent.
    xyz : (N, 3) ndarray
        The atomic cartesian coordinates of the substituent.
    """
    lib = SubLib()
    return lib.get_sub(label)

def subst(elem, xyz, sublbl, isub, ibond=None, pl=None, vec=None):
    """Returns a molecular geometry with an specified atom replaced by
    substituent.

    Labels are case-insensitive. The index isub gives the position to be
    substituted. If specified, ibond gives the atom bonded to the
    substituent. Otherwise, the nearest atom to isub is used. The
    orientation of the substituent can be given as a vector (the plane
    normal) or an index (the plane containing isub, ibond and pl).

    If isub is given as a list, the entire list of atoms is removed
    and the first index is treated as the position of the substituent.

    Parameters
    ----------
    elem : (N,) array_like
        The atomic symbols of the unsubstituted molecule.
    xyz : (N, 3) array_like
        The atomic cartesian coordinates of the unsubstituted molecule.
    sublbl : str
        The substituent label.
    isub : int or list
        The atomic index (or indices) to be replaced by the substituent.
    ibond : int, optional
        The atomic index of the atom bonded to position isub. If None
        (default), the nearest atom is chosen.
    pl : int or array_like, optional
        The atomic index or vector defining the xz-plane of the
        substituent. If an index is given, the plane normal to the
        isub-ibond-pl plane is used. If None (default), the plane
        is arbitrarily set to [1, 1, 1] and the bond axis is projected
        out.
    vec : (N, 3) array_like, optional
        The atomic cartesian vectors of the unsubstitued molecule. Default
        is None.

    Returns
    -------
    new_elem : (N,) ndarray
        The atomic symbols of the substituted molecule.
    new_xyz : (N, 3) ndarray
        The atomic cartesian coordinates of the substituted molecule.
    new_vec : (N, 3) ndarray
        The atomic cartesian vectors of the substituted molecule.
        Substituent atoms are all set of zero. If vec is None, new_vec
        is all zeros.
    """
    elem = np.array(elem)
    xyz = np.atleast_2d(xyz)
    if not isinstance(isub, int):
        ipos = isub[0]
    else:
        isub = [isub]
        ipos = isub[0]

    if ibond is None:
        dist = np.linalg.norm(xyz - xyz[ipos], axis=1)
        dist[ipos] += np.max(dist)
        ibond = np.argmin(dist)
    elif ibond == ipos:
        raise ValueError('sub and bond indices cannot be the same')

    ax = con.unit_vec(xyz[ipos] - xyz[ibond])
    if pl is None:
        # choose an arbitrary axis and project out the bond axis
        pl = np.ones(3)
        pl -= np.dot(pl, ax) * ax
    elif isinstance(pl, int):
        if pl == ipos:
            raise ValueError('plane and sub indices cannot be the same')
        elif pl == ibond:
            raise ValueError('plane and bond indices cannot be the same')
        pl = np.cross(xyz[ipos] - xyz[ibond], xyz[pl] - xyz[ibond])

    sub_el, sub_xyz = import_sub(sublbl)
    if elem[ipos] == sub_el[0]:
        blen = np.linalg.norm(xyz[ipos] - xyz[ibond])
    else:
        blen = con.get_covrad(elem[ibond]) + con.get_covrad(sub_el[0])

    # rotate to correct orientation and displace to correct position
    sub_xyz = displace.align_axis(sub_xyz, 'Z', ax)
    sub_pl = displace.align_axis([0., 1., 0.], 'Z', ax)
    sub_xyz = displace.align_axis(sub_xyz, sub_pl, pl)
    sub_xyz += xyz[ibond] + blen * ax

    # build the final geometry
    ind1 = [i for i in range(ipos) if i not in isub[1:]]
    ind2 = [i for i in range(ipos+1, len(elem)) if i not in isub[1:]]
    new_elem = np.hstack((elem[ind1], sub_el, elem[ind2]))
    new_xyz = np.vstack((xyz[ind1], sub_xyz, xyz[ind2]))
    if vec is None:
        return new_elem, new_xyz, None
    else:
        new_vec = np.vstack((vec[ind1], np.zeros((len(sub_el), 3)), vec[ind2]))
        return new_elem, new_xyz, new_vec


"""Module containing tools for EM-Bright classification of
compact binaries using trained supervised classifier
"""
import h5py

import numpy as np
from scipy.interpolate import interp1d
from astropy import cosmology, units as u

from . import (
    EOS_BAYES_FACTORS,
    PACKAGE_FILENAMES,
    computeDiskMass,
    utils
)

_classifiers = {
    eos: utils._open_and_return_clfs(
        PACKAGE_FILENAMES[f'{eos}.pickle']
    )
    for eos in EOS_BAYES_FACTORS
}
"""Classifiers keyed on EOS names. Order (clf_ns, clf_em)"""
assert set(_classifiers) == set(EOS_BAYES_FACTORS), "Inconsistency in"
" number of trained classifiers."

def mchirp(m1, m2):
    return(m1 * m2)**(3./5.)/(m1 + m2)**(1./5.)

def q(m1, m2):
    return m2/m1 if m2 < m1 else m1/m2

def source_classification(m1, m2, chi1, chi2, snr,
                          ns_classifier=None,
                          emb_classifier=None):
    """
    Computes ``HasNS`` and ``HasRemnant`` probabilities
    from point mass, spin and signal to noise ratio
    estimates.

    Parameters
    ----------
    m1 : float
        primary mass
    m2 : float
        secondary mass
    chi1 : float
        dimensionless primary spin
    chi2 : float
        dimensionless secondary spin
    snr : float
        signal to noise ratio of the signal
    ns_classifier : object, optional
        pickled object for NS classification
    emb_classifier : object, optional
        pickled object for EM brightness classification

    Returns
    -------
    tuple
        (P_NS, P_EMB) predicted values.

    Notes
    -----
    By default the classifiers, trained based
    on different nuclear equations of state (EoSs)
    are downloaded from the project page:
    https://git.ligo.org/deep.chatterjee/em-bright.
    The methodology is described in arXiv:1911.00116.
    The score from each classifier is weighted based on
    the bayes factors of individual EoSs as mentioned in
    Table I of arXiv:2104.08681.
    However, if the trained classifiers are supplied
    via ``ns_classifier`` and ``emb_classifier``,
    the score is reported based on the classifier instead
    of re-weighting the score.
    Examples
    --------
    >>> from ligo.em_bright import em_bright
    >>> em_bright.source_classification(2.0 ,1.0 ,0. ,0. ,10.0)
    (1.0, 1.0)
    """
    features = [[m1, m2, chi1, chi2, snr]]
    try:
        # custom classifiers supplied
        return (
            ns_classifier.predict_proba(features).T[1][0],
            emb_classifier.predict_proba(features).T[1][0]
        )
    except AttributeError as e:
        msg, *_ = e.args
        if msg != """'NoneType' object has no attribute 'predict_proba'""":
            raise

    reweighted_ns_score = reweighted_emb_score = 0.
    for eosname, bayes_factor in EOS_BAYES_FACTORS.items():
        ns_classifier, emb_classifier = _classifiers[eosname]
        reweighted_ns_score += ns_classifier.predict_proba(
            features).T[1][0] * bayes_factor
        reweighted_emb_score += emb_classifier.predict_proba(
            features).T[1][0] * bayes_factor
    return reweighted_ns_score, reweighted_emb_score

def get_redshifts(distances, N=10000):
    """
    Compute redshift using the Planck15 cosmology.

    Parameters
    ----------
    distances: float or numpy.ndarray
              distance(s) in Mpc

    N : int, optional
      Number of steps for the computation of the interpolation function

    Example
    -------
    >>> distances = np.linspace(10, 100, 10)
    >>> em_bright.get_redshifts(distances)
    array([0.00225566, 0.00450357, 0.00674384, 0.00897655,
           0.01120181, 0.0134197 , 0.01563032, 0.01783375
           0.02003009, 0.02221941])

    Notes
    -----
    This function accepts HDF5 posterior samples file and computes
    redshift by interpolating the distance-redshift relation.
    """
    function = cosmology.Planck15.luminosity_distance
    min_dist = np.min(distances)
    max_dist = np.max(distances)
    z_min = cosmology.z_at_value(func=function, fval=min_dist*u.Mpc)
    z_max = cosmology.z_at_value(func=function, fval=max_dist*u.Mpc)
    z_steps = np.linspace(z_min - (0.1*z_min), z_max + (0.1*z_min), N)
    lum_dists = cosmology.Planck15.luminosity_distance(z_steps)
    s = interp1d(lum_dists, z_steps)
    redshifts = s(distances)
    return redshifts

def source_classification_pe(posterior_samples_file, hdf5=True,
                             threshold=3.0, sourceframe=True):
    """
    Compute ``HasNS`` and ``HasRemnant`` probabilities from posterior
    samples.

    Parameters
    ----------
    posterior_samples_file : str
        Posterior samples file

    hdf5 : bool, optional
        Supply when not using HDF5 format

    threshold : float, optional
        Maximum neutron star mass for `HasNS` computation

    sourceframe : bool, optional
        Supply to use detector frame quantities

    Returns
    -------
    tuple
        (P_NS, P_EMB) predicted values.

    Examples
    --------
    >>> from ligo.em_bright import em_bright
    >>> em_bright.source_classification_pe('posterior_samples.hdf5')
    (1.0, 0.9616727412238634)
    >>> em_bright.source_classification_pe('posterior_samples.dat', hdf5=False)  # noqa:E501
    (0.0, 0.0)
    """
    if hdf5:
        with h5py.File(posterior_samples_file, 'r') as data:
            engine = list(data['lalinference'].keys())[0]
            samples = data['lalinference'][engine]['posterior_samples'][()]
        mc_det_frame = samples['mc']
        lum_dist = samples['dist']
        redshifts = get_redshifts(lum_dist)
        if sourceframe:
            mc = mc_det_frame/(1 + redshifts)
        else:
            mc = mc_det_frame

    else:
        samples = np.recfromtxt(posterior_samples_file, names=True)
        if sourceframe:
            mc = samples['mc_source']
        else:
            mc = samples['mc']

    q = samples['q']
    m1 = mc * (1 + q)**(1/5) * (q)**(-3/5)
    m2 = mc * (1 + q)**(1/5) * (q)**(2/5)

    try:
        chi1 = samples['a1'] * np.cos(samples['tilt1'])
        chi2 = samples['a2'] * np.cos(samples['tilt2'])
    except ValueError:
        # for aligned-spin PE, a1, a2 is the z component
        chi1 = samples['a1']
        chi2 = samples['a2']

    M_rem = computeDiskMass.computeDiskMass(m1, m2, chi1, chi2)
    prediction_ns = np.sum(m2 <= threshold)/len(m2)
    prediction_em = np.sum(M_rem > 0)/len(M_rem)

    return prediction_ns, prediction_em


from __future__ import division, print_function
import numpy as np
import sys
import os
import random

import LFPy
from kcsd import sample_data_path
import neuron
morphology_directory = os.path.join(sample_data_path,'morphology')
"""Electrode grid is 2D. If z is the zero dimention x=x, y=y.
If x is the zero dimension x=0, y=x, z=y.
If y is the zero dimention, z=x & x=y"""

"""cell_types = {'Ballstick:1,
'Y_shaped':2,
'Morpho1':3,
'Agasbogas':4,
'Mainen':5,
'User_defined:6',
'Gang_simple':7,
'Domi':8,
'test':9}
electrode_orientation = {'x':1, 'y':2, 'z':3}
electrode_distribute = {'Grid':1, 'Random':2, 'Hexagonal':3, 'Domi':4}
LFPy_sim = {'Random':1,
'Y_symmetric':2,
'Mainen':3,
'Oscill':4,
'Const':5,
'Sine':6 }
"""
rm = 4

class CellModel():
    MORPHOLOGY_FILES = {
        1: os.path.join(morphology_directory,"ballstick.hoc"),
        2: os.path.join(morphology_directory,"villa.hoc"),
        3: os.path.join(morphology_directory,"morpho1.swc"),
        4: os.path.join(morphology_directory,"neuron_agasbogas.swc"),
        5: os.path.join(morphology_directory,"Mainen_swcLike.swc"),
        6: os.path.join(morphology_directory,"retina_ganglion.swc"),
        7: os.path.join(morphology_directory,"Badea2011Fig2Du.CNG.swc"),
        8: os.path.join(morphology_directory,"DomiCell.swc"),
        9: os.path.join(morphology_directory,"Test.swc"),
        10: os.path.join(morphology_directory,"L5_Mainen96_LFPy.hoc"),
    }
    CELL_PARAMETERS = {
        'Ra': 123,
        # start time of simulation, recorders start at t=0
        'tstart': 0.,
        'passive': True,
        'passive_parameters': {'e_pas': -65,
        'g_pas': 1./30000},
        # initial crossmembrane potential
        'v_init': -65,
        'nsegs_method': 'fixed_length',
        'max_nsegs_length': 10,
        'custom_code': [],  # will run this file
        'dt': 0.25,
    }
    SYNAPSE_PARAMETERS = {
        # idx to be set later
        'e': 0.,  # reversal potential
        'syntype': 'ExpSyn',  # synapse type
        'tau': 2.,
        'weight': .04,  # synaptic weight
        'record_current': True,
    }
    SIMULATION_PARAMETERS = {
        'rec_imem': True,
    }
    ELECTRODE_PARAMETERS = {
        'method': 'linesource',

    }
    POINT_PROCESS = {
        'idx': 0,
        'pptype': 'IClamp',
        }

    def __init__(self, **kwargs):
        self.cell_parameters = self.CELL_PARAMETERS.copy()
        self.synapse_parameters = self.SYNAPSE_PARAMETERS.copy()
        self.simulation_parameters = self.SIMULATION_PARAMETERS.copy()
        self.electrode_parameters = self.ELECTRODE_PARAMETERS.copy()
        self.point_process = self.POINT_PROCESS.copy()
        self.cell_name = kwargs.pop('cell_name', 'cell_1')
        self.path = kwargs.pop('path', 'simulation')
        self.stimulus = kwargs.pop('stimulus', 'random')
        dt = kwargs.pop('dt', 0.5)
        self.cell_parameters['dt'] = dt
        eldistribute = kwargs.pop('electrode_distribution', 1)
        # according to which axis
        orientation = kwargs.pop('electrode_orientation', 2)
        colnb = kwargs.pop('colnb', 4)
        rownb = kwargs.pop('rownb', 4)
        xmin = kwargs.pop('xmin', 0)
        xmax = kwargs.pop('xmax', 200)
        ymin = kwargs.pop('ymin', -100)
        ymax = kwargs.pop('ymax', -500)
        tstop = kwargs.pop('tstop', 850)
        self.cell_parameters['tstop'] = tstop
        cell_electrode_dist = kwargs.pop('electrode_distance', 50)
        triside = 2*kwargs.pop('triside', 60)
        ssNB = kwargs.pop('seed', 123456)
        custom_code = kwargs.pop('custom_code', [])
        morphology_no = kwargs.pop('morphology', 1)
        self.sigma = kwargs.pop('sigma', 0.3)
        self.n_pre_syn = kwargs.pop('n_presyn', 1000)
        self.n_synapses = kwargs.pop('n_syn', 1000)
        self.new_path = os.path.join(self.path, self.cell_name)
        np.random.seed(ssNB)
        self.synapse_parameters['weight'] = kwargs.pop('weight', 0.04)
        if kwargs:
            raise TypeError('Invalid keyword arguments:', kwargs.keys())
        try:
            morphology = self.MORPHOLOGY_FILES[morphology_no]
        except AttributeError:
            sys.exit('Unknown morphology %d\n', morphology)
        if morphology_no == 2:
                self.make_y_shaped()
        self.make_cell(morphology, custom_code)
        self.setup_LFPy_2D_grid(eldistribute,
                                orientation,
                                colnb,
                                rownb,
                                xmin,
                                xmax,
                                ymin,
                                ymax,
                                cell_electrode_dist,
                                triside,
                                ssNB)
        self.add_electrodes()
        self.pre_syn_pick = np.empty((1,1))

    def make_y_shaped(self):
        self.cell_parameters['passive_parameters'] = {'g_pas':1./30000,
                                                      'e_pas':-65}
        self.cell_parameters['cm'] = 1.
        self.cell_parameters['Ra'] = 100.
        self.simulation_parameters['rec_vmem'] = True

    def stationary_poisson(self, nsyn, lambd, tstart, tstop):
        '''Generates nsyn stationary possion processes with
        rate lambda between tstart and tstop'''
        interval_s = (tstop - tstart)*.001
        spiketimes = []
        for i in range(nsyn):
            spikecount = np.random.poisson(interval_s*lambd)
            spikevec = np.empty(spikecount)
            if spikecount == 0:
                spiketimes.append(spikevec)
            else:
                spikevec = tstart\
                           + (tstop - tstart)*np.random.random(spikecount)
                spiketimes.append(np.sort(spikevec))
        return spiketimes

    def make_cell(self, morphology, custom_code=[]):
        self.cell_parameters['morphology'] = morphology

        for code in custom_code:
            self.cell_parameters['custom_code'].append(custom_code)

        self.cell = LFPy.Cell(**self.cell_parameters)

        if not morphology.endswith('.hoc'):
            if not self.cell_parameters['custom_code']:
               for section in self.cell.allseclist:
                   if 'soma' in section.name() or 'axon' in section.name():
                       section.insert('hh')
                       print('Inserting Hodgkin-Huxley channels into %s' % section.name())

        self.cell.set_pos(x = LFPy.cell.neuron.h.x3d(0),
                          y = LFPy.cell.neuron.h.y3d(0),
                          z = LFPy.cell.neuron.h.z3d(0))
        return self.cell

    def save_morphology_to_file(self):
        segments = self.cell.get_idx()
        nseg = len(segments)
        self.morphology = np.zeros((nseg+1, 7))
        coords = np.array((self.cell.xstart,
                           self.cell.ystart,
                           self.cell.zstart)).T
        ends = np.array((self.cell.xend,
                         self.cell.yend,
                         self.cell.zend)).T
        segdiam = self.cell.diam
        parents = {}
        self.morphology[0, 0] = 1
        self.morphology[0, 1] = 1
        self.morphology[0, 2:5] = coords[0]
        self.morphology[0, 5] = segdiam[0]
        self.morphology[0, 6] = -1

        for section in neuron.h.allsec():
            parents[section.name()] = section.parentseg()
        for sec in neuron.h.allsec():
            secn = sec.name()
            idxs = self.cell.get_idx(secn)

            for i, idx in enumerate(idxs):
                self.morphology[idx+1, 0] = idx+2
                self.morphology[idx+1, 2:5] = ends[idx]
                self.morphology[idx+1, 5] = segdiam[idx]
                if 'soma' in secn:
                    self.morphology[idx+1, 1] = 1
                elif 'dend' in secn:
                    self.morphology[idx+1, 1] = 3
                elif 'apic' in secn:
                    self.morphology[idx+1, 1] = 3
                elif 'axon' in secn:
                    self.morphology[idx+1, 1] = 2
                elif 'basal' in secn:
                    self.morphology[idx+1, 1] = 4
                else:
                    self.morphology[idx+1, 1] = 5
                if i == 0:
                    if not parents[secn]:
                        self.morphology[idx+1, 6] = 1
                    else:
                        x = parents[secn].x
                        how_many = len(self.cell.get_idx(parents[secn].sec.name()))
                        par_idx = int(x*how_many)
                        if par_idx > how_many - 1:
                            par_idx = how_many -1
                        par = self.cell.get_idx(parents[secn].sec.name())[par_idx]
                        self.morphology[idx+1, 6] = par + 2
                else:
                    self.morphology[idx+1, 6] = idx+1
        morph_path = os.path.join(self.new_path, 'morphology')
        if not os.path.exists(morph_path):
            print("Creating", morph_path)
            os.makedirs(morph_path)
        fname = os.path.join(morph_path, self.cell_name) + '.swc'
        print('Saving morphology to', fname)
        np.savetxt(fname,
                   self.morphology,
                   header='',
                   fmt=['%d', '%d', '%6.2f', '%6.2f', '%6.2f', '%6.2f', '%d'])

    def find_parent(self, i, coords, ends):
        for j, end in enumerate(ends):
            check_parent = np.isclose(coords[i], end)
            if check_parent[0] and check_parent[1] and check_parent[2]:
                return j

    def add_electrodes(self):
        self.electrode_parameters['x'] = self.ele_coordinates[:, 0],
        self.electrode_parameters['y'] = self.ele_coordinates[:, 1],
        self.electrode_parameters['z'] = self.ele_coordinates[:, 2],
        self.electrode_parameters['sigma'] = self.sigma
        electrode = LFPy.RecExtElectrode(**self.electrode_parameters)
        self.simulation_parameters['electrode'] = electrode

    def setup_LFPy_2D_grid(self,
                           eldistribute,
                           orientation,
                           colnb,
                           rownb,
                           xmin,
                           xmax,
                           ymin,
                           ymax,
                           cellelectrodedist,
                           triside,
                           ssNB):
        if orientation == 1:
            i, j, k = 1, 2, 0
        if orientation == 2:
            i, j, k = 2, 0, 1
        if orientation == 3:
            i, j, k = 0, 1, 2
        self.ele_coordinates = np.ones((rownb*colnb, 3))*cellelectrodedist

        if eldistribute == 1:  # grid
            linspace = np.linspace(xmin, xmax, rownb)
            self.ele_coordinates[:, i] = np.array(colnb*list(linspace))
            self.ele_coordinates[:, j] = np.repeat(np.linspace(ymin,
                                                               ymax,
                                                               colnb),
                                                   rownb)
        elif eldistribute == 2:  # random
            self.ele_coordinates[:, i] = np.random.uniform(low=xmin,
                                                           high=xmax,
                                                           size=rownb * colnb)
            self.ele_coordinates[:, j] = np.random.uniform(low=ymin,
                                                           high=ymax,
                                                           size=rownb * colnb)
        elif eldistribute == 3:
            assert (rownb % 2 == 0)
            triheight = triside*np.cos(np.pi/6)
            rownb = rownb//2
            triX1 = xmin + triside*np.arange(1, colnb+1) - triside
            triX2 = xmin - triside/2 + triside*np.arange(1, colnb+1)
            triY1 = ymin + 2*triheight*np.arange(1, rownb+1) - triheight
            triY2 = ymin + 2*triheight*np.arange(1, rownb+1)
            grid1 = [[], []]
            grid2 = [[], []]
            for l2 in range(len(triY1)):
                for l1 in range(len(triX1)):
                    grid1[0].append(triX1[l1])
                    grid1[1].append(triY1[l2])
                    grid2[0].append(triX2[l1])
                    grid2[1].append(triY2[l2])
            Xcoord = grid1[0] + grid2[0]
            Ycoord = grid1[1] + grid2[1]
            self.ele_coordinates[:, i] = Xcoord
            self.ele_coordinates[:, j] = Ycoord
        elif eldistribute == 4:
            self.ele_coordinates = np.loadtxt(os.path.join(sample_data_path, 'ElcoordsDomi14.txt'))
        if not os.path.exists(self.new_path):
            print("Creating", self.new_path)
            os.makedirs(self.new_path)
            self.add_electrodes()

    def constant_current_injection(self, amp, idx=0):
        self.point_process['idx'] = idx
        self.point_process['amp'] = amp
        self.point_process['dur'] = self.cell.tstop
        self.point_process['delay'] = 2
        stimulus = LFPy.StimIntElectrode(self.cell, **self.point_process)

    def cosine_current_injection(self, tstop=850):
        pre_syn_sptimes = self.stationary_poisson(nsyn=self.n_pre_syn,
                                                  lambd=2,
                                                  tstart=0,
                                                  tstop=400)
        l = np.arange(self.n_pre_syn)
        pre_syn_pick = np.random.permutation(l)[0:self.n_synapses]
        pars = {}
        for i_syn in range(self.n_synapses):
            syn_idx = int(self.cell.get_rand_idx_area_norm())
            spike_times = pre_syn_sptimes[pre_syn_pick[i_syn]]
            if syn_idx in pars:
                pars[syn_idx].extend(list(spike_times))
            else:
                pars[syn_idx] = list(spike_times)
        for syn_idx in pars:
            self.synapse_parameters.update({'idx': syn_idx})
            synapse = LFPy.Synapse(self.cell, **self.synapse_parameters)
            synapse.set_spike_times(np.array(pars[syn_idx]))

        TimesStim = np.arange(tstop)
        stim = np.array(3.6*np.sin(2.*3.141*6.5*TimesStim/1000.))
        for istim in range(tstop):
            pointprocess = {
                'idx' : 0,
                'pptype': 'IClamp',
                'record_current' : True,
                'amp': stim[istim],
                'dur' : 1.,
                'delay': istim,
                }
            stimulus = LFPy.StimIntElectrode(self.cell, **pointprocess)

    def distal_cosine_current_injection(self, tstop=850):
        pre_syn_sptimes = self.stationary_poisson(nsyn=self.n_pre_syn,
                                                  lambd=2,
                                                  tstart=0,
                                                  tstop=400)
        l = np.arange(self.n_pre_syn)
        pre_syn_pick = np.random.permutation(l)[0:self.n_synapses]
        pars = {}
        for i_syn in range(self.n_synapses):
            syn_idx = int(self.cell.get_rand_idx_area_norm())
            spike_times = pre_syn_sptimes[pre_syn_pick[i_syn]]
            if syn_idx in pars:
                pars[syn_idx].extend(list(spike_times))
            else:
                pars[syn_idx] = list(spike_times)
        for syn_idx in pars:
            self.synapse_parameters.update({'idx': syn_idx})
            synapse = LFPy.Synapse(self.cell, **self.synapse_parameters)
            synapse.set_spike_times(np.array(pars[syn_idx]))

        TimesStim = np.arange(850)
        stim = np.array(3.6*np.sin(2.*3.141*6.5*TimesStim/1000.))
        all_idxs = self.cell.get_idx()

        for istim in range(tstop):
            pointprocess = {
                'idx' : max(all_idxs),
                'pptype': 'IClamp',
                'record_current' : True,
                'amp': stim[istim],
                'dur' : 1.,
                'delay': istim,
                }
            stimulus = LFPy.StimIntElectrode(self.cell, **pointprocess)

    def random_synaptic_input(self,
                              lambd=2,
                              tstart=0,
                              tstop=70):

        self.synapse_parameters['idx'] = 0
        pre_syn_sptimes = self.stationary_poisson(self.n_pre_syn,
                                                  lambd,
                                                  tstart,
                                                  tstop)
        l = np.arange(self.n_pre_syn)
        self.pre_syn_pick = np.random.permutation(l)[0:self.n_synapses]

        for i_syn in range(self.n_synapses):
            syn_idx = int(self.cell.get_rand_idx_area_norm())
            self.synapse_parameters.update({'idx': syn_idx})
            synapse = LFPy.Synapse(self.cell, **self.synapse_parameters)

            synapse.set_spike_times(pre_syn_sptimes[self.pre_syn_pick[i_syn]])

    def y_shaped_symmetric_input(self):
        self.synapse_parameters['idx'] = 0
        pre_syn_sptimes = [np.array([5., 25., 60.]), np.array([5., 45., 60.])]
        syn_no = [65, 33]

        for i, i_syn in enumerate(syn_no):
            new_pars = self.synapse_parameters.copy()
            new_pars['idx'] = i_syn
            synapse = LFPy.Synapse(self.cell, **new_pars)
            synapse.set_spike_times(pre_syn_sptimes[i])

    def sine_synaptic_input(self, tstop=None):
        if not tstop:
            tstop = self.cell.tstop
        frequencies = np.arange(0.5, 13, 0.5)
        i = 0
        distance = 0
        nseg = self.cell.get_idx()
        freq_step = sum(self.cell.length)/len(frequencies)
        for j, istim in enumerate(nseg):
            distance += self.cell.length[j]
            if distance > (i + 1)*freq_step:
                i += 1
            freq = frequencies[i]
            pointprocess = {
                'idx': istim,
                'pptype': 'SinSyn',
                'pkamp':  3.6,
                'freq': freq,
                'phase': -np.pi/2,
                'dur': self.cell.tstop,
            }
            stimulus = LFPy.StimIntElectrode(self.cell, **pointprocess)

    def simulate(self, stimulus=None):
        if stimulus:
            self.stimulus = stimulus
        if self.stimulus == 'constant':
            self.constant_current_injection(amp=10, idx=0)
        elif self.stimulus == 'random':
            self.random_synaptic_input()
        elif self.stimulus == 'sine':
            self.sine_synaptic_input()
        elif self.stimulus == 'symmetric':
            self.y_shaped_symmetric_input()
        elif self.stimulus == 'oscillatory':
            self.cosine_current_injection(tstop=self.cell_parameters['tstop'])
        elif self.stimulus == 'distal_oscillatory':
            self.distal_cosine_current_injection(tstop=self.cell_parameters['tstop'])
        self.cell.simulate(**self.simulation_parameters)

    def save_LFP(self, directory=''):
        self.simulation_parameters['electrode'].calc_lfp()
        LFP_path = os.path.join(self.new_path, directory)
        if not os.path.exists(LFP_path):
            print("Creating", LFP_path)
            os.makedirs(LFP_path)
        fname = os.path.join(LFP_path, 'MyLFP')
        np.savetxt(fname, self.simulation_parameters['electrode'].LFP)

    def save_electrode_pos(self, directory=''):
        electr = np.hstack((self.ele_coordinates[:, 0],
                            self.ele_coordinates[:, 1],
                            self.ele_coordinates[:, 2]))
        elcoord_x_y_x_path = os.path.join(self.new_path, directory)
        if not os.path.exists(elcoord_x_y_x_path):
            print("Creating", elcoord_x_y_x_path)
            os.makedirs(elcoord_x_y_x_path)
        fname = os.path.join(elcoord_x_y_x_path, 'elcoord_x_y_z')
        np.savetxt(fname, electr)

    def save_somav(self, directory=''):
        if directory:
            new_path = directory
        else:
            new_path = self.new_path
        np.savetxt(os.path.join(new_path, 'somav.txt'),
                   self.cell.somav)

    def save_tvec(self, directory=''):
        if directory:
            new_path = directory
        else:
            new_path = self.new_path
        np.savetxt(os.path.join(new_path, 'tvec.txt'),
                   self.cell.tvec)

    def save_for_R_kernel(self, directory=''):
        if directory:
            new_path = directory
        else:
            new_path = self.new_path
        self.save_LFP(directory)
        self.save_electrode_pos(directory)
        self.save_somav(directory)
        coords = np.hstack((self.cell.xmid,
                            self.cell.ymid,
                            self.cell.zmid))
        np.savetxt(os.path.join(new_path, 'coordsmid_x_y_z'),
                   coords)
        # coordinates of the segment's beginning
        coordsstart = np.hstack((self.cell.xstart,
                                 self.cell.ystart,
                                 self.cell.zstart))
        np.savetxt(os.path.join(new_path, 'coordsstart_x_y_z'),
                   coordsstart)
        # coordinates of the segment's end
        coordsend = np.hstack((self.cell.xend,
                               self.cell.yend,
                               self.cell.zend))
        np.savetxt(os.path.join(new_path, 'coordsend_x_y_z'),
                   coordsend)
        # diameter of the segments
        segdiam = np.hstack((self.cell.diam))
        np.savetxt(os.path.join(new_path, 'segdiam_x_y_z'),
                   segdiam)
        # time in the simulation
        np.savetxt(os.path.join(new_path, 'time'),
                   self.cell.tvec)
        # let's write to file the simulation locations
        np.savetxt('synapse_locations', self.pre_syn_pick)
        self.save_memb_curr()
        self.save_seg_length()

    def save_memb_curr(self, directory=''):
        if directory:
            new_path = directory
        else:
            new_path = self.new_path
        np.savetxt(os.path.join(new_path, 'membcurr'),
                   self.cell.imem)

    def save_seg_length(self, directory=''):
        if directory:
            new_path = directory
        else:
            new_path = self.new_path
        np.savetxt(os.path.join(new_path, 'seglength'),
                   self.cell.length)

    def save_skCSD_python(self):
        self.save_morphology_to_file()
        self.save_LFP('LFP')
        self.save_electrode_pos('electrode_positions')

    def return_paths_skCSD_python(self):
        return self.new_path

if __name__ == '__main__':
    c = CellModel(morphology=7,
                  cell_name='Gang_simple',
                  electrode_distribution=4,
                  colnb=1,
                  rownb=8,
                  xmin=-500,
                  xmax=500,
                  ymin=-500,
                  ymax=500)
    c.simulate()
    c.save_skCSD_python()
    c.save_for_R_kernel()


#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
   This file belong to https://github.com/snolfi/evorobotpy
   and has been written by Stefano Nolfi and Paolo Pagliuca, stefano.nolfi@istc.cnr.it, paolo.pagliuca@istc.cnr.it
   salimans.py include an implementation of the OpenAI-ES algorithm described in
   Salimans T., Ho J., Chen X., Sidor S & Sutskever I. (2017). Evolution strategies as a scalable alternative to reinforcement learning. arXiv:1703.03864v2
   requires es.py, policy.py, and evoalgo.py
"""

import numpy as np
from numpy import zeros, ones, dot, sqrt
import math
import time
from mpi4py import MPI
from evoalgo import EvoAlgo
from utils import ascendent_sort
import sys
import os
import configparser

# Parallel implementation of Open-AI-ES algorithm developed by Salimans et al. (2017)
# the workers evaluate a fraction of the population in parallel
# the master post-evaluate the best sample of the last generation and eventually update the input normalization vector

class Algo(EvoAlgo):
    def __init__(self, env, policy, seed, fileini, filedir):
        EvoAlgo.__init__(self, env, policy, seed, fileini, filedir)

    def loadhyperparameters(self):

        if os.path.isfile(self.fileini):

            config = configparser.ConfigParser()
            config.read(self.fileini)
            self.maxsteps = 1000000
            self.stepsize = 0.01
            self.batchSize = 20
            self.noiseStdDev = 0.02
            self.wdecay = 0
            self.symseed = 1
            self.saveeach = 60
            options = config.options("ALGO")
            for o in options:
                found = 0
                if o == "maxmsteps":
                    self.maxsteps = config.getint("ALGO","maxmsteps") * 1000000
                    found = 1
                if o == "stepsize":
                    self.stepsize = config.getfloat("ALGO","stepsize")
                    found = 1
                if o == "noisestddev":
                    self.noiseStdDev = config.getfloat("ALGO","noiseStdDev")
                    found = 1
                if o == "samplesize":
                    self.batchSize = config.getint("ALGO","sampleSize")
                    found = 1
                if o == "wdecay":
                    self.wdecay = config.getint("ALGO","wdecay")
                    found = 1
                if o == "symseed":
                    self.symseed = config.getint("ALGO","symseed")
                    found = 1
                if o == "saveeach":
                    self.saveeach = config.getint("ALGO","saveeach")
                    found = 1

                if found == 0:
                    print("\033[1mOption %s in section [ALGO] of %s file is unknown\033[0m" % (o, self.fileini))
                    print("available hyperparameters are: ")
                    print("maxmsteps [integer]       : max number of (million) steps (default 1)")
                    print("stepsize [float]          : learning stepsize (default 0.01)")
                    print("samplesize [int]          : popsize/2 (default 20)")
                    print("noiseStdDev [float]       : samples noise (default 0.02)")
                    print("wdecay [0/2]              : weight decay (default 0), 1 = L1, 2 = L2")
                    print("symseed [0/1]             : same environmental seed to evaluate symmetrical samples [default 1]")
                    print("saveeach [integer]        : save file every N minutes (default 60)")

                    sys.exit()
        else:
            print("\033[1mERROR: configuration file %s does not exist\033[0m" % (self.fileini))

    def setProcess(self, n_workers, comm, rank):
        self.loadhyperparameters()               # load parameters
        self.n_workers = n_workers               # number of workers, includes the master
        self.rank = rank                         # worker id
        self.comm = comm                         #
        self.center = np.copy(self.policy.get_trainable_flat())  # the initial centroid
        self.nparams = len(self.center)          # number of adaptive parameters
        self.n_worker_samples = int(self.batchSize / (self.n_workers - 1)) # number of sample evaluated by each worker
        self.id = (self.rank - 1)                # id of the process (master has id -1)
        self.cgen = 0                            # currrent generation
        self.fitness = ones(self.n_workers * (self.n_worker_samples * 2)) # vector of fitness filled by the master and the workers
        self.evals = zeros(self.n_workers, dtype=np.int32)  #vector of evaluation steps filled by the master and by the workers
        self.samplefitness = zeros(self.batchSize * 2) # the fitness of the samples
        self.samples = None                      # the random samples
        self.m = zeros(self.nparams)             # Adam: momentum vector
        self.v = zeros(self.nparams)             # Adam: second momentum vector (adam)
        self.epsilon = 1e-08                     # Adam: To avoid numerical issues with division by zero...
        self.beta1 = 0.9                         # Adam: beta1
        self.beta2 = 0.999                       # Adam: beta2
        self.bestgfit = -99999999                # the best generalization fitness
        self.bfit = 0                            # the fitness of the best sample
        self.gfit = 0                            # the postevaluation fitness of the best sample of last generation
        self.rs = None                           # random number generator
        if self.policy.normalize == 1:           # normalization vector
            self.normvector = np.arange(self.n_workers * (self.policy.ninputs * 2), dtype=np.float64)  # normalization vector broadcasted to workers
        self.inormepisodes = self.batchSize * 2 * self.policy.ntrials / 100.0 # number of normalization episode for generation (1% of generation episodes)
        self.tnormepisodes = 0.0                 # total epsidoes in which normalization data should be collected so far
        self.normepisodes = 0                    # numer of episodes in which normalization data has been actually collected so far

    def savedata(self):
        # save best postevaluated so far
        fname = self.filedir + "/bestgS" + str(self.seed)
        np.save(fname, self.bestgsol)
        # save best so far
        fname = self.filedir + "/bestS" + str(self.seed)
        np.save(fname, self.bestsol)
        # save statistics
        fname = self.filedir + "/statS" + str(self.seed)
        np.save(fname, self.stat)
        # save summary statistics
        fname = self.filedir + "/S" + str(self.seed) + ".fit"
        fp = open(fname, "w")
        fp.write('Seed %d (%.1f%%) gen %d msteps %d bestfit %.2f bestgfit %.2f bestsam %.2f avgfit %.2f paramsize %.2f \n' %
             (self.seed, self.steps / float(self.maxsteps) * 100, self.cgen, self.steps / 1000000, self.bestfit, self.bestgfit, self.bfit, self.avgfit, self.avecenter))
        fp.close()

    def evaluate(self):
        global none_val
        seed_worker = self.seed + self.cgen * self.batchSize  # Set the seed for current generation (master and workers have the same seed)
        self.rs = np.random.RandomState(seed_worker)
        self.samples = self.rs.randn(self.batchSize, self.nparams)
        self.cgen += 1
        fitness_worker = ones(self.n_worker_samples * 2)
        ceval = 0

        if self.rank == 0:                                      # Master: postevaluate the best sample of last generation
            gfit = 0
            if self.bestsol is not None:
                self.policy.set_trainable_flat(self.bestsol)
                self.tnormepisodes += self.inormepisodes
                normalizationdatacollected = False
                for t in range(self.policy.nttrials):
                    if self.policy.normalize == 1 and self.normepisodes < self.tnormepisodes:
                        self.policy.nn.normphase(1)
                        self.normepisodes += 1
                        normalizationdatacollected = True
                    else:
                        self.policy.nn.normphase(0)
                    eval_rews, eval_length = self.policy.rollout(1, seed=(self.seed + 100000 + t))
                    gfit += eval_rews
                    ceval += eval_length
                self.updateBestg(gfit / self.policy.nttrials, self.bestsol)
                if normalizationdatacollected:
                    self.policy.nn.updateNormalizationVectors()  # update the normalization vectors with the new data collected
            else:
                self.policy.nn.getNormalizationVectors()         # update the normalization vector accessible in python with that initialized by evonet
        else:
            candidate = np.arange(self.nparams, dtype=np.float64)
            for b in range(self.n_worker_samples):               # Worker (evaluate a fraction of the population)
                for bb in range(2):
                    if (bb == 0):
                        candidate = self.center + self.samples[(self.id * self.n_worker_samples) + b,:] * self.noiseStdDev
                    else:
                        candidate = self.center - self.samples[(self.id * self.n_worker_samples) + b,:] * self.noiseStdDev
                    self.policy.set_trainable_flat(candidate)
                    self.policy.nn.normphase(0)   # workers never collect normalization data
                    eval_rews, eval_length = self.policy.rollout(self.policy.ntrials, seed=(self.seed + (self.cgen * self.batchSize) + (self.id * self.n_worker_samples) + b))
                    fitness_worker[b*2+bb] = eval_rews
                    ceval += eval_length
        ceval = np.asarray([ceval], dtype=np.int32)
        return fitness_worker, ceval

    def optimize(self):

        fitness, index = ascendent_sort(self.samplefitness)       # sort the fitness
        self.avgfit = np.average(fitness)                         # compute the average fitness

        self.bfit = fitness[(self.batchSize * 2) - 1]
        bidx = index[(self.batchSize * 2) - 1]
        if ((bidx % 2) == 0):                                     # regenerate the genotype of the best samples
            bestid = int(bidx / 2)
            self.bestsol = self.center + self.samples[bestid] * self.noiseStdDev
        else:
            bestid = int(bidx / 2)
            self.bestsol = self.center - self.samples[bestid] * self.noiseStdDev

        if self.rank == 0:
            self.updateBest(self.bfit, self.bestsol)              # Stored if it is the best obtained so far

        popsize = self.batchSize * 2                              # compute a vector of utilities [-0.5,0.5]
        utilities = zeros(popsize)
        for i in range(popsize):
            utilities[index[i]] = i
        utilities /= (popsize - 1)
        utilities -= 0.5

        weights = zeros(self.batchSize)                           # Assign the weights (utility) to samples on the basis of their fitness rank
        for i in range(self.batchSize):
            idx = 2 * i
            weights[i] = (utilities[idx] - utilities[idx + 1])    # merge the utility of symmetric samples

        g = 0.0
        i = 0
        while i < self.batchSize:                                 # Compute the gradient (the dot product of the samples for their utilities)
            gsize = -1
            if self.batchSize - i < 500:                          # if the popsize is larger than 500, compute the gradient for multiple sub-populations
                gsize = self.batchSize - i
            else:
                gsize = 500
            g += dot(weights[i:i + gsize], self.samples[i:i + gsize,:])
            i += gsize
        g /= popsize                                              # normalize the gradient for the popsize

        if self.wdecay == 1:
            globalg = -g + 0.005 * self.center                    # apply weight decay
        else:
            globalg = -g

        # adam stochastic optimizer
        a = self.stepsize * sqrt(1.0 - self.beta2 ** self.cgen) / (1.0 - self.beta1 ** self.cgen)
        self.m = self.beta1 * self.m + (1.0 - self.beta1) * globalg
        self.v = self.beta2 * self.v + (1.0 - self.beta2) * (globalg * globalg)
        dCenter = -a * self.m / (sqrt(self.v) + self.epsilon)

        self.center += dCenter                                    # move the center in the direction of the momentum vectors
        self.avecenter = np.average(np.absolute(self.center))

    def update_normvector(self):
        if self.rank > 0:                                                # workers overwrite their normalization vector with the vector received from the master
            for i in range(self.policy.ninputs * 2):
                self.policy.normvector[i] = np.copy(self.normvector[i])
            self.policy.nn.setNormalizationVectors()

    def run(self):

        start_time = time.time()
        last_save_time = start_time
        elapsed = 0
        self.steps = 0
        if self.rank ==0:
            print("Salimans: seed %d maxmsteps %d batchSize %d stepsize %lf noiseStdDev %lf wdecay %d symseed %d nparams %d" % (self.seed, self.maxsteps / 1000000, self.batchSize, self.stepsize, self.noiseStdDev, self.wdecay, self.symseed, self.nparams))

        while (self.steps < self.maxsteps):

            fitness_worker, weval = self.evaluate()   # evaluate sample (each worker evaluate a fraction of the population)
            self.comm.Allgatherv(fitness_worker, [self.fitness, MPI.DOUBLE])                # brodcast fitness value to all workers
            self.comm.Allgatherv(weval, [self.evals, MPI.INT])                              # broadcast number of steps performed to all workers
            if self.policy.normalize == 1:
                self.comm.Allgatherv(self.policy.normvector, [self.normvector, MPI.DOUBLE]) # broadcast normalization vector (it is update from the master, the vectors of the workers are ignored)

            self.samplefitness = self.fitness[(self.n_worker_samples * 2):] # Merge the fitness of the workers by discarding the vector returned by the master
            self.steps += np.sum(self.evals)          # Update the total number of steps performed so far

            self.optimize()                           # estimate the gradient and move the centroid in the gradient direction

            self.stat = np.append(self.stat, [self.steps, self.bestfit, self.bestgfit, self.bfit, self.avgfit, self.avecenter])  # store performance across generations
            if self.rank == 0 and ((time.time() - last_save_time) > (self.saveeach * 60)):
                self.savedata()                       # save data on files
                last_save_time = time.time()

            if self.policy.normalize == 1:
                self.update_normvector()              # the workers overwrite their normalization vector with the vector received from the master

            if self.rank == 0:
                print('Seed %d (%.1f%%) gen %d msteps %d bestfit %.2f bestgfit %.2f bestsam %.2f avg %.2f weightsize %.2f' %
                      (self.seed, self.steps / float(self.maxsteps) * 100, self.cgen, self.steps / 1000000, self.bestfit, self.bestgfit, self.bfit, self.avgfit, self.avecenter))

        if self.rank == 0:
            self.savedata()                           # save data at the end of evolution

        # print simulation time
        end_time = time.time()
        print('Simulation time: %dm%ds ' % (divmod(end_time - start_time, 60)))


import numpy as np
import numba

import strax
export, __all__ = strax.exporter()

# Hardcoded numbers:
TRIAL_COUNTER_NEIGHBORING_RECORDS = 100  # Trial counter when looking for hitlet data.
NO_FWXM = -42  # Value in case FWXM cannot be found.

# ----------------------
# Hitlet building:
# ----------------------
@export
def concat_overlapping_hits(hits, extensions, pmt_channels, start, end):
    """
    Function which concatenates hits which may overlap after left and
    right hit extension. Assumes that hits are sorted correctly.

    Note:
        This function only updates time, length and record_i of the hit.
        (record_i is set according to the first hit)

    :param hits: Hits in records.
    :param extensions: Tuple of the left and right hit extension.
    :param pmt_channels: Tuple of the detectors first and last PMT
    :param start: Startime of the chunk
    :param end: Endtime of the chunk

    :returns:
        array with concataneted hits.
    """
    # Getting channel map and compute the number of channels:
    first_channel, last_channel = pmt_channels
    nchannels = last_channel - first_channel + 1

    # Buffer for concat_overlapping_hits, if specified in
    # _concat_overlapping_hits numba crashes.
    last_hit_in_channel = np.zeros(nchannels,
                                   dtype=(strax.hit_dtype
                                          + [(('End time of the interval (ns since unix epoch)',
                                               'endtime'), np.int64)]))

    if len(hits):
        hits = _concat_overlapping_hits(hits, extensions, first_channel, last_hit_in_channel, start, end)
    return hits

@strax.utils.growing_result(strax.hit_dtype, chunk_size=int(1e4))
@numba.njit(nogil=True, cache=True)
def _concat_overlapping_hits(hits,
                             extensions,
                             first_channel,
                             last_hit_in_channel,
                             start=0,
                             end=float('inf'),
                             _result_buffer=None):
    buffer = _result_buffer
    offset = 0

    le, re = extensions
    dt = hits['dt'][0]
    assert np.all(hits['dt'] == dt), 'All hits must have the same dt!'

    for h in hits:
        st = h['time'] - int(le * h['dt'])
        et = strax.endtime(h) + int(re * h['dt'])
        hc = h['channel']
        r_i = h['record_i']

        lhc = last_hit_in_channel[hc - first_channel]
        # Have not found any hit in this channel yet:
        if lhc['time'] == 0:
            lhc['time'] = max(st, start)
            lhc['endtime'] = min(et, end)
            lhc['channel'] = hc
            lhc['record_i'] = h['record_i']
            lhc['dt'] = dt

        # Checking if events overlap:
        else:
            if lhc['endtime'] >= st:
                # Yes, so we have to update only the end_time:
                lhc['endtime'] = et
            else:
                # No, this means we have to save the previous data and update lhc:
                res = buffer[offset]
                res['time'] = lhc['time']
                res['length'] = (lhc['endtime'] - lhc['time']) // lhc['dt']
                res['channel'] = lhc['channel']
                res['record_i'] = lhc['record_i']
                res['dt'] = lhc['dt']
                offset += 1
                if offset == len(buffer):
                    yield offset
                    offset = 0

                # Updating current last hit:
                lhc['time'] = st
                lhc['endtime'] = et
                lhc['channel'] = hc
                lhc['record_i'] = r_i

    # We went through so now we have to save all remaining hits:
    mask = last_hit_in_channel['time'] != 0
    for lhc in last_hit_in_channel[mask]:
        res = buffer[offset]
        res['time'] = lhc['time']
        res['channel'] = lhc['channel']
        res['length'] = (lhc['endtime'] - lhc['time']) // lhc['dt']
        res['record_i'] = lhc['record_i']
        res['dt'] = lhc['dt']
        offset += 1
        if offset == len(buffer):
            yield offset
            offset = 0
    yield offset

@export
@numba.njit(nogil=True, cache=True)
def refresh_hit_to_hitlets(hits, hitlets):
    """
    Function which copies basic hit information into a new hitlet array.
    """
    nhits = len(hits)
    for ind in range(nhits):
        h_new = hitlets[ind]
        h_old = hits[ind]

        h_new['time'] = h_old['time']
        h_new['length'] = h_old['length']
        h_new['channel'] = h_old['channel']
        h_new['area'] = h_old['area']
        h_new['record_i'] = h_old['record_i']
        h_new['dt'] = h_old['dt']

@export
@numba.njit(nogil=True, cache=True)
def get_hitlets_data(hitlets, records, to_pe):
    """
    Function which searches for every hitlet in a given chunk the
    corresponding records data.

    :param hitlets: Hitlets found in a chunk of records.
    :param records: Records of the chunk.
    :param to_pe: Array with area conversion factors from adc/sample to
        pe/sample

    Note:
        hitlets must have a "data" and "area" field.

    The function updates the hitlet fields time, length (if necessary
    e.g. hit was extended in regions of now records) and area
    according to the found data.
    """

    rlink = strax.record_links(records)
    for h in hitlets:
        data, start_time = get_single_hitlet_data(h, records, *rlink)
        h['length'] = len(data)
        h['data'][:len(data)] = data * to_pe[h['channel']]
        h['time'] = start_time
        h['area'] = np.sum(data * to_pe[h['channel']])

@export
@numba.njit(nogil=True, cache=True)
def get_single_hitlet_data(hitlet, records, prev_r, next_r):
    """
    Function which gets the data of a single hit or hitlet. The data is
    returned according to the objects time and length (LE/RE is not
    included in case of a hit.).

    In case the hit or hitlet is extended into non-recorded regions
    the data gets chopped.

    :param hitlet: Hits or hitlets.
    :param records: Records
    :param prev_r: Index of the previous record seen from the current
        record. (Return of strax.record_links)
    :param next_r: Index of the next record seen by the current record.
        (Return of strax.record_links)
    :return:
        np.ndarray: Samples of the hitlet [ADC]
        int: Start time of the hitlet. (In case data gets chopped on the
            left)
    """
    temp_data = np.zeros(hitlet['length'], dtype=np.float64)

    # Lets get the starting record and the corresponding data:
    r_i = hitlet['record_i']
    r = records[r_i]
    data, (p_start_i, p_end_i) = _get_thing_data(hitlet, r)
    temp_data[p_start_i:p_end_i] = data

    # We have to store the first and last index of the so far found data
    data_start = p_start_i
    data_end = p_end_i

    if not (p_end_i - p_start_i == hitlet['length']):
        # We have not found the entire data yet....
        # Starting with data before our current record:
        trial_counter = 0
        prev_r_i = r_i
        while p_start_i:
            # We are still searching for data in a previous record.
            temp_prev_r_i = prev_r[prev_r_i]
            if temp_prev_r_i == -1:
                # There is no (more) previous record. So stop here and keep
                # last pre_r_i
                break
            prev_r_i = temp_prev_r_i

            # There is a previous record:
            r = records[prev_r_i]
            data, (p_start_i, end) = _get_thing_data(hitlet, r)
            if not end:
                # If end is zero this means we have not found any
                # overlap which should not have happened. In case of an
                # overlap start and end should reflect the start and end
                # sample of the hitlet for which we found data.
                print('Data found for this record:', data,
                      'Start index', p_start_i,
                      'End index:', end)
                raise ValueError('This is odd found previous record, but no'
                                 ' overlapping indices.')
            temp_data[p_start_i:data_start] = data
            data_start = p_start_i

            if trial_counter > TRIAL_COUNTER_NEIGHBORING_RECORDS:
                raise RuntimeError('Tried too hard. There are more than'
                                   '100 successive records. This is odd...')
            trial_counter += 1

        # Now we have to do the very same for records in the future:
        # Almost the same code as above can I change this?
        trial_counter = 0
        next_r_i = r_i
        while hitlet['length'] - p_end_i:
            # We are still searching for data in a next record.
            temp_next_r_i = next_r[next_r_i]
            if temp_next_r_i == -1:
                # There is no (more) previous record. So stop here and keep
                # last next_r_i
                break
            next_r_i = temp_next_r_i
            # There is a next record:
            r = records[next_r_i]
            data, (start, p_end_i) = _get_thing_data(hitlet, r)
            if not start:
                # If start is zero this means we have not found any
                # overlap which should not have happened. In case of an
                # overlap start and end should reflect the start and end
                # sample of the hitlet for which we found data.
                print('Data found for this record:', data,
                      'Start index', start,
                      'End index:', p_end_i)
                raise ValueError('This is odd found the next record, but no'
                                 ' overlapping indicies.')
            temp_data[data_end:p_end_i] = data
            data_end = p_end_i

            if trial_counter > TRIAL_COUNTER_NEIGHBORING_RECORDS:
                raise RuntimeError('Tried too hard. There are more than'
                                   '100 successive records. This is odd...')
            trial_counter += 1

    # In some cases it might have happened that due to the left and right hit extension
    # we extended our hitlet into regions without any data so we have to chop
    # "time" according to the data we found....
    time = hitlet['time'] + data_start * hitlet['dt']
    temp_data = temp_data[data_start:data_end] + r['baseline'] % 1
    return temp_data, time

@numba.njit(nogil=True, cache=True)
def _get_thing_data(thing, container):
    """
    Function which returns data for some overlapping indices of a thing
    in a container.

    Note:
        Thing must be of the interval dtype kind.
    """
    overlap_hit_i, overlap_record_i = strax.overlap_indices(thing['time']//thing['dt'],
                                                            thing['length'],
                                                            container['time']//container['dt'],
                                                            container['length'])
    data = container['data'][overlap_record_i[0]:overlap_record_i[1]]
    return data, overlap_hit_i

# ----------------------
# Hitlet splitting:
# ----------------------
@export
def update_new_hitlets(hitlets, records, next_ri, to_pe):
    """
    Function which computes the hitlet data area and record_i after
    splitting.

    :param hitlets: New hitlets received after splitting.
    :param records: Records of the chunk.
    :param next_ri: Index of next record for current record record_i.
    :param  to_pe: ADC to PE conversion factor array (of n_channels).
    """
    _update_record_i(hitlets, records, next_ri)
    get_hitlets_data(hitlets, records, to_pe)

@numba.njit(cache=True, nogil=True)
def _update_record_i(new_hitlets, records, next_ri):
    """
    Function which updates the record_i value of the new hitlets.

    Notes:
        Assumes new_hitlets to be sorted in time.
    """
    for ind, hit in enumerate(new_hitlets):

        updated = False
        counter = 0
        current_ri = hit['record_i']
        while not updated:
            r = records[current_ri]
            # Hitlet must only partially be contained in record_i:
            time = hit['time']
            end_time = strax.endtime(hit)
            start_in = (r['time'] <= time) & (time < strax.endtime(r))
            end_in = (r['time'] < end_time) & (end_time <= strax.endtime(r))
            if start_in or end_in:
                hit['record_i'] = current_ri
                break
            else:
                last_ri = current_ri
                current_ri = next_ri[current_ri]
                counter += 1

            if current_ri == -1:
                print('Record:\n', r, '\nHit:\n', hit)
                raise ValueError('Was not able to find record_i')

            if counter > TRIAL_COUNTER_NEIGHBORING_RECORDS:
                print(ind, last_ri)
                raise RuntimeError('Tried too often to find correct record_i.')

# ----------------------
# Hitlet properties:
# ----------------------
@export
@numba.njit(cache=True, nogil=True)
def hitlet_properties(hitlets):
    """
    Computes additional hitlet properties such as amplitude, FHWM, etc.
    """
    for h in hitlets:
        dt = h['dt']
        data = h['data'][:h['length']]

        if np.any(data):
            # Compute amplitude
            amp_ind = np.argmax(data)
            amp_time = int(amp_ind * dt)
            height = data[amp_ind]

            h['amplitude'] = height
            h['time_amplitude'] = amp_time

            # Computing FWHM:
            left_edge, right_edge = get_fwxm(h, 0.5)
            width = right_edge - left_edge

            # Computing FWTM:
            left_edge_low, right_edge = get_fwxm(h, 0.1)
            width_low = right_edge - left_edge_low

            h['fwhm'] = width
            h['left'] = left_edge
            h['low_left'] = left_edge_low
            h['fwtm'] = width_low

@export
@numba.njit(cache=True, nogil=True)
def get_fwxm(hitlet, fraction=0.5):
    """
    Estimates the left and right edge of a specific height percentage.

    :param hitlet: Single hitlet
    :param fraction: Level for which the width shall be computed.
    :returns: Two floats, left edge and right edge in ns

    Notes:
        The function searches for the last sample below and above the
        specified height level on the left and right hand side of the
        maximum. When the samples are found the width is estimated based
        upon a linear interpolation between the respective samples. In
        case, that the samples cannot be found for either one of the
        sides the corresponding outer most bin edges are used: left 0;
        right last sample + 1.
    """
    data = hitlet['data'][:hitlet['length']]

    index_maximum = np.argmax(data)
    max_val = data[index_maximum] * fraction
    if np.all(data > max_val) or np.all(data == 0):
        # In case all samples are larger, FWXM is not definition.
        return np.nan, np.nan

    pre_max = data[:index_maximum]  # Does not include maximum
    post_max = data[1 + index_maximum:]  # same

    if len(pre_max) and np.any(pre_max <= max_val):
        # First the left edge:

        lbi, lbs = _get_fwxm_boundary(pre_max[::-1], max_val)  # Reversing data starting at sample
        # before maximum and go left
        lbi = (index_maximum - 1) - lbi  # start sample minus samples we went to the left
        m = data[lbi + 1] - lbs  # divided by 1 sample
        left_edge = lbi + (max_val - lbs) / m + 0.5
    else:
        # There is no data before the maximum:
        left_edge = 0

    if len(post_max) and np.any(post_max <= max_val):
        # Now the right edge:
        rbi, rbs = _get_fwxm_boundary(post_max, max_val)  # Starting after maximum and go right
        rbi += 1 + index_maximum  # sample to the right plus start
        m = data[rbi - 1] - rbs
        right_edge = rbi - (max_val - rbs) / m + 0.5
    else:
        right_edge = len(data)

    left_edge = left_edge * hitlet['dt']
    right_edge = right_edge * hitlet['dt']
    return left_edge, right_edge

@numba.njit(cache=True, nogil=True)
def _get_fwxm_boundary(data, max_val):
    """
    Returns sample position and height for the last sample which
    amplitude is below the specified value.

    If no sample can be found returns position and value of last sample
    seen.

    Note:
        For FWHM we assume that we start at the maximum.
    """
    ind = None
    s = None
    for i, d in enumerate(data):
        if d <= max_val:
            ind = i
            s = d
            return ind, s
    return len(data)-1, data[-1]

@export
def conditional_entropy(hitlets, template='flat', square_data=False):
    """
    Function which estimates the conditional entropy based on the
    specified template.

    In order to compute the conditional entropy each hitlet will be
    aligned such that its maximum falls into the same sample as for the
    template. If the maximum is ambiguous the first maximum is taken.

    :param hitlets: Hitlets for which the entropy shall be computed.
        Can be any data_kind which offers the fields data and length.
    :param template: Template to compare the data with. Can be either
        specified as "flat" to use a flat distribution or as a numpy
        array containing any normalized template.
    :param square_data: If true data will be squared and normalized
        before estimating the entropy. Otherwise the data will only be
        normalized.
    :returns: Array containing the entropy values for each hitlet.

    Note:
        The template has to be normalized such that its total area is 1.
        Independently of the specified options, only samples for which
        the content is greater zero are used to compute the entropy.

        In case of the non-squared case negative samples are omitted in
        the calculation.
    """
    if not isinstance(template, np.ndarray) and template != 'flat':
        raise ValueError('Template input not understood. Must be either a numpy array,\n',
                         'or "flat".')

    if 'data' not in hitlets.dtype.names:
        raise ValueError('"hitlets" must have a field "data".')

    if isinstance(template, str) and template == 'flat':
        template = np.empty(0, dtype=np.float32)
        flat = True
    else:
        flat = False
    res = _conditional_entropy(hitlets, template, flat=flat,
                               square_data=square_data)
    return res

@numba.njit(cache=True, nogil=True)
def _conditional_entropy(hitlets, template, flat=False, square_data=False):
    res = np.zeros(len(hitlets), dtype=np.float32)
    for ind, h in enumerate(hitlets):
        # h['data'][:] generates just a view and not a copy of the data
        # Since the data of hitlets should not be modified we have
        # to use a copy instead.... See also https://stackoverflow.com/questions/4370745/view-onto-a-numpy-array
        hitlet = np.copy(h['data'][:h['length']])

        # Squaring and normalizing:
        if square_data:
            hitlet[:] = hitlet * hitlet
        if np.sum(hitlet):
            hitlet[:] = hitlet / np.sum(hitlet)
        else:
            # If there is no area we cannot normalize
            res[ind] = np.nan
            continue

        if flat:
            # Take out values which are smaller euqal zero since:
            # lim x_i --> 0, x_i * log(x_i) --> 0
            # and log not defined for negative values.
            m = hitlet > 0
            hitlet = hitlet[m]
            len_hitlet = len(hitlet)

            template = np.ones(len_hitlet, dtype=np.float32)
            template = template / np.sum(template)

            e = - np.sum(hitlet * np.log(hitlet / template))
        else:
            # In case of a template we take out zeros and negative values
            # once we populated the buffer. Otherwise we miss-align
            # template and buffer.

            # create a buffers to align data and template:
            len_hitlet = len(hitlet)
            len_template = len(template)
            length = np.max(np.array([len_hitlet, len_template]))
            length = length * 2 + 1
            buffer = np.zeros((2, length), dtype=np.float32)

            # align data and template and compute entropy:
            si = length // 2 - np.argmax(hitlet)
            ei = si + len(hitlet)
            buffer[0, si:ei] = hitlet[:]

            si = length // 2 - np.argmax(template)
            ei = si + len(template)
            buffer[1, si:ei] = template[:]

            # Remove zeros from buffers:
            m_hit = (buffer[0] > 0)
            m_temp = (buffer[1] > 0)
            m = m_hit & m_temp
            e = - np.sum(buffer[0][m] * np.log(buffer[0][m] / buffer[1][m]))
        res[ind] = e
    return res


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Dec  10 15:24:21 2020

@author: L.I.Vazquez-Salazar
@email: litzavazquezs@gmail.com

Example of a file for the calculation and plotting of KL divergency.
"""
import numpy as np
import pandas as pd
import os
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib
from metrics import Metrics

#Read all the files related to the bonds with C.
bonds = []
for path, dirs, files in os.walk("."):
    for file in files:
        if file.startswith('bonds_C') and file.endswith(".csv"):
            bonds.append(path+"/"+file)

def read_files(array):
    '''
    Read the list of files and separate them by target or test.

    Parameters
    ----------
    array : List of files

    Returns
    -------
    values : Dictionary that refeers to each of the dataframes read it.

    '''
    values = {}
    for item in array:
        if item.find('QM9.csv') != -1:
            values['QM9'] = pd.read_csv(item)
        elif item.find('PC9.csv') != -1:
            values['PC9'] = pd.read_csv(item)
        elif item.find('ANIE.csv') != -1:
            values['ANIE'] = pd.read_csv(item)
        elif item.find('PC9_9p.csv') != -1:
            values['PC9 opt'] = pd.read_csv(item)
        elif item.find('QM9_9p.csv') != -1:
            values['QM9 opt'] = pd.read_csv(item)
        elif item.find('ANIE_9p.csv') != -1:
            values['ANIE opt'] = pd.read_csv(item)

    return values

bonds_read = read_files(bonds)

def getdata(dictionary,key):
    '''
    Separate the data for a type of bond for the different databases evaluated.
    This considers only the reference distribution.

    Parameters
    ----------
    dictionary : Created by read_files. It contains the df for each of the databases.
    key : It is name of the bond on the databases

    Returns
    -------
    df : new database for the type of bond selected.

    '''
    X_ANIE = dictionary['ANIE'][key]
    X_PC9 = dictionary['PC9'][key]
    X_QM9 = dictionary['QM9'][key]

    dictionary = {'QM9':X_QM9, 'PC9':X_PC9, 'ANI-1E':X_ANIE}
    df = pd.DataFrame(dictionary)
    return df

def getdata_test(dictionary,key):
    '''
    Separate the data for a type of bond for the different databases evaluated.
    This considers only the target distribution.

    Parameters
    ----------
    dictionary : Created by read_files. It contains the df for each of the databases.
    key : It is name of the bond on the databases

    Returns
    -------
    df : new database for the type of bond selected.

    '''
    X_Tauto_opt_ANIE = dictionary['ANIE opt'][key]
    X_Tauto_opt_PC9 = dictionary['PC9 opt'][key]
    X_Tauto_opt_QM9 = dictionary['QM9 opt'][key]

    dictionary = {'Test set opt QM9':X_Tauto_opt_QM9, 'Test set opt PC9':X_Tauto_opt_PC9
                  , 'Test set opt ANIE':X_Tauto_opt_ANIE }

    df = pd.DataFrame(dictionary)
    return df

bonds_CC = getdata(bonds_read,'CC')
test_CC = getdata_test(bonds_read,'CC')

bonds_CO = getdata(bonds_read, 'CO')
test_CO = getdata_test(bonds_read, 'CO')

bonds_CN = getdata(bonds_read, 'CN')
test_CN = getdata_test(bonds_read, 'CN')

def compute_KL(reference,target):
    '''
    Read the created databases for a specific type of bond and then compute
    the value of the KL divergency.

    Parameters
    ----------
    reference : Dataframe of the values of the bond lenghts of a specific type of bond
                on the reference databases.
    target : Datafram of the values of the bond lenghts of a specific type of bond
                on the target database.

    Returns
    -------
    X_values : Evaluated values.
    KL_values : Values of the cummulative KL

    '''
    mt = Metrics(reference,target)
    data_QM9 = mt.get_data('QM9','Test set opt QM9')
    data_PC9 = mt.get_data('PC9','Test set opt PC9')
    data_ANIE = mt.get_data('ANI-1E','Test set opt ANIE')

    QM9_x, QM9_kl = mt.KL_divergence_cum(data_QM9)
    PC9_x, PC9_kl = mt.KL_divergence_cum(data_PC9)
    ANIE_x, ANIE_kl = mt.KL_divergence_cum(data_ANIE)

    KL_values = [QM9_kl,PC9_kl,ANIE_kl]
    X_values = [QM9_x[1:],PC9_x[1:], ANIE_x[1:]]
    return X_values, KL_values

X_CC, KL_CC = compute_KL(bonds_CC,test_CC)
X_CN, KL_CN = compute_KL(bonds_CN,test_CN)
X_CO, KL_CO = compute_KL(bonds_CO,test_CO)

x_value = np.linspace(1,2,num=999)
# =============================================================================
#%%
colors = ['firebrick','darkorange','deepskyblue']
colors_rgba = []
for i in colors:
    x = matplotlib.colors.to_rgba(i)
    colors_rgba.append(x)
names = ['QM9 KL', 'PC9 KL', 'ANIE KL']

fig, axs = plt.subplots(nrows=3, ncols=3, figsize=(15,15),sharex='col')
plt.rc('font', family='sans-serif', size=16)
ax1 = sns.kdeplot(data=bonds_CC, ax=axs[0][0], palette=colors_rgba, legend=True, common_norm=False)
ax1.set_title('C-C Bond distribution')
ax1.text(0.85,5,('Reference'),fontsize=35,rotation='vertical')
ax1.set_ylim(0,18)
ax2 = sns.kdeplot(data=test_CC,ax=axs[1][0],fill=False,common_norm=False,legend=False,palette=colors_rgba)
ax2.text(0.85,4.5,('Target'),fontsize=35,rotation='vertical')
ax2.set_ylim(0,18)
ax3 = axs[2][0]
for i in range(0,3):
    ax3.plot(X_CC[i],KL_CC[i], label=names[i],color=colors[i])

ax3.set_xlim(1.1,1.7)
ax3.set_ylabel(r'$D_{KL}(P||Q)$')
ax3.set_xlabel(r'r($\AA$)')
ax3.text(0.85,0.1,('KL-div'),fontsize=35,rotation='vertical')

ax4 = sns.kdeplot(data=bonds_CN, ax=axs[0][1], palette=colors_rgba, legend=True, common_norm=False,)
ax4.set_title('C-N Bond distribution')
ax4.set_ylabel('')
ax4.set_ylim(0,18)
ax5 = sns.kdeplot(data=test_CN,ax=axs[1][1],fill=False,common_norm=False,legend=False,palette=colors_rgba)
ax5.set_ylim(0,18)
ax5.set_ylabel('')
ax6 = axs[2][1]
for i in range(0,3):
    ax6.plot(X_CN[i],KL_CN[i], label=names[i],color=colors[i])

ax6.set_xlim(1.1,1.7)

ax6.set_xlabel(r'r($\AA$)')

ax7 = sns.kdeplot(data=bonds_CO, ax=axs[0][2], palette=colors_rgba, legend=True, common_norm=False)
ax7.set_title('C-O Bond distribution')
ax7.set_ylabel('')
ax7.set_ylim(0,18)

ax8 = sns.kdeplot(data=test_CO,ax=axs[1][2],fill=False,common_norm=False,legend=False,palette=colors_rgba)
ax8.set_ylabel('')
ax8.set_ylim(0,18)

ax9 = axs[2][2]
for i in range(0,3):
    ax9.plot(X_CO[i],KL_CO[i], label=names[i],color=colors[i])

ax9.set_xlim(1.1,1.7)

ax9.set_xlabel(r'r($\AA$)')

plt.subplots_adjust(hspace=0)
#plt.savefig('C-bonds-dist-KL_9p.pdf',bbox_inches='tight')
# plt.show()