from .chain_collection import ChainCollection
import numpy as np
import pandas as pd
from .chain import calculate_charge
from abpytools.utils import DataLoader
from operator import itemgetter
from .fab import Fab
from .helper_functions import germline_identity_pd, to_numbering_table
from .base import CollectionBase
import os
import json
from .utils import (json_FabCollection_formatter, pb2_FabCollection_formatter, pb2_FabCollection_parser,
                    json_FabCollection_parser)
from .flags import *

if BACKEND_FLAGS.HAS_PROTO:
    from abpytools.core.formats import FabCollectionProto

class FabCollection(CollectionBase):

    def __init__(self, fab=None, heavy_chains=None, light_chains=None, names=None):
        """
        Fab object container that handles combinations of light/heavy Chain pairs.

        Args:
            fab (list):
            heavy_chains (ChainCollection):
            light_chains (ChainCollection):
            names (list):
        """
        # check if it's a Chain object
        if heavy_chains is None and light_chains is None and fab is None:
            raise ValueError('Provide a list of Chain objects or an ChainCollection object')

        # check if fab object is a list and if all object are abpytools.Fab objects
        if isinstance(fab, list) and all(isinstance(fab_i, Fab) for fab_i in fab):
            self._fab = fab
            self._light_chains = ChainCollection([x[0] for x in self._fab])
            self._heavy_chains = ChainCollection([x[1] for x in self._fab])

        if fab is None and (heavy_chains is not None and light_chains is not None):

            if isinstance(heavy_chains, list):
                self._heavy_chains = ChainCollection(antibody_objects=heavy_chains)

            elif isinstance(heavy_chains, ChainCollection):
                self._heavy_chains = heavy_chains

            else:
                raise ValueError('Provide a list of Chain objects or an ChainCollection object')

            if isinstance(light_chains, list):
                self._light_chains = ChainCollection(antibody_objects=light_chains)

            elif isinstance(light_chains, ChainCollection):
                self._light_chains = light_chains

            else:
                raise ValueError('Provide a list of Chain objects or an ChainCollection object')

            if len(self._light_chains.loading_status()) == 0:
                self._light_chains.load()

            if len(self._heavy_chains.loading_status()) == 0:
                self._heavy_chains.load()

            if self._light_chains.n_ab != self._heavy_chains.n_ab:
                raise ValueError('Number of heavy chains must be the same of light chains')

        if isinstance(names, list) and all(isinstance(name, str) for name in names):
            if len(names) == self._heavy_chains.n_ab:
                self._names = names
            else:
                raise ValueError(
                    'Length of name list must be the same as length of heavy_chains/light chains lists')

        elif names is None:
            self._names = ['{} - {}'.format(heavy, light) for heavy, light in zip(self._heavy_chains.names,
                                                                                self._light_chains.names)]

        else:
            raise ValueError("Names expected a list of strings, instead got {}".format(type(names)))

        self._n_ab = self._light_chains.n_ab
        self._pair_sequences = [heavy + light for light, heavy in zip(self._heavy_chains.sequences,
                                                                      self._light_chains.sequences)]

        # keep the name of the heavy and light chains internally to keep everything in the right order
        self._internal_heavy_name = self._heavy_chains.names
        self._internal_light_name = self._light_chains.names

    # even though it makes more sense to draw all these values from the base Fab objects this is much slower
    # whenever self._n_ab > 1 it makes more sense to use the self._heavy_chain and self._light_chain containers
    # in all the methods
    # in essence the abpytools.Fab object is just a representative building block that could in future just
    # cache data and would then represent a speed up in the calculations
    def molecular_weights(self, monoisotopic=False):

        return [heavy + light for heavy, light in zip(self._heavy_chains.molecular_weights(monoisotopic=monoisotopic),
                                                      self._light_chains.molecular_weights(monoisotopic=monoisotopic))]

    def extinction_coefficients(self, extinction_coefficient_database='Standard', reduced=False, normalise=False,
                                **kwargs):

        heavy_ec = self._heavy_chains.extinction_coefficients(
            extinction_coefficient_database=extinction_coefficient_database,
            reduced=reduced)
        light_ec = self._light_chains.extinction_coefficients(
            extinction_coefficient_database=extinction_coefficient_database,
            reduced=reduced)

        if normalise:
            return [(heavy + light) / mw for heavy, light, mw in
                    zip(heavy_ec, light_ec, self.molecular_weights(**kwargs))]
        else:
            return [heavy + light for heavy, light in zip(heavy_ec, light_ec)]

    def hydrophobicity_matrix(self):

        return np.column_stack((self._heavy_chains.hydrophobicity_matrix(), self._light_chains.hydrophobicity_matrix()))

    def charge(self):

        return np.column_stack((self._heavy_chains.charge, self._light_chains.charge))

    def total_charge(self, ph=7.4, pka_database='Wikipedia'):

        available_pi_databases = ["EMBOSS", "DTASetect", "Solomon", "Sillero", "Rodwell", "Wikipedia", "Lehninger",
                                  "Grimsley"]
        assert pka_database in available_pi_databases, \
            "Selected pI database {} not available. Available databases: {}".format(pka_database,
                                                                                ' ,'.join(available_pi_databases))

        data_loader = DataLoader(data_type='AminoAcidProperties', data=['pI', pka_database])
        pka_data = data_loader.get_data()

        return [calculate_charge(sequence=seq, ph=ph, pka_values=pka_data) for seq in self.sequences]

    def igblast_local_query(self, file_path, chain):

        if chain.lower() == 'light':
            self._light_chains.igblast_local_query(file_path=file_path)
        elif chain.lower() == 'heavy':
            self._heavy_chains.igblast_local_query(file_path=file_path)
        else:
            raise ValueError('Specify if the data being loaded is for the heavy or light chain')

    def igblast_server_query(self, **kwargs):
        self._light_chains.igblast_server_query(**kwargs)
        self._heavy_chains.igblast_server_query(**kwargs)

    def numbering_table(self, as_array=False, region='all', chain='both', **kwargs):

        return to_numbering_table(as_array=as_array, region=region, chain=chain,
                                  heavy_chains_numbering_table=self._heavy_chains.numbering_table,
                                  light_chains_numbering_table=self._light_chains.numbering_table,
                                  names=self.names, **kwargs)

    def _germline_pd(self):

        # empty dictionaries return false, so this condition checks if any of the values are False
        if all([x for x in self._light_chains.germline_identity.values()]) is False:
            # this means there is no information about the germline,
            # by default it will run a web query
            self._light_chains.igblast_server_query()
        if all([x for x in self._heavy_chains.germline_identity.values()]) is False:
            self._heavy_chains.igblast_server_query()

        heavy_chain_germlines = self._heavy_chains.germline
        light_chain_germlines = self._light_chains.germline

        data = np.array([[heavy_chain_germlines[x][0] for x in self._internal_heavy_name],
                         [heavy_chain_germlines[x][1] for x in self._internal_heavy_name],
                         [light_chain_germlines[x][0] for x in self._internal_light_name],
                         [light_chain_germlines[x][1] for x in self._internal_light_name]]).T

        df = pd.DataFrame(data=data,
                          columns=pd.MultiIndex.from_tuples([('Heavy', 'Assignment'),
                                                             ('Heavy', 'Score'),
                                                             ('Light', 'Assignment'),
                                                             ('Light', 'Score')]),
                          index=self.names)

        df.loc[:, (slice(None), 'Score')] = df.loc[:, (slice(None), 'Score')].apply(pd.to_numeric)

        return df

    def save_to_json(self, path, update=True):
        with open(os.path.join(path + '.json'), 'w') as f:
            fab_data = json_FabCollection_formatter(self)
            json.dump(fab_data, f, indent=2)

    def save_to_pb2(self, path, update=True):
        proto_parser = FabCollectionProto()
        try:
            with open(os.path.join(path + '.pb2'), 'rb') as f:
                proto_parser.ParseFromString(f.read())
        except IOError:
            # Creating new file
            pass

        pb2_FabCollection_formatter(self, proto_parser)

        with open(os.path.join(path + '.pb2'), 'wb') as f:
            f.write(proto_parser.SerializeToString())

    def save_to_fasta(self, path, update=True):
        raise NotImplementedError

    @classmethod
    def load_from_json(cls, path, n_threads=20, verbose=True, show_progressbar=True):
        with open(path, 'r') as f:
            data = json.load(f)

        fab_objects = json_FabCollection_parser(data)

        fab_collection = cls(fab=fab_objects)

        return fab_collection

    @classmethod
    def load_from_pb2(cls, path, n_threads=20, verbose=True, show_progressbar=True):
        with open(path, 'rb') as f:
            proto_parser = FabCollectionProto()
            proto_parser.ParseFromString(f.read())

        fab_objects = pb2_FabCollection_parser(proto_parser)

        fab_collection = cls(fab=fab_objects)

        return fab_collection

    @classmethod
    def load_from_fasta(cls, path, numbering_scheme=NUMBERING_FLAGS.CHOTHIA, n_threads=20,
                        verbose=True, show_progressbar=True):
        raise NotImplementedError

    def _get_names_iter(self, chain='both'):
        if chain == 'both':
            for light_chain, heavy_chain in zip(self._light_chains, self._heavy_chains):
                yield f"{light_chain.name}-{heavy_chain.name}"
        elif chain == 'light':
            for light_chain in self._light_chains:
                yield light_chain.name
        elif chain == 'heavy':
            for heavy_chain in self._heavy_chains:
                yield heavy_chain.name
        else:
            raise ValueError(f"Unknown chain type ({chain}), available options are:"
                             f"both, light or heavy.")

    @property
    def regions(self):
        heavy_regions = self._heavy_chains.ab_region_index()
        light_regions = self._light_chains.ab_region_index()

        return {name: {CHAIN_FLAGS.HEAVY_CHAIN: heavy_regions[heavy],
                       CHAIN_FLAGS.LIGHT_CHAIN: light_regions[light]} for name, heavy, light in
                zip(self.names, self._internal_heavy_name, self._internal_light_name)}

    @property
    def names(self):
        return self._names

    @property
    def sequences(self):
        return self._pair_sequences

    @property
    def aligned_sequences(self):
        return [heavy + light for light, heavy in
                zip(self._heavy_chains.aligned_sequences,
                    self._light_chains.aligned_sequences)]

    @property
    def n_ab(self):
        return self._n_ab

    @property
    def germline_identity(self):
        return self._germline_identity()

    @property
    def germline(self):
        return self._germline_pd()

    def _string_summary_basic(self):
        return "abpytools.FabCollection Number of sequences: {}".format(self._n_ab)

    def __len__(self):
        return self._n_ab

    def __repr__(self):
        return "<%s at 0x%02x>" % (self._string_summary_basic(), id(self))

    def __getitem__(self, indices):
        if isinstance(indices, int):
            return Fab(heavy_chain=self._heavy_chains[indices],
                       light_chain=self._light_chains[indices],
                       name=self.names[indices], load=False)
        else:
            return FabCollection(heavy_chains=list(itemgetter(*indices)(self._heavy_chains)),
                                 light_chains=list(itemgetter(*indices)(self._light_chains)),
                                 names=list(itemgetter(*indices)(self._names)))

    def _germline_identity(self):

        # empty dictionaries return false, so this condition checks if any of the values are False
        if all([x for x in self._light_chains.germline_identity.values()]) is False:
            # this means there is no information about the germline,
            # by default it will run a web query
            self._light_chains.igblast_server_query()
        if all([x for x in self._heavy_chains.germline_identity.values()]) is False:
            self._heavy_chains.igblast_server_query()

        return germline_identity_pd(self._heavy_chains.germline_identity,
                                    self._light_chains.germline_identity,
                                    self._internal_heavy_name,
                                    self._internal_light_name,
                                    self._names)

    def get_object(self, name):

        """

        :param name: str
        :return:
        """

        if name in self.names:
            index = self.names.index(name)
            return self[index]
        else:
            raise ValueError('Could not find sequence with specified name')


import numpy
import xraylib

"""
transfocator_id30b : transfocator for id13b:
        It can:

            1) guess the lens configuration (number of lenses for each type) for a given photon energy
            and target image size. Use transfocator_compute_configuration() for this task

            2) for a given transfocator configuration, compute the main optical parameters
                (image size, focal distance, focal position and divergence).
                Use transfocator_compute_parameters() for this task

            3) Performs full ray tracing. Use id30b_ray_tracing() for this task

        Note that for the optimization and parameters calculations the transfocator configuration is
        given in keywords. For ray tracing calculations many parameters of the transfocator are hard coded
        with the values of id30b

        See main program for examples.

        Dependencies:
            Numpy
            xraylib (to compute refracion indices)
            Shadow (for ray tracing only)
            matplotlib (for some plots of ray=tracing)

        Side effects:
            When running ray tracing some files are created.

        MODIFICATION HISTORY:
           2015-03-25 srio@esrf.eu, written

"""

__author__ = "Manuel Sanchez del Rio"
__contact__ = "srio@esrf.eu"
__copyright__ = "ESRF, 2015"

def transfocator_compute_configuration(photon_energy_ev,s_target,\
            symbol=["Be","Be","Be"], density=[1.845,1.845,1.845],\
            nlenses_max = [15,3,1], nlenses_radii = [500e-4,1000e-4,1500e-4], lens_diameter=0.05, \
            sigmaz=6.46e-4, alpha = 0.55, \
            tf_p=5960, tf_q=3800, verbose=1 ):
    """
    Computes the optimum transfocator configuration for a given photon energy and target image size.

    All length units are cm

    :param photon_energy_ev: the photon energy in eV
    :param s_target:       the target image size in cm.
    :param symbol:         the chemical symbol of the lens material of each type. Default symbol=["Be","Be","Be"]
    :param density:        the density of each type of lens. Default: density=[1.845,1.845,1.845]
    :param nlenses_max:    the maximum allowed number of lenases for each type of lens. nlenses_max = [15,3,1]
    :param nlenses_radii:  the radii in cm of each type of lens. Default: nlenses_radii = [500e-4,1000e-4,1500e-4]
    :param lens_diameter:    the physical diameter (acceptance) in cm of the lenses. If different for each type of lens,
                            consider the smaller one. Default: lens_diameter=0.05
    :param sigmaz:         the sigma (standard deviation) of the source in cm
    :param alpha:          an adjustable parameter in [0,1](see doc). Default: 0.55 (it is 0.76 for pure Gaussian beams)
    :param tf_p:           the distance source-transfocator in cm
    :param tf_q:           the distance transfocator-image in cm
    :param:verbose:        set to 1 for verbose text output
    :return:               a list with the number of lenses of each type.

    """
    if s_target < 2.35*sigmaz*tf_q/tf_p:
        print("Source size FWHM is: %f um"%(1e4*2.35*sigmaz))
        print("Maximum Demagnifications is: %f um"%(tf_p/tf_q))
        print("Minimum possible size is: %f um"%(1e4*2.35*sigmaz*tf_q/tf_p))
        print("Error: redefine size")
        return None

    deltas = [(1.0 - xraylib.Refractive_Index_Re(symbol[i],photon_energy_ev*1e-3,density[i])) \
              for i in range(len(symbol))]

    focal_q_target = _tansfocator_guess_focal_position( s_target, p=tf_p, q=tf_q, sigmaz=sigmaz, alpha=alpha, \
                                           lens_diameter=lens_diameter,method=2)

    focal_f_target = 1.0 / (1.0/focal_q_target + 1.0/tf_p)
    div_q_target = alpha  * lens_diameter / focal_q_target

    #corrections for extreme cases
    source_demagnified = 2.35*sigmaz*focal_q_target/tf_p
    if source_demagnified > lens_diameter: source_demagnified = lens_diameter

    s_target_calc = numpy.sqrt( (div_q_target*(tf_q-focal_q_target))**2 + source_demagnified**2)

    nlenses_target = _transfocator_guess_configuration(focal_f_target,deltas=deltas,\
                                                   nlenses_max=nlenses_max,radii=nlenses_radii, )
    if verbose:
        print("transfocator_compute_configuration: focal_f_target: %f"%(focal_f_target))
        print("transfocator_compute_configuration: focal_q_target: %f cm"%(focal_q_target))
        print("transfocator_compute_configuration: s_target: %f um"%(s_target_calc*1e4))
        print("transfocator_compute_configuration: nlenses_target: ",nlenses_target)

    return nlenses_target

def transfocator_compute_parameters(photon_energy_ev, nlenses_target,\
            symbol=["Be","Be","Be"], density=[1.845,1.845,1.845],\
            nlenses_max = [15,3,1], nlenses_radii = [500e-4,1000e-4,1500e-4], lens_diameter=0.05, \
            sigmaz=6.46e-4, alpha = 0.55, \
            tf_p=5960, tf_q=3800 ):
    """
    Computes the parameters of the optical performances of a given transgocator configuration.

    returns a l

    All length units are cm

    :param photon_energy_ev:
    :param nlenses_target: a list with the lens configuration, i.e. the number of lenses of each type.
    :param symbol:         the chemical symbol of the lens material of each type. Default symbol=["Be","Be","Be"]
    :param density:        the density of each type of lens. Default: density=[1.845,1.845,1.845]
    :param nlenses_max:    the maximum allowed number of lenases for each type of lens. nlenses_max = [15,3,1]
                           TODO: remove (not used)
    :param nlenses_radii:  the radii in cm of each type of lens. Default: nlenses_radii = [500e-4,1000e-4,1500e-4]
    :param lens_diameter:    the physical diameter (acceptance) in cm of the lenses. If different for each type of lens,
                            consider the smaller one. Default: lens_diameter=0.05
    :param sigmaz:         the sigma (standard deviation) of the source in cm
    :param alpha:          an adjustable parameter in [0,1](see doc). Default: 0.55 (it is 0.76 for pure Gaussian beams)
    :param tf_p:           the distance source-transfocator in cm
    :param tf_q:           the distance transfocator-image in cm
    :return:               a list with parameters (image_siza, lens_focal_distance,
                           focal_position from transfocator center, divergence of beam after the transfocator)
    """

    deltas = [(1.0 - xraylib.Refractive_Index_Re(symbol[i],photon_energy_ev*1e-3,density[i])) \
              for i in range(len(symbol))]
    focal_f = _transfocator_calculate_focal_distance( deltas=deltas,\
                                                       nlenses=nlenses_target,radii=nlenses_radii)
    focal_q = 1.0 / (1.0/focal_f - 1.0/tf_p)
    div_q = alpha  * lens_diameter / focal_q
    #corrections
    source_demagnified = 2.35*sigmaz*focal_q/tf_p
    if source_demagnified > lens_diameter: source_demagnified = lens_diameter
    s_target = numpy.sqrt( (div_q*(tf_q-focal_q))**2 + (source_demagnified)**2 )
    return (s_target,focal_f,focal_q,div_q)

def transfocator_nlenses_to_slots(nlenses,nlenses_max=None):
    """
    converts the transfocator configuration from a list of the number of lenses of each type,
    into a list of active (1) or inactive (0) actuators for the slots.

    :param nlenses: the list with number of lenses (e.g., [5,2,0]
    :param nlenses_max: the maximum number of lenses of each type, usually powers of two minus one.
                        E.g. [15,3,1]
    :return: a list of on (1) and off (0) slots, e.g., [1, 0, 1, 0, 0, 1, 0]
            (first type: 1*1+0*2+1*4+0*8=5, second type: 0*1+1*2=2, third type: 0*1=0)
    """
    if nlenses_max == None:
        nlenses_max = nlenses

    ss = []
    for i,iopt in enumerate(nlenses):
        if iopt > nlenses_max[i]:
            print("Error: i:%d, nlenses: %d, nlenses_max: %d"%(i,iopt,nlenses_max[i]))
        ncharacters = len("{0:b}".format(nlenses_max[i]))

        si = list( ("{0:0%db}"%(ncharacters)).format(int(iopt)) )
        si.reverse()
        ss += si

    on_off = [int(i) for i in ss]
    #print("transfocator_nlenses_to_slots: nlenses_max: ",nlenses_max," nlenses: ",nlenses," slots: ",on_off)
    return on_off

def _transfocator_calculate_focal_distance(deltas=[0.999998],nlenses=[1],radii=[500e-4]):

    inverse_focal_distance = 0.0
    for i,nlensesi in enumerate(nlenses):
        if nlensesi > 0:
            focal_distance_i = radii[i] / (2.*nlensesi*deltas[i])
            inverse_focal_distance += 1.0/focal_distance_i
    if inverse_focal_distance == 0:
        return 99999999999999999999999999.
    else:
        return 1.0/inverse_focal_distance

def _tansfocator_guess_focal_position( s_target, p=5960., q=3800.0, sigmaz=6.46e-4, \
                                      alpha=0.66, lens_diameter=0.05, method=2):
    x = 1e15
    if method == 1: # simple sum
        AA = 2.35*sigmaz/p
        BB = -(s_target + alpha * lens_diameter)
        CC = alpha*lens_diameter*q

        cc = numpy.roots([AA,BB,CC])
        x = cc[1]
        return x

    if method == 2: # sum in quadrature
        AA = ( (2.35*sigmaz)**2)/(p**2)
        BB = 0.0
        CC = alpha**2 * lens_diameter**2 - s_target**2
        DD = - 2.0 * alpha**2 * lens_diameter**2 * q
        EE = alpha**2 * lens_diameter**2 * q**2
        cc = numpy.roots([AA,BB,CC,DD,EE])
        for i,cci in enumerate(cc):
            if numpy.imag(cci) == 0:
                return numpy.real(cci)

    return x

def _transfocator_guess_configuration(focal_f_target,deltas=[0.999998],nlenses_max=[15],radii=[500e-4]):

    nn = len(nlenses_max)

    ncombinations = (1+nlenses_max[0]) * (1+nlenses_max[1]) * (1+nlenses_max[2])

    icombinations = 0
    aa = numpy.zeros((3,ncombinations),dtype=int)
    bb = numpy.zeros(ncombinations)
    for i0 in range(1+nlenses_max[0]):
        for i1 in range(1+nlenses_max[1]):
            for i2 in range(1+nlenses_max[2]):
                aa[0,icombinations] = i0
                aa[1,icombinations] = i1
                aa[2,icombinations] = i2
                bb[icombinations] = focal_f_target - _transfocator_calculate_focal_distance(deltas=deltas,nlenses=[i0,i1,i2],radii=radii)
                icombinations += 1
    bb1 = numpy.abs(bb)
    ibest = bb1.argmin()

    return (aa[:,ibest]).tolist()

#
#
#

def id30b_ray_tracing(emittH=4e-9,emittV=1e-11,betaH=35.6,betaV=3.0,number_of_rays=50000,\
                      density=1.845,symbol="Be",tf_p=1000.0,tf_q=1000.0,lens_diameter=0.05,\
                      slots_max=None,slots_on_off=None,photon_energy_ev=14000.0,\
                      slots_lens_thickness=None,slots_steps=None,slots_radii=None,\
                      s_target=10e-4,focal_f=10.0,focal_q=10.0,div_q=1e-6):

    #=======================================================================================================================
    # Gaussian undulator source
    #=======================================================================================================================

    import Shadow
    #import Shadow.ShadowPreprocessorsXraylib as sx

    sigmaXp = numpy.sqrt(emittH/betaH)
    sigmaZp = numpy.sqrt(emittV/betaV)
    sigmaX = emittH/sigmaXp
    sigmaZ = emittV/sigmaZp

    print("\n\nElectron sizes H:%f um, V:%fu m;\nelectron divergences: H:%f urad, V:%f urad"%\
          (sigmaX*1e6, sigmaZ*1e6, sigmaXp*1e6, sigmaZp*1e6))

    # set Gaussian source
    src = Shadow.Source()
    src.set_energy_monochromatic(photon_energy_ev)
    src.set_gauss(sigmaX*1e2,sigmaZ*1e2,sigmaXp,sigmaZp)

    print("\n\nElectron sizes stored H:%f um, V:%f um;\nelectron divergences: H:%f urad, V:%f urad"%\
          (src.SIGMAX*1e4,src.SIGMAZ*1e4,src.SIGDIX*1e6,src.SIGDIZ*1e6))

    src.apply_gaussian_undulator(undulator_length_in_m=2.8, user_unit_to_m=1e-2, verbose=1)

    print("\n\nElectron sizes stored (undulator) H:%f um, V:%f um;\nelectron divergences: H:%f urad, V:%f urad"%\
          (src.SIGMAX*1e4,src.SIGMAZ*1e4,src.SIGDIX*1e6,src.SIGDIZ*1e6))

    print("\n\nSource size in vertical FWHM: %f um\n"%\
          (2.35*src.SIGMAZ*1e4))

    src.NPOINT = number_of_rays
    src.ISTAR1 = 0 # 677543155

    src.write("start.00")

    # create source
    beam = Shadow.Beam()
    beam.genSource(src)
    beam.write("begin.dat")
    src.write("end.00")

    #=======================================================================================================================
    # complete the (detailed) transfocator description
    #=======================================================================================================================

    print("\nSetting detailed Transfocator for ID30B")

    slots_nlenses = numpy.array(slots_max)*numpy.array(slots_on_off)
    slots_empty = (numpy.array(slots_max)-slots_nlenses)

    #
    ####interactive=True, SYMBOL="SiC",DENSITY=3.217,FILE="prerefl.dat",E_MIN=100.0,E_MAX=20000.0,E_STEP=100.0

    Shadow.ShadowPreprocessorsXraylib.prerefl(interactive=False,E_MIN=2000.0,E_MAX=55000.0,E_STEP=100.0,\
                                              DENSITY=density,SYMBOL=symbol,FILE="Be2_55.dat" )

    nslots = len(slots_max)
    prerefl_file = ["Be2_55.dat" for i in range(nslots)]

    print("slots_max:     ",slots_max)
    #print("slots_target:  ",slots_target)
    print("slots_on_off:  ",slots_on_off)
    print("slots_steps:   ",slots_steps)
    print("slots_radii:   ",slots_radii)
    print("slots_nlenses: ",slots_nlenses)
    print("slots_empty:   ",slots_empty)

    #calculate distances, nlenses and slots_empty

    # these are distances p and q with TF length removed
    tf_length = numpy.array(slots_steps).sum()  #tf length in cm
    tf_fs_before = tf_p - 0.5*tf_length     #distance from source to center of transfocator
    tf_fs_after  = tf_q - 0.5*tf_length     # distance from center of transfocator to image

    # for each slot, these are the empty distances before and after the lenses
    tf_p0 = numpy.zeros(nslots)
    tf_q0 = numpy.array(slots_steps) - (numpy.array(slots_max) * slots_lens_thickness)
    # add now the p q distances
    tf_p0[0]  += tf_fs_before
    tf_q0[-1] += tf_fs_after

    print("tf_p0:   ",tf_p0)
    print("tf_q0:   ",tf_q0)
    print("tf_length: %f cm"%(tf_length))

    # build transfocator
    tf = Shadow.CompoundOE(name='TF ID30B')

    tf.append_transfocator(tf_p0.tolist(), tf_q0.tolist(), \
                    nlenses=slots_nlenses.tolist(), radius=slots_radii, slots_empty=slots_empty.tolist(),\
                    thickness=slots_lens_thickness, prerefl_file=prerefl_file,\
                    surface_shape=4, convex_to_the_beam=0, diameter=lens_diameter,\
                    cylinder_angle=0.0,interthickness=50e-4,use_ccc=0)

    itmp = input("SHADOW Source complete. Do you want to run SHADOR trace? [1=Yes,0=No]: ")
    if str(itmp) != "1":
        return

    #trace system
    tf.dump_systemfile()
    beam.traceCompoundOE(tf,write_start_files=0,write_end_files=0,write_star_files=0, write_mirr_files=0)

    #write only last result file
    beam.write("star_tf.dat")
    print("\nFile written to disk: star_tf.dat")

    #
    # #ideal calculations
    #

    print("\n\n\n")
    print("=============================================== TRANSFOCATOR OUTPUTS ==========================================")

    print("\nTHEORETICAL results: ")

    print("REMIND-----With these lenses we obtained (analytically): ")
    print("REMIND-----  focal_f: %f cm"%(focal_f))
    print("REMIND-----  focal_q: %f cm"%(focal_q))
    print("REMIND-----  s_target: %f um"%(s_target*1e4))

    demagnification_factor = tf_p/focal_q
    theoretical_focal_size = src.SIGMAZ*2.35/demagnification_factor

    # analyze shadow results
    print("\nSHADOW results: ")
    st1 = beam.get_standard_deviation(3,ref=0)
    st2 = beam.get_standard_deviation(3,ref=1)

    print("  stDev*2.35: unweighted: %f um, weighted: %f um "%(st1*2.35*1e4,st2*2.35*1e4))

    tk = beam.histo1(3, nbins=75, ref=1, nolost=1, write="HISTO1")

    print("  Histogram FWHM: %f um "%(1e4*tk["fwhm"]))

    print("  Transmitted intensity: %f (source was: %d) (transmission is %f %%) "%(beam.intensity(nolost=1), src.NPOINT, beam.intensity(nolost=1)/src.NPOINT*100))

    #scan around image
    xx1 = numpy.linspace(0.0,1.1*tf_fs_after,11) # position from TF exit plane

    #xx0 = focal_q - tf_length*0.5
    xx0 = focal_q - tf_length*0.5 # position of focus from TF exit plane

    xx2 = numpy.linspace(xx0-100.0,xx0+100,21) # position from TF exit plane
    xx3 = numpy.array([tf_fs_after])

    xx = numpy.concatenate(([-0.5*tf_length],xx1,xx2,[tf_fs_after]))

    xx.sort()

    f = open("id30b.spec","w")
    f.write("#F id30b.spec\n")
    f.write("\n#S 1 calculations for id30b transfocator\n")
    f.write("#N 8\n")
    labels = "  %18s  %18s   %18s   %18s    %18s    %18s    %18s    %18s"%\
             ("pos from source","pos from image","[pos from TF]", "pos from TF center", "pos from focus",\
              "fwhm shadow(stdev)","fwhm shadow(histo)","fwhm theoretical")

    f.write("#L "+labels+"\n")

    out = numpy.zeros((8,xx.size))
    for i,pos in enumerate(xx):
        beam2 = beam.duplicate()
        beam2.retrace(-tf_fs_after+pos)
        fwhm1 = 2.35*1e4*beam2.get_standard_deviation(3,ref=1,nolost=1)
        tk = beam2.histo1(3, nbins=75, ref=1, nolost=1)
        fwhm2 = 1e4*tk["fwhm"]
        #fwhm_th = 1e4*transfocator_calculate_estimated_size(pos,diameter=diameter,focal_distance=focal_q)
        fwhm_th2 = 1e4*numpy.sqrt( (div_q*(pos+0.5*tf_length-focal_q))**2 + theoretical_focal_size**2 )
        #fwhm_th2 = 1e4*( numpy.abs(div_q*(pos-focal_q+0.5*tf_length)) + theoretical_focal_size )

        out[0,i] = tf_fs_before+tf_length+pos
        out[1,i] = -tf_fs_after+pos
        out[2,i] = pos
        out[3,i] = pos+0.5*tf_length
        out[4,i] = pos+0.5*tf_length-focal_q

        out[5,i] = fwhm1
        out[6,i] = fwhm2
        out[7,i] = fwhm_th2

        f.write(" %18.3f   %18.3f   %18.3f   %18.3f    %18.3f    %18.3f    %18.3f    %18.3f \n"%\
                (tf_fs_before+tf_length+pos,\
                 -tf_fs_after+pos,\
                 pos,\
                 pos+0.5*tf_length,\
                 pos+0.5*tf_length-focal_q,\
                 fwhm1,fwhm2,fwhm_th2))

    f.close()
    print("File with beam evolution written to disk: id30b.spec")

    #
    # plots
    #
    itmp = input("Do you want to plot the intensity distribution and beam evolution? [1=yes,0=No]")
    if str(itmp) != "1":
        return

    import matplotlib.pylab as plt

    plt.figure(1)
    plt.plot(out[1,:],out[5,:],'blue',label="fwhm shadow(stdev)")
    plt.plot(out[1,:],out[6,:],'green',label="fwhm shadow(histo1)")
    plt.plot(out[1,:],out[7,:],'red',label="fwhm theoretical")

    plt.xlabel("Distance from image plane [cm]")
    plt.ylabel("spot size [um] ")
    ax = plt.subplot(111)
    ax.legend(bbox_to_anchor=(1.1, 1.05))

    print("Kill graphic to continue.")
    plt.show()

    Shadow.ShadowTools.histo1(beam,3,nbins=75,ref=1,nolost=1,calfwhm=1)

    input("<Enter> to finish.")
    return None

def id30b_full_simulation(photon_energy_ev=14000.0,s_target=20.0e-4,nlenses_target=None):

    if nlenses_target == None:
        force_nlenses = 0
    else:
        force_nlenses = 1

    #
    # define lens setup (general)
    #
    xrl_symbol = ["Be","Be","Be"]
    xrl_density = [1.845,1.845,1.845]
    lens_diameter = 0.05
    nlenses_max = [15,3,1]
    nlenses_radii = [500e-4,1000e-4,1500e-4]

    sigmaz=6.46e-4
    alpha = 0.55

    tf_p = 5960 # position of the TF measured from the center of the transfocator
    tf_q = 9760 - tf_p # position of the image plane measured from the center of the transfocator

    if s_target < 2.35*sigmaz*tf_q/tf_p:
        print("Source size FWHM is: %f um"%(1e4*2.35*sigmaz))
        print("Maximum Demagnifications is: %f um"%(tf_p/tf_q))
        print("Minimum possible size is: %f um"%(1e4*2.35*sigmaz*tf_q/tf_p))
        print("Error: redefine size")
        return

    print("================================== TRANSFOCATOR INPUTS ")
    print("Photon energy: %f eV"%(photon_energy_ev))
    if force_nlenses:
        print("Forced_nlenses: ",nlenses_target)
    else:
        print("target size: %f cm"%(s_target))

    print("materials: ",xrl_symbol)
    print("densities: ",xrl_density)
    print("Lens diameter: %f cm"%(lens_diameter))
    print("nlenses_max:",nlenses_max,"nlenses_radii: ",nlenses_radii)

    print("Source size (sigma): %f um, FWHM: %f um"%(1e4*sigmaz,2.35*1e4*sigmaz))
    print("Distances: tf_p: %f cm, tf_q: %f cm"%(tf_p,tf_q))
    print("alpha: %f"%(alpha))

    print("========================================================")

    if force_nlenses != 1:
        nlenses_target =  transfocator_compute_configuration(photon_energy_ev,s_target,\
            symbol=xrl_symbol,density=xrl_density,\
            nlenses_max=nlenses_max, nlenses_radii=nlenses_radii, lens_diameter=lens_diameter, \
            sigmaz=sigmaz, alpha=alpha, \
            tf_p=tf_p,tf_q=tf_q, verbose=1)

    (s_target,focal_f,focal_q,div_q) = \
        transfocator_compute_parameters(photon_energy_ev, nlenses_target,\
                                        symbol=xrl_symbol,density=xrl_density,\
                                        nlenses_max=nlenses_max, nlenses_radii=nlenses_radii, \
                                        lens_diameter=lens_diameter,\
                                        sigmaz=sigmaz, alpha=alpha,\
                                        tf_p=tf_p,tf_q=tf_q)

    slots_max   = [  1,  2,  4,  8,   1,   2,   1]  # slots
    slots_on_off  = transfocator_nlenses_to_slots(nlenses_target,nlenses_max=nlenses_max)

    print("=============================== TRANSFOCATOR SET")
    #print("deltas: ",deltas)

    if force_nlenses != 1:
        print("nlenses_target (optimized): ",nlenses_target)
    else:
        print("nlenses_target (forced): ",nlenses_target)

    print("With these lenses we obtain: ")
    print("  focal_f: %f cm"%(focal_f))
    print("  focal_q: %f cm"%(focal_q))
    print("  s_target: %f um"%(s_target*1e4))
    print("  slots_max:    ",slots_max)
    print("  slots_on_off: ",slots_on_off)
    print("==================================================")

    # for theoretical calculations use the focal position and distances given by the target nlenses

    itmp = input("Start SHADOW simulation? [1=yes,0=No]: ")
    if str(itmp) != "1":
        return

    #=======================================================================================================================
    # Inputs
    #=======================================================================================================================

    emittH = 3.9e-9
    emittV = 10e-12
    betaH = 35.6
    betaV = 3.0
    number_of_rays = 50000

    nslots = len(slots_max)
    slots_lens_thickness = [0.3 for i in range(nslots)]   #total thickness of a single lens in cm
    # for each slot, positional gap  of the first lens in cm
    slots_steps    = [  4,   4, 1.9, 6.1,   4,   4, slots_lens_thickness[-1]]
    slots_radii   = [.05, .05, .05, .05, 0.1, 0.1, 0.15]  # radii of the lenses in cm
    AAA= 333
    id30b_ray_tracing(emittH=emittH,emittV=emittV,betaH=betaH,betaV=betaV,number_of_rays=number_of_rays,\
                      density=xrl_density[0],symbol=xrl_symbol[0],tf_p=tf_p,tf_q=tf_q,lens_diameter=lens_diameter,\
                      slots_max=slots_max,slots_on_off=slots_on_off,photon_energy_ev=photon_energy_ev,\
                      slots_lens_thickness=slots_lens_thickness,slots_steps=slots_steps,slots_radii=slots_radii,\
                      s_target=s_target,focal_f=focal_f,focal_q=focal_q,div_q=div_q)

def main():

    # this performs the full simulation: calculates the optimum configuration and do the ray-tracing

    itmp = input("Enter: \n  0 = optimization calculation only \n  1 = full simulation (ray tracing) \n?> ")
    photon_energy_kev = float(input("Enter photon energy in keV: "))
    s_target_um = float(input("Enter target focal dimension in microns: "))
    if str(itmp) == "1":

        id30b_full_simulation(photon_energy_ev=photon_energy_kev*1e3,s_target=s_target_um*1e-4,nlenses_target=None)
        #id30b_full_simulation(photon_energy_ev=14000.0,s_target=20.0e-4,nlenses_target=[3,1,1])
    else:
        #this performs the calculation of the optimizad configuration

        nlenses_optimum = transfocator_compute_configuration(photon_energy_kev*1e3,s_target_um*1e-4,\
                symbol=["Be","Be","Be"], density=[1.845,1.845,1.845],\
                nlenses_max = [15,3,1], nlenses_radii = [500e-4,1000e-4,1500e-4], lens_diameter=0.05, \
                sigmaz=6.46e-4, alpha = 0.55, \
                tf_p=5960, tf_q=3800, verbose=0 )
        print("Optimum lens configuration is: ",nlenses_optimum)

        if nlenses_optimum == None:
            return

        print("Activate slots: ",transfocator_nlenses_to_slots(nlenses_optimum,nlenses_max=[15,3,1]))

        # this calculates the parameters (image size, etc) for a given lens configuration

        (size, f, q_f, div) = transfocator_compute_parameters(photon_energy_kev*1e3, nlenses_optimum,\
                symbol=["Be","Be","Be"], density=[1.845,1.845,1.845],\
                nlenses_max = [15,3,1], nlenses_radii = [500e-4,1000e-4,1500e-4], lens_diameter=0.05, \
                sigmaz=6.46e-4, alpha = 0.55, \
                tf_p=5960, tf_q=3800 )
        print("For given configuration ",nlenses_optimum," we get: ")
        print("  size: %f cm, focal length: %f cm, focal distance: %f cm, divergence: %f rad: "%(size, f, q_f, div))

if __name__ == "__main__":
    main()


import os
import pandas as pd
import matplotlib.pyplot as plt

import seaborn
seaborn.set()

path= os.path.expanduser("~/Desktop/ece671/udpt8")
num_files = len([f for f in os.listdir(path)if os.path.isfile(os.path.join(path, f))])
print(num_files)
u8=[]
i=0
def file_len(fname):
    with open(fname) as f:
        for i, l in enumerate(f):
            pass
    return i + 1
while i<(num_files/2) :
 #      df+=[]
    j=i+1

    path ="/home/vetri/Desktop/ece671/udpt8/ftotal."+str(j)+".csv"
    y = file_len(path)
    #    except:       pass
    #df.append(pd.read_csv(path,header=None))
#       a+=[]
    #y=len(df[i].index)-1 #1 row added by default so that table has a entry
    if y<0:
        y=0

    u8.append(y)

    i+=1
print(u8)
path= os.path.expanduser("~/Desktop/ece671/udpnone")
num_files = len([f for f in os.listdir(path)if os.path.isfile(os.path.join(path, f))])
print(num_files)
i=0
j=0
u=[]
while i<(num_files/2):
    j=i+1

    path ="/home/vetri/Desktop/ece671/udpnone/ftotal."+str(j)+".csv"
    y = file_len(path)
    #    except:       pass
    #df.append(pd.read_csv(path,header=None))
#       a+=[]
    #y=len(df[i].index)-1 #1 row added by default so that table has a entry
    if y<0:
        y=0
    u.append(y)
    i+=1

print(u)
path= os.path.expanduser("~/Desktop/ece671/tcpnone")
num_files = len([f for f in os.listdir(path)if os.path.isfile(os.path.join(path, f))])
print(num_files)
i=0
j=0
t=[]
while i<(num_files/2):
    j=i+1

    path ="/home/vetri/Desktop/ece671/tcpnone/ftotal."+str(j)+".csv"
    y = file_len(path)
    #    except:       pass
    #df.append(pd.read_csv(path,header=None))
#       a+=[]
    #y=len(df[i].index)-1 #1 row added by default so that table has a entry
    if y<0:
        y=0
    t.append(y)
    i+=1

print(t)
path= os.path.expanduser("~/Desktop/ece671/tcpt8")
num_files = len([f for f in os.listdir(path)if os.path.isfile(os.path.join(path, f))])
print(num_files)
i=0
j=0
t8=[]
while i<(num_files/2):
    j=i+1

    path ="/home/vetri/Desktop/ece671/tcpt8/ftotal."+str(j)+".csv"
    y = file_len(path)
    #    except:       pass
    #df.append(pd.read_csv(path,header=None))
#       a+=[]
    #y=len(df[i].index)-1 #1 row added by default so that table has a entry
    if y<0:
        y=0
    t8.append(y)
    i+=1

print(t8)
#plt.figure(figsize=(4, 5))
plt.plot(list(range(1,len(u8)+1)),u8, '.-',label="udpt8")
plt.plot(list(range(1,len(u)+1)),u, '.-',label="udpnone")
plt.plot(list(range(1,len(t)+1)),t, '.-',label="tcpnone")
plt.plot(list(range(1,len(t8)+1)),t8, '.-',label="tcpt8")
plt.title("Total Flows Present after 1st flow")
plt.xlabel("time(s)")
plt.ylabel("flows")
#plt.frameon=True
plt.legend()
plt.show()


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tools to allow different io formats."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.contrib.learn.python.learn.learn_io.dask_io import extract_dask_data
from tensorflow.contrib.learn.python.learn.learn_io.dask_io import extract_dask_labels
from tensorflow.contrib.learn.python.learn.learn_io.dask_io import HAS_DASK
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import queue_parsed_features
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import read_batch_examples
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import read_batch_features
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import read_batch_record_features
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import read_keyed_batch_examples
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import read_keyed_batch_examples_shared_queue
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import read_keyed_batch_features
from tensorflow.contrib.learn.python.learn.learn_io.graph_io import read_keyed_batch_features_shared_queue
from tensorflow.contrib.learn.python.learn.learn_io.numpy_io import numpy_input_fn
from tensorflow.contrib.learn.python.learn.learn_io.pandas_io import extract_pandas_data
from tensorflow.contrib.learn.python.learn.learn_io.pandas_io import extract_pandas_labels
from tensorflow.contrib.learn.python.learn.learn_io.pandas_io import extract_pandas_matrix
from tensorflow.contrib.learn.python.learn.learn_io.pandas_io import HAS_PANDAS
from tensorflow.contrib.learn.python.learn.learn_io.pandas_io import pandas_input_fn
from tensorflow.contrib.learn.python.learn.learn_io.generator_io import generator_input_fn


"""
Benchmark SGD prediction time with dense/sparse coefficients.

Invoke with
-----------

$ kernprof.py -l sparsity_benchmark.py
$ python -m line_profiler sparsity_benchmark.py.lprof

Typical output
--------------

input data sparsity: 0.050000
true coef sparsity: 0.000100
test data sparsity: 0.027400
model sparsity: 0.000024
r^2 on test data (dense model) : 0.233651
r^2 on test data (sparse model) : 0.233651
Wrote profile results to sparsity_benchmark.py.lprof
Timer unit: 1e-06 s

File: sparsity_benchmark.py
Function: benchmark_dense_predict at line 51
Total time: 0.532979 s

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    51                                           @profile
    52                                           def benchmark_dense_predict():
    53       301          640      2.1      0.1      for _ in range(300):
    54       300       532339   1774.5     99.9          clf.predict(X_test)

File: sparsity_benchmark.py
Function: benchmark_sparse_predict at line 56
Total time: 0.39274 s

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    56                                           @profile
    57                                           def benchmark_sparse_predict():
    58         1        10854  10854.0      2.8      X_test_sparse = csr_matrix(X_test)
    59       301          477      1.6      0.1      for _ in range(300):
    60       300       381409   1271.4     97.1          clf.predict(X_test_sparse)
"""

from scipy.sparse.csr import csr_matrix
import numpy as np
from sklearn.linear_model.stochastic_gradient import SGDRegressor
from sklearn.metrics import r2_score

np.random.seed(42)

def sparsity_ratio(X):
    return np.count_nonzero(X) / float(n_samples * n_features)

n_samples, n_features = 5000, 300
X = np.random.randn(n_samples, n_features)
inds = np.arange(n_samples)
np.random.shuffle(inds)
X[inds[int(n_features / 1.2):]] = 0  # sparsify input
print("input data sparsity: %f" % sparsity_ratio(X))
coef = 3 * np.random.randn(n_features)
inds = np.arange(n_features)
np.random.shuffle(inds)
coef[inds[n_features/2:]] = 0  # sparsify coef
print("true coef sparsity: %f" % sparsity_ratio(coef))
y = np.dot(X, coef)

# add noise
y += 0.01 * np.random.normal((n_samples,))

# Split data in train set and test set
n_samples = X.shape[0]
X_train, y_train = X[:n_samples / 2], y[:n_samples / 2]
X_test, y_test = X[n_samples / 2:], y[n_samples / 2:]
print("test data sparsity: %f" % sparsity_ratio(X_test))

###############################################################################
clf = SGDRegressor(penalty='l1', alpha=.2, fit_intercept=True, n_iter=2000)
clf.fit(X_train, y_train)
print("model sparsity: %f" % sparsity_ratio(clf.coef_))

def benchmark_dense_predict():
    for _ in range(300):
        clf.predict(X_test)

def benchmark_sparse_predict():
    X_test_sparse = csr_matrix(X_test)
    for _ in range(300):
        clf.predict(X_test_sparse)

def score(y_test, y_pred, case):
    r2 = r2_score(y_test, y_pred)
    print("r^2 on test data (%s) : %f" % (case, r2))

score(y_test, clf.predict(X_test), 'dense model')
benchmark_dense_predict()
clf.sparsify()
score(y_test, clf.predict(X_test), 'sparse model')
benchmark_sparse_predict()


from __future__ import print_function, division
import matplotlib
import logging
from sys import stdout
matplotlib.use('Agg') # Must be before importing matplotlib.pyplot or pylab!
from neuralnilm import (Net, RealApplianceSource,
                        BLSTMLayer, DimshuffleLayer,
                        BidirectionalRecurrentLayer)
from neuralnilm.source import standardise, discretize, fdiff, power_and_fdiff
from neuralnilm.experiment import run_experiment, init_experiment
from neuralnilm.net import TrainingError
from neuralnilm.layers import MixtureDensityLayer
from neuralnilm.objectives import scaled_cost, mdn_nll, scaled_cost_ignore_inactive, ignore_inactive
from neuralnilm.plot import MDNPlotter

from lasagne.nonlinearities import sigmoid, rectify, tanh
from lasagne.objectives import mse
from lasagne.init import Uniform, Normal
from lasagne.layers import (LSTMLayer, DenseLayer, Conv1DLayer,
                            ReshapeLayer, FeaturePoolLayer, RecurrentLayer)
from lasagne.updates import nesterov_momentum, momentum
from functools import partial
import os
import __main__
from copy import deepcopy
from math import sqrt
import numpy as np
import theano.tensor as T

NAME = os.path.splitext(os.path.split(__main__.__file__)[1])[0]
PATH = "/homes/dk3810/workspace/python/neuralnilm/figures"
SAVE_PLOT_INTERVAL = 250
GRADIENT_STEPS = 100
SEQ_LENGTH = 512

source_dict = dict(
    filename='/data/dk3810/ukdale.h5',
    appliances=[
        ['fridge freezer', 'fridge', 'freezer'],
        'hair straighteners',
        'television'
        # 'dish washer',
        # ['washer dryer', 'washing machine']
    ],
    max_appliance_powers=[300, 500, 200, 2500, 2400],
    on_power_thresholds=[5] * 5,
    max_input_power=5900,
    min_on_durations=[60, 60, 60, 1800, 1800],
    min_off_durations=[12, 12, 12, 1800, 600],
    window=("2013-06-01", "2014-07-01"),
    seq_length=SEQ_LENGTH,
    output_one_appliance=False,
    boolean_targets=False,
    train_buildings=[1],
    validation_buildings=[1],
    skip_probability=0.0,
    n_seq_per_batch=16,
    subsample_target=4,
    include_diff=False,
    clip_appliance_power=True,
    target_is_prediction=False,
    independently_center_inputs = True,
    standardise_input=True,
    standardise_targets=True,
    input_padding=0,
    lag=0,
    reshape_target_to_2D=False,
    input_stats={'mean': np.array([ 0.05526326], dtype=np.float32),
                 'std': np.array([ 0.12636775], dtype=np.float32)},
    target_stats={
        'mean': np.array([ 0.04066789,  0.01881946,
                           0.24639061,  0.17608672,  0.10273963],
                         dtype=np.float32),
        'std': np.array([ 0.11449792,  0.07338708,
                       0.26608968,  0.33463112,  0.21250485],
                     dtype=np.float32)}
)

N = 50
net_dict = dict(
    save_plot_interval=SAVE_PLOT_INTERVAL,
#    loss_function=partial(ignore_inactive, loss_func=mdn_nll, seq_length=SEQ_LENGTH),
#    loss_function=lambda x, t: mdn_nll(x, t).mean(),
    loss_function=lambda x, t: mse(x, t).mean(),
#    loss_function=partial(scaled_cost, loss_func=mse),
    updates_func=momentum,
    learning_rate=1e-02,
    learning_rate_changes_by_iteration={
         500: 5e-03
       #  4000: 1e-03,
       # 6000: 5e-06,
       # 7000: 1e-06
       # 2000: 5e-06
        # 3000: 1e-05
        # 7000: 5e-06,
        # 10000: 1e-06,
        # 15000: 5e-07,
        # 50000: 1e-07
    },
    do_save_activations=True
)

def callback(net, epoch):
    net.source.reshape_target_to_2D = True
    net.plotter = MDNPlotter(net)
    net.generate_validation_data_and_set_shapes()
    net.loss_function = lambda x, t: mdn_nll(x, t).mean()
    net.learning_rate.set_value(1e-05)

def exp_a(name):
    # 3 appliances
    global source
    source_dict_copy = deepcopy(source_dict)
    source_dict_copy['reshape_target_to_2D'] = False
    source = RealApplianceSource(**source_dict_copy)
    source.reshape_target_to_2D = False
    net_dict_copy = deepcopy(net_dict)
    net_dict_copy.update(dict(
        experiment_name=name,
        source=source
    ))
    N = 50
    net_dict_copy['layers_config'] = [
        {
            'type': BidirectionalRecurrentLayer,
            'num_units': N,
            'gradient_steps': GRADIENT_STEPS,
            'W_in_to_hid': Normal(std=1.),
            'nonlinearity': tanh
        },
        {
            'type': FeaturePoolLayer,
            'ds': 4, # number of feature maps to be pooled together
            'axis': 1, # pool over the time axis
            'pool_function': T.max
        },
        {
            'type': BidirectionalRecurrentLayer,
            'num_units': N,
            'gradient_steps': GRADIENT_STEPS,
            'W_in_to_hid': Normal(std=1/sqrt(N)),
            'nonlinearity': tanh
        },
        {
            'type': DenseLayer,
            'W': Normal(std=1/sqrt(N)),
            'num_units': source.n_outputs,
            'nonlinearity': None
        }
    ]
    net_dict_copy['layer_changes'] = {
        1001: {
            'remove_from': -2,
            'callback': callback,
            'new_layers': [
                {
                    'type': MixtureDensityLayer,
                    'num_units': source.n_outputs,
                    'num_components': 2
                }
            ]
        }
    }

    net = Net(**net_dict_copy)
    return net

def main():
    #     EXPERIMENTS = list('abcdefghijklmnopqrstuvwxyz')
    EXPERIMENTS = list('a')
    for experiment in EXPERIMENTS:
        full_exp_name = NAME + experiment
        func_call = init_experiment(PATH, experiment, full_exp_name)
        logger = logging.getLogger(full_exp_name)
        try:
            net = eval(func_call)
            run_experiment(net, epochs=None)
        except KeyboardInterrupt:
            logger.info("KeyboardInterrupt")
            break
        except Exception as exception:
            logger.exception("Exception")
            raise
        finally:
            logging.shutdown()

if __name__ == "__main__":
    main()


"""
=================================================
Generate a functional label from source estimates
=================================================

Threshold source estimates and produce a functional label. The label
is typically the region of interest that contains high values.
Here we compare the average time course in the anatomical label obtained
by FreeSurfer segmentation and the average time course from the
functional label. As expected the time course in the functional
label yields higher values.

"""
# Author: Luke Bloy <luke.bloy@gmail.com>
#         Alex Gramfort <alexandre.gramfort@telecom-paristech.fr>
# License: BSD (3-clause)

import numpy as np
import matplotlib.pyplot as plt

import mne
from mne.minimum_norm import read_inverse_operator, apply_inverse
from mne.datasets import sample

print(__doc__)

data_path = sample.data_path()
subjects_dir = data_path + '/subjects'
fname_inv = data_path + '/MEG/sample/sample_audvis-meg-oct-6-meg-inv.fif'
fname_evoked = data_path + '/MEG/sample/sample_audvis-ave.fif'
subjects_dir = data_path + '/subjects'
subject = 'sample'

snr = 3.0
lambda2 = 1.0 / snr ** 2
method = "dSPM"  # use dSPM method (could also be MNE or sLORETA)

# Compute a label/ROI based on the peak power between 80 and 120 ms.
# The label bankssts-lh is used for the comparison.
aparc_label_name = 'bankssts-lh'
tmin, tmax = 0.080, 0.120

# Load data
evoked = mne.read_evokeds(fname_evoked, condition=0, baseline=(None, 0))
inverse_operator = read_inverse_operator(fname_inv)
src = inverse_operator['src']  # get the source space

# Compute inverse solution
stc = apply_inverse(evoked, inverse_operator, lambda2, method,
                    pick_ori='normal')

# Make an STC in the time interval of interest and take the mean
stc_mean = stc.copy().crop(tmin, tmax).mean()

# use the stc_mean to generate a functional label
# region growing is halted at 60% of the peak value within the
# anatomical label / ROI specified by aparc_label_name
label = mne.read_labels_from_annot(subject, parc='aparc',
                                   subjects_dir=subjects_dir,
                                   regexp=aparc_label_name)[0]
stc_mean_label = stc_mean.in_label(label)
data = np.abs(stc_mean_label.data)
stc_mean_label.data[data < 0.6 * np.max(data)] = 0.

func_labels, _ = mne.stc_to_label(stc_mean_label, src=src, smooth=True,
                                  subjects_dir=subjects_dir, connected=True)

# take first as func_labels are ordered based on maximum values in stc
func_label = func_labels[0]

# load the anatomical ROI for comparison
anat_label = mne.read_labels_from_annot(subject, parc='aparc',
                                        subjects_dir=subjects_dir,
                                        regexp=aparc_label_name)[0]

# extract the anatomical time course for each label
stc_anat_label = stc.in_label(anat_label)
pca_anat = stc.extract_label_time_course(anat_label, src, mode='pca_flip')[0]

stc_func_label = stc.in_label(func_label)
pca_func = stc.extract_label_time_course(func_label, src, mode='pca_flip')[0]

# flip the pca so that the max power between tmin and tmax is positive
pca_anat *= np.sign(pca_anat[np.argmax(np.abs(pca_anat))])
pca_func *= np.sign(pca_func[np.argmax(np.abs(pca_anat))])

###############################################################################
# plot the time courses....
plt.figure()
plt.plot(1e3 * stc_anat_label.times, pca_anat, 'k',
         label='Anatomical %s' % aparc_label_name)
plt.plot(1e3 * stc_func_label.times, pca_func, 'b',
         label='Functional %s' % aparc_label_name)
plt.legend()
plt.show()

###############################################################################
# plot brain in 3D with PySurfer if available
brain = stc_mean.plot(hemi='lh', subjects_dir=subjects_dir)
brain.show_view('lateral')

# show both labels
brain.add_label(anat_label, borders=True, color='k')
brain.add_label(func_label, borders=True, color='b')