import numpy as np
import configparser  as cnfgp
import re
import scipy.special    as spsp
import scipy.stats      as spst
import scipy.optimize   as spo

# This is a library with different functions needed to "fit" the data using Multiple Channel Markov Chain Monte-Carlo (MCMCMC) fitting



# this is f(x) = 1 + erf(x), a bit more stable implementation
OnePlusErf = lambda x: 2.*spst.norm.cdf(x, scale=1.0/np.sqrt(2.))

# here we will define all the basic functions for fitting
# all of them will require the t0, tcc, and tr parameters, but not all of them will process them

# this is the 'const' type function
def ConstFunc(x, t0=None, tcc=None, tr=None):
    try:
        return np.full(x.shape, fill_value=1.0)
    except:
        return 1.0

# this is the 'gauss' type function
def GaussFunc(x, t0, tcc, tr=None):
    return np.exp( -((x - t0)/tcc)**2 )

# this is the 'switch' type function with 'ProbeIsPump' being False
def SwitchDirectFunc(x, t0, tcc, tr=None):
    return OnePlusErf(+(x-t0)/tcc) 

# this is the 'switch' type function with 'ProbeIsPump' being True
def SwitchBackwardFunc(x, t0, tcc, tr=None):
    return OnePlusErf(-(x-t0)/tcc) 

# this is the 'trans'(transient) type function with 'ProbeIsPump' being False
def TransDirectFunc(x, t0, tcc, tr):
    if tr > 0.02 * tcc:
        return np.exp(-(x-t0)/tr) * OnePlusErf((+(x-t0)/tcc) - (0.5*tcc/tr)) * np.exp(0.25 * (tcc/tr)**2)
    else:
        return GaussFunc(x, t0, tcc, tr=None)

# this is the 'trans'(transient) type function with 'ProbeIsPump' being True
def TransBackwardFunc(x, t0, tcc, tr):
    if tr > 0.02 * tcc:
        return np.exp(+(x-t0)/tr) * OnePlusErf((-(x-t0)/tcc) - (0.5*tcc/tr)) * np.exp(0.25 * (tcc/tr)**2)
    else:
        return GaussFunc(x, t0, tcc, tr=None)



# this function converts the channel type and the 'ProbeIsPump' to a certain function
# Type (str) = 'const'/'gauss'/'switch'/'trans'
# ProbeIsPump is bool
def GetSomeFunc(Type, ProbeIsPump):
    if   Type == 'const' :  # constant 
        return ConstFunc
    elif Type == 'gauss' :  # Gaussian
        return GaussFunc
    elif Type == 'switch':  # Switch Behavior
        if ProbeIsPump:
            return SwitchBackwardFunc
        else:
            return SwitchDirectFunc
    elif Type == 'trans' :  # Transient peak/dip
        if ProbeIsPump:
            return TransBackwardFunc
        else:
            return TransDirectFunc
    else:
        raise RuntimeError("Unknown channel type '%s' requested" % (Type) )


    
    
    
    
# This is a function to read the variables definitions
# Requires: FileName (str) -- name of the config file contatining the information on the variables
# Returns: dictionary with keys being the names of the variables, and the containts are:
# "minvalue"/"maxvalue" -- are the lower/upper boundaries of the values for this variable
# "inivalue" -- is the starting value for this parameter, it is replaced by the "value"
# "maxdisp" -- is the maximal displacement during the generation of the new configuration, if not given, set to None, which will trigger the sampling from the uniform distribution
# if this parameter is not given,  
def readVariablesDefinitions(FileName):
    config = cnfgp.ConfigParser()  # initialize configparser object
    #config.optionxform = str      # this is to make everything case-sensitive
    config.read(FileName)          # 

    AllTheVars = dict()  # this dictionary of dictionaries is gonna be returned

    AcceptableKeys = ("IniValue".lower(), "MinValue".lower(), "MaxValue".lower(), "MaxDisp".lower() )
    print("Acceptable keys for variables are "+str(AcceptableKeys))
    for name in config.sections(): # get all the variables.
        Error = RuntimeError("Variable %s misdefined" % (name))

        VarVals = dict() # this is the dictionary with the variable settings
        for key in config[name]: # scan through the keys of the 
            #print(key)
            if key in AcceptableKeys:
                Val = float(config[name][key])
                if not key == "inivalue":
                    VarVals.update({key: Val})
                else:
                    VarVals.update({"value": Val})
            else:
                 print("Parameter %s for the variable %s was ignored" % (key,name))

        #print(VarVals)
        # check input parameters:
        # 1. MinValue and MaxValue should be provided
        if (not ("MinValue".lower() in VarVals.keys())) or (not ("MaxValue".lower() in VarVals.keys())):
            print("Min/Max values should be both provided for variable %s" % (name) )
            raise Error

        # check if MinValue <= IniValue <= MaxValue    
        if ("value" in VarVals.keys()) and ( (VarVals["value"] < VarVals["minvalue"]) or (VarVals["value"] > VarVals["maxvalue"])):
            print("Initial value is not in between MinValue and MaxValue")
            raise Error

        # if no inivalue is given, generate it from the uniform distribution
        if not "value" in VarVals.keys():
            VarVals.update({"value": np.random.uniform(low=VarVals["minvalue"], high=VarVals["maxvalue"], size=None)})

        if not "maxdisp" in VarVals.keys():
            VarVals.update({"maxdisp": None})

        AllTheVars.update({name.lower(): VarVals})

    return AllTheVars



# This is a function to read the variables definitions
# Requires: FileName (str) -- name of the config file contatining the information on the variables
# Returns: dictionary with keys being the names of the channels, and the containts are:
#          1. type = 'const'/'gauss'/'switch'/'trans' (required)
#          2. t0   (the t0 value), not needed for the 'const'
#          3. tcc  (the cross-correlation time), not needed for the 'const'
#          4. tr   (the relaxation=decay time), not needed for 'const' and 'gauss'
#          5. probeispump (if not given assumed to be False)
#          if t0/tcc/tr are not given, the fill value is None
def readChannelsDefinitions(FileName):
    config = cnfgp.ConfigParser()  # initialize configparser object
    #config.optionxform = str      # this is to make everything case-sensitive
    config.read(FileName)          # 

    AllTheChannels = dict()  # this dictionary of dictionaries is gonna be returned

    for name in config.sections(): # get all the channels. 
        Error = RuntimeError("Channel %s misdefined" % (name))
        ChanDef = dict() # this is the dictionary with the channel settings
        for key in config[name]: # scan through the keys of the channel
            # there are four possible types of functions:
            # 'const' -- constant = const
            # 'gauss' -- Gaussian cross-correlation = exp(-(t/tcc)**2)
            # 'switch'    -- switch between pump/probe and probe/pump (1 + erf(t/tcc))
            # 'trans'     -- transient increase/dip = exp((tcc/(2*tr))**2) * exp(-t/tr) * [ 1 + erf(t/tcc - tcc/(2*tr))]
            if key == "Type".lower():
                Val = str(config[name][key]).lower()
                ChanDef.update({key: Val})
            elif key == "T0".lower():
                Val = str(config[name][key]).lower()
                ChanDef.update({key: Val})
            elif key == "Tcc".lower():
                Val = str(config[name][key]).lower()
                ChanDef.update({key: Val})
            elif key == "Tr".lower():
                Val = str(config[name][key]).lower()
                ChanDef.update({key: Val})
            #elif key == "PumpIsPump".lower():
            #    Val = sconfig[name].getboolean(key)
            #    ChanDef.update({key: Val})
            elif key == "ProbeIsPump".lower():
                Val = config[name].getboolean(key)
                ChanDef.update({key: Val})
            else:
                 print("Parameter %s for the channel %s was ignored" % (key,name))

        # if the 'type' undefined, then it's an error
        if not "Type".lower() in ChanDef.keys():
            print("'type' parameter is obligatory for the channels. Not specified for %s" % name)
            raise Error
        
        # if this parameter is not given, by default it is considered to be False
        if not "ProbeIsPump".lower() in ChanDef.keys():
            ChanDef.update({"ProbeIsPump".lower() : False})

        # if this parameter is not given, by default it is considered to None 
        if not "Tr".lower() in ChanDef.keys():
            ChanDef.update({"Tr".lower() : None})
        
        # if this parameter is not given, by default it is considered to None 
        if not "Tcc".lower() in ChanDef.keys():
            ChanDef.update({"Tcc".lower() : None})

        # if this parameter is not given, by default it is considered to None 
        if not "T0".lower() in ChanDef.keys():
            ChanDef.update({"T0".lower() : None})
        
        AllTheChannels.update({name.lower(): ChanDef})    
    
    return AllTheChannels



# This is a function to read the variables definitions
# Requires: FileName (str) -- name of the config file contatining the information on the variables
# Returns: dictionary with keys being the names of the channels, and the containts are:
def readDataSetsDefinitions(FileName):
    config = cnfgp.ConfigParser()  # initialize configparser object
    #config.optionxform = str      # this is to make everything case-sensitive
    config.read(FileName)          # 

    AllTheDataSets = dict()  # this dictionary of dictionaries is gonna be returned

    for name in config.sections(): # get all the channels. 
        Error = RuntimeError("Data set %s misdefined" % (name))
        DataDef = dict() # this is the dictionary with the channel settings        
        for key in config[name]: # scan through the keys of the channel
            if key == "File".lower():
                Val = str(config[name][key])
                DataDef.update({key: Val})

                Data = np.loadtxt(Val)
                # load the data, and store that in the dictionary
                # We need at least three columns:
                #   1. x-values, 2. y-values, 3. errors of the y-values
                #   if the file has more than 3 columns, it is treated as x-y-xerrors-yerrors-... (i.e. 1st,2nd,and 4th columns are taken)
                if Data.shape[1] < 3:
                    print("Data set %s (file %s) has less than three columns (x+y+error)" % (name,Val))
                    raise Error
                elif Data.shape[1] == 3:
                    x = Data[:,0]
                    y = Data[:,1]
                    yerr = Data[:,2]
                else:
                    x = Data[:,0]
                    y = Data[:,1]
                    yerr = Data[:,3]

                DataDef.update({"x": x})
                DataDef.update({"y": y})
                DataDef.update({"yerr": yerr})

            elif key == "Channels".lower():
                Val = [ a.lower() for a in  re.split(r'[;,\s]\s*', config[name][key])]
                
                DataDef.update({key: Val})
            else:
                 print("Parameter %s for the dataset %s was ignored" % (key,name))

        if (not 'file' in DataDef.keys()) or (not 'channels' in DataDef.keys()):
            print("Problems with data %s: both 'file' and 'channels' are obligatory components" % (name))
            raise Error
        
        AllTheDataSets.update({name.lower(): DataDef})    
    
    return AllTheDataSets


# filter out of all the channels only those that were requested
# DataSets is the dictionary produced by the readDataSetsDefinitions
# Channels is the dictionary produced by the readChannelsDefinitions
def GetOnlyTheChannels(DataSets, Channels):
    TheChosenChannels = set() 
    for dataName in DataSets.keys():
        for key in DataSets[dataName]['channels']:
            TheChosenChannels.add(key)
            if not key in Channels.keys():
                print("Channel %s was asked for in the %s data, but is not defined in " % (key,dataName) + str(Channels.keys()) )
                raise RuntimeError("Failed to choose only the relevant channels")

    return TheChosenChannels

# filter out of all the variables that were not requested
# ChosenChannels is the set produced by the function GetOnlyTheChannels
# Channels is the dictionary produced by the readChannelsDefinitions
# Variables is the dictionary produced by the readVariablesDefinitions
def GetOnlyTheVariables(ChosenChannels, Channels, Variables):
    TheChosenVariables = set()
    Error1 = RuntimeError("Failed to choose only the relevant variables")
    Error2 = RuntimeError("Relevant variable was not defined for a channel")
    for chan in ChosenChannels:
        print(chan, Channels[chan])
        t0  = Channels[chan]['t0']
        tcc = Channels[chan]['tcc']
        tr  = Channels[chan]['tr']

        # get the t0 parameter of the channel
        if not t0 is None:
            if t0 in Variables.keys():
                TheChosenVariables.add(t0)
            else:
                print("Variable %s was requested as t0 in channel %s, but was not defined" % (t0,chan) )
                raise Error1
        else:
            if not Channels[chan]['type'] == 'const':
                print("the channel %s of type %s should have t0 defined" % (chan,Channels[chan]['type']))
                raise Error2


        # get the tcc parameter of the channel
        if not tcc is None:
            if tcc in Variables.keys():
                TheChosenVariables.add(tcc)
            else:
                print("Variable %s was requested as tcc in channel %s, but was not defined" % (tcc,chan) )
                raise Error1
        else:
            if not Channels[chan]['type'] == 'const': 
                print("the channel %s of type %s should have tcc defined" % (chan,Channels[chan]['type']))
                raise Error2


        # get the tcc parameter of the channel
        if not tr is None:
            if tr in Variables.keys():
                TheChosenVariables.add(tr)
            else:
                print("Variable %s was requested as tr in channel %s, but was not defined" % (tr,chan) )
                raise Error1
        else:
            if Channels[chan]['type'] == 'trans': 
                print("the channel %s of type %s should have tr defined" % (chan,Channels[chan]['type']))
                raise Error2
    
    return TheChosenVariables    


# A function to get directly a map of the variables with an index (i.e. get them in an ordered list)
def GetVariableMap(DataSets, Channels, Variables):
    TheChosenChannels = GetOnlyTheChannels(DataSets, Channels)
    TheChosenVariables = GetOnlyTheVariables(TheChosenChannels, Channels, Variables)
    return list(TheChosenVariables) # use list, not np.array, to use the list.index() method later, when making the map, the actual data vectors will be np.arrays anyways





#### This class will be stroring the actual data, do the linear LSQ fits, and calculation of the LSQ value
class DataSetHandler():

    # just a wrapper for procedure of the parameters indexes array extensions
    # Arr2Extend is self.__t0ind/self.__tccind/self.__tccind
    # Param is the value of the corresponding channel parameter
    # Map is VariableMap
    def __GetTheIndex(self, Arr2Extend, Param, Map):
        if not Param is None:
            Arr2Extend.append(Map.index(Param))
        else:
            Arr2Extend.append(None)



    # Initiation of this handler requires:
    #   1. the actual data set (i.e. particular dict from the DataSets dictionary obtained with readDataSetsDefinitions function)
    #   2. Name of the data set (str, whatever)
    #   3. Channels definitions (dict produced by the readChannelsDefinitions function)
    #   4. Variable map obtained by the GetVariableMap function
    #   5./6. xmin/xmax -- are the minimal/maximal values of the theoretical plot of the function
    #   7. Npts is the number of theoretical points in the map
    def __init__(self, DataSet, DataSetName, Channels, VariableMap, xmin=None, xmax=None, Npts=100, xMinMaxPart=0.1):
        self.name = str(DataSetName)
        self.__x = DataSet['x']      # x-values
        self.__y = DataSet['y']      # y-values
        self.__w = 1.0/(DataSet['yerr']**2)   # weights of the data points


        if xmin is None:
            xmin = min(self.__x) - (max(self.__x) - min(self.__x)) * xMinMaxPart

        if xmax is None:
            xmax = max(self.__x) + (max(self.__x) - min(self.__x)) * xMinMaxPart



        # we will use lists for channels for simplicity, and since we do not expect more than 10 channels per data set (it is already probably unfittable)
        self.__chNames = []  # names of the channels (for printing)
        self.__func    = []  # functions defining the channel (stores actual functions)
        self.__t0ind   = []  # indexes of the t0 parameters for the channels in the VariableMap
        self.__tccind  = []  # indexes of the tcc parameters for the channels in the VariableMap
        self.__trind   = []  # indexes of the tr parameters for the channels in the VariableMap
        #for key in Channels.keys():
        for key in DataSet["channels"]:
            chan = Channels[key]

            self.__chNames.append(key)
            self.__func.append(GetSomeFunc(Type = chan['type'], ProbeIsPump  = chan['probeispump']))
            
            self.__GetTheIndex(Arr2Extend = self.__t0ind, Param=chan['t0'], Map=VariableMap)
            self.__GetTheIndex(Arr2Extend = self.__tccind, Param=chan['tcc'], Map=VariableMap)
            self.__GetTheIndex(Arr2Extend = self.__trind, Param=chan['tr'], Map=VariableMap)


        self.__NCh = len(self.__func) # number of channels for this data set


        self.__Xt = np.linspace(xmin, xmax, Npts)  # storage of the fine grid theoretical function values
        self.__Yt = np.zeros(self.__Xt.shape)      # storage of the total theoretical fit function
        self.__YtComponents = np.zeros( (self.__NCh,self.__Xt.shape[0]), dtype = float )  # storage of the components of the fitted function

        
        self.__F = np.zeros( (self.__NCh,self.__x.shape[0]), dtype = float ) # this is gonna be an array of the current values of the channels, shape is (number of channels, number of x-points)
        self.__A = np.zeros( self.__NCh, dtype = float ) # these are the amplitudes of the channels
        self.__LSQ = 0.0 # this is where the least-squares value is stored



    # 
    def __getValueFromVector(self, index, Vector):
        if index is None:
            return None
        else:
            return Vector[index]

    # auxillary function to get the value of the 
    def __getNonLinParamsForChannel(self, chanIndex, Vector):
        t0  = self.__getValueFromVector(index = self.__t0ind[chanIndex], Vector = Vector)
        tcc = self.__getValueFromVector(index = self.__tccind[chanIndex], Vector = Vector)
        tr  = self.__getValueFromVector(index = self.__trind[chanIndex], Vector = Vector)
        return t0,tcc,tr

    # this functions triggers solving of the linear LSQ procedure for the data sets with defined self.__F and self.__y
    def _updateAmplitudes(self):
        AMat = np.matmul(self.__F,np.matmul(np.diag(self.__w),self.__F.T))
        if AMat.shape != (self.__NCh,self.__NCh):
            raise RuntimeError("Flip the transposing of the F-matrix!")
        BVec = np.matmul(self.__F,np.matmul(np.diag(self.__w),self.__y))
        try:
            self.__A = np.linalg.solve(AMat, BVec)
        except:
            #raise RuntimeError("The linear least squares for the dataset "+self.name+" could not be solved: probably some channels are linearly dependent")
            print("The linear least squares for the dataset "+self.name+" could not be solved: probably some channels are linearly dependent")
            self.__A = np.zeros(BVec.shape) # if no LSQ solution was obtained, just set every channel cross-section to zero
            
            
    # this function updates the self.__F values of the channels with the new set of nonlinear parameters
    def _updateFunctions(self, VariablesVector):
        for n in range(0,self.__NCh):
            t0,tcc,tr = self.__getNonLinParamsForChannel(chanIndex = n, Vector = VariablesVector) 
            self.__F[n] = self.__func[n](x = self.__x, t0 = t0, tcc = tcc, tr = tr)

    # Update the theoretical plotting values for the new 
    def _updateTheorFunction(self, VariablesVector):
        self.__Yt = np.zeros(self.__Xt.shape)
        for n in range(0,self.__NCh):
            t0,tcc,tr = self.__getNonLinParamsForChannel(chanIndex = n, Vector = VariablesVector)
            self.__YtComponents[n] = self.__A[n] * self.__func[n](x = self.__Xt, t0 = t0, tcc = tcc, tr = tr)         
            self.__Yt += self.__YtComponents[n]

    


    # Update Functions and Amplitudes based onthe new nonlinear parameters
    def update(self, VariablesVector):
        self._updateFunctions(VariablesVector = VariablesVector)
        self._updateAmplitudes()
        self._updateTheorFunction(VariablesVector = VariablesVector)


    # return list of the Channel names
    def getChNames(self):
        return np.array(self.__chNames)
  
    # return the Channel amplitudes
    def getAmp(self):
        return np.array(self.__A)

    # return the total number of channels
    def getNumOfCh(self):
        return int(self.__NCh)
        
    # return the Channel names and amplitudes in the form of the dictionary
    def getAmpDict(self):
        res = dict()
        for n in range(0,self.__NCh):
            res.update({self.__chNames[n]: self.__A[n]})
        return res

    # get the experimental X-values
    def getXValues(self):
        return np.array(self.__x)
 
    # get the experimental Y-values
    def getYValues(self):
        return np.array(self.__y)


    # get the experimental Y-error values
    def getErrValues(self):
        return np.array(1./np.sqrt(self.__w))


    # get the theoretical signal, 
    # GiveExpPoints regulate whether a finer grid is used, or the experimental values are computed only
    def getCurrentYtheor(self, GiveExpPoints = False):
        if GiveExpPoints:
            #print(self.__F)
            #print(self.__A)
            return np.matmul(self.__F.T, self.__A)
        else:
            return np.array(self.__Yt)

    # get the theoretical signal x-axis
    def getCurrentXtheor(self, GiveExpPoints = False):
        if GiveExpPoints:
            return np.array(self.__x)
        else:
            return np.array(self.__Xt)

    # get the 
    def getCurrentYtheorComponent(self, n, GiveExpPoints = False):
        if GiveExpPoints:
            #print(self.__F)
            #print(self.__A)
            return self.__A[n]*self.__F[n]
        else:
            return np.array(self.__YtComponents[n])


    # function to return the LSQ values for the current set of amplitudes
    def getLSQ(self):
        Yth = self.getCurrentYtheor(GiveExpPoints = True)
        return 0.5 * np.sum(self.__w * (Yth - self.__y)**2)

            

#TODO: variable handler
class VariablesHandler():

    def __getIniValue(self, varName, Variables):
        val = Variables[varName]["value"]
        if not val is None:
            return val
        else:
            return np.random.uniform(low=Variables[varName]["minvalue"], high=Variables[varName]["maxvalue"])
        
    def __init__(self, DataSets, Channels, Variables):

        self.__VariableMap = GetVariableMap(DataSets, Channels, Variables)
        self.__NVars = len(self.__VariableMap)
                
        self.__VarVec = np.array([self.__getIniValue(a, Variables=Variables) for a in self.__VariableMap  ]) # vector with the current values of the variables

        self.__Bounds  = np.array([(Variables[a]["minvalue"], Variables[a]["maxvalue"]) for  a in self.__VariableMap   ]) # distribution bounds for the variables
        
        self.__Steps  = np.array([Variables[a]["maxdisp"] for a in self.__VariableMap])        # the steps of changing the parameters, None means drawing the step from the uniform distribution
    
    def getMap(self):
        return list(self.__VariableMap)
  
    def getVarVec(self):
        return np.array(self.__VarVec)

    def getBounds(self):
        return tuple(self.__Bounds)
 

    # function to reset the values of the variables vector
    def updateVarVec(self, VarVec):
        if len(VarVec) == len(self.__VarVec):
            ItIsGood = True
            for n,bounds in enumerate(self.__Bounds): # check if the 
                if (bounds[0]>VarVec[n]) or (bounds[1]<VarVec[n]):
                    ItIsGood = False
            if not ItIsGood:
                raise RuntimeError("the newly assigned vector does not match with the pre-defined boundaries of the variables")

            self.__VarVec = np.array(VarVec) 
        else:
            raise RuntimeError("the new assigned variables vector does not have the same shape as before")

    # method to generate a trial value for a single variable
    def __genSingleVar(self, n):
        dxmax = self.__Steps[n]
        
        if dxmax is None:
            res = np.random.uniform(low=self.__Bounds[n][0], high=self.__Bounds[n][1])
        else:
            dx = np.random.uniform(low=-dxmax, high=dxmax)
            res = self.__VarVec[n] + dx
            # the new value is manually set to be within the defined boundaries
            res = max(self.__Bounds[n][0], res)
            res = min(self.__Bounds[n][1], res)
        return res
            
    # generate the value of a trial vector from the chosen distribution
    def genTrialVarVec(self):
        return np.array([self.__genSingleVar(n)   for n in range(0,self.__NVars)])
            
            
    def getDictOfVars(self, x = None):
        res = dict()
        if not x is None:
            for n in range(0,self.__NVars):
                res.update({self.__VariableMap[n]: x[n]})
        else:
            for n in range(0,self.__NVars):
                res.update({self.__VariableMap[n]: self.__VarVec[n]})            
        return res        


# class to store the regularization factors
class AdditionalConstraints():
    def __readInpFile(self, FileName, VariableMap):
        InpFile = open(FileName, "r")
        for line in InpFile:
            word = line.split()
            if len(word) >= 4: 
                # format for the repellent correction (i.e. addition of Alpha_{ij}/|x_i - x_j| to LSQ) is the following:
                # " repellent   parameter1name  parameter2name   alphaValue(float >= 0.)  "
                if word[0].lower() == 'repellent':
                    par1n = word[1].lower()
                    par2n = word[2].lower()
                    if (par1n in VariableMap) and (par2n in VariableMap):
                        ind1 = VariableMap.index(par1n)
                        ind2 = VariableMap.index(par2n)
                        alpha = float(word[3])
                        self.__RInd.append( (ind1,ind2) )
                        self.__Alpha.append(alpha)
                        

                # format for the known value (i.e. addition of 0.5*(x-x0)**2/errorX**2 to LSQ) is the following:
                # " constraint   parameter  x0(float)    errorX(float>0.)  "
                if word[0].lower() == 'constraint':
                    parn = word[1].lower()
                    print("adding constraint for the variable "+parn)
                    if (parn in VariableMap):  
                        ind = VariableMap.index(parn)           
                        self.__CInd.append(ind)
                        x0  = float(word[2])
                        err = float(word[3])

                        self.__X0.append(x0)
                        self.__WX.append(1./err**2)
        InpFile.close()

        self.__RInd  = np.array(self.__RInd)
        self.__Alpha = np.array(self.__Alpha)

        self.__CInd  = np.array(self.__CInd)
        self.__X0    = np.array(self.__X0)
        self.__WX    = np.array(self.__WX)



    # requires: FileName with the constrains
    #           VariableMap (same as before)
    def __init__(self, FileName, VariableMap):
        self.__CInd = []  # here the constraint parameters indexes in the VariableMap will be stored
        self.__X0     = []  # here the known X0 values for the constraining parameters are stored
        self.__WX     = []  # here the weights of the W are given (1/errorX**2)

        self.__RInd  = []  # here the pairs of indexes of the repelled parameters will be stored
        self.__Alpha  = []  # here the pairs of indexes of the repelled parameters will be stored        

        self.__readInpFile(FileName = FileName, VariableMap = VariableMap)

        self.__NC = len(self.__X0)
        self.__NR = len(self.__Alpha) 

    def _totalConstr(self, VarVec):
        if self.__NC>0:
            X = VarVec[self.__CInd]
            return 0.5 * np.dot(self.__WX,  (X - self.__X0)**2)
        else:
            return 0.0

    def _totalRepell(self, VarVec):
        if self.__NR>0:
            X1 = VarVec[self.__RInd[:,0]]
            X2 = VarVec[self.__RInd[:,1]]
            return np.dot(self.__Alpha,  1.0/np.abs(X1 - X2))
        else:
            return 0.0

    def getRegValue(self, VarVec):
        return self._totalConstr(VarVec) + self._totalRepell(VarVec)
    
    
    
    
## Function to perform the global fit of the data based on the channels and constraints present
## ListOfData -- is a list of DataSetHandler instances
## VarHand -- is a variables handler (VariablesHandler instance)
## AddConstr is the handler of the additional constraints (AdditionalConstraints instance)
## Method -- is the choice of three global optimization methods from SciPy.optimize 
##           ('differential_evolution' = 'DiffEvo', 'shgo' = 'SHGO', and 'dual_annealing' = 'DualAnn' ). UPD: "SHGO" is depreciated
## AdditionalArguments -- is the dictionary of additional arguments of the Scipy.optimize for the possible global opt methods
def DoTheGlobalFit(ListOfData, VarHand, AddConstr=None, Method = 'DiffEvo', 
                   AdditionalArguments=dict(), ReturnDictionary=False,
                   AddInfoToFiles = None):
    if Method.lower() == 'diffevo':
        OptMethod = spo.differential_evolution
#    elif Method.lower() == 'shgo':
#        OptMethod = spo.shgo
    elif Method.lower() == 'dualann':
        OptMethod = spo.dual_annealing
    else:
        raise RuntimeError("Unknown global optimization routine "+str(Method)+" was requested")
    
    Bounds = VarHand.getBounds()
    
    # a cover function to get the Least-Squares value for the fitting
    def PureLSQ(x):
        val = 0.0
        for data in ListOfData:
            data.update(VariablesVector = x)
            val += data.getLSQ()
        return val
    
    # regularization term    
    def LSQConstrVal(x):
        if not AddConstr is None:
            return AddConstr.getRegValue(VarVec = x)
        else:
            return 0.0    
    
    # total LSQ to be minimized
    def TotalLSQ(x):
        return PureLSQ(x) + LSQConstrVal(x)
    
    Res = OptMethod(func = TotalLSQ, bounds=Bounds, **AdditionalArguments) 

    # save how the theoretical fitted curves look like
    if AddInfoToFiles is None:
        AddInfoToFiles = ""
    else:
        AddInfoToFiles = "_"+str(AddInfoToFiles)   


    # save all the fitted pump-probe dependencies
    for data in ListOfData:
        OutFileName = data.name+"_gfit_"+Method.lower()+AddInfoToFiles+".res"
        OutFileNameResiduals = data.name+"_residuals_gfit_"+Method.lower()+AddInfoToFiles+".res"

        X = data.getCurrentXtheor(GiveExpPoints = False)
        Y = data.getCurrentYtheor(GiveExpPoints = False)
        ChNames = data.getChNames()
 
        # save the residuals as well
        XRes = data.getXValues() # experimental x-points
        YRes = data.getCurrentYtheor(GiveExpPoints = True) - data.getYValues() # the residual is the "Theory - Experiment"
        YErr = data.getErrValues() # errors in y-axis

        Arr2Stack = [X,Y]
        # also save how the particular channels look like
        FileHeader = "  PPDelay[$1]  YTheorTotal[$2] "
        for n,chn in enumerate(ChNames):
            FileHeader+=" %s[$%i] " % (chn,n+3)
            Arr2Stack = Arr2Stack + [data.getCurrentYtheorComponent(n=n,GiveExpPoints = False)]

        np.savetxt(OutFileName, np.stack(Arr2Stack, axis=-1), header=FileHeader)
        np.savetxt(OutFileNameResiduals, np.stack([XRes,YRes,YErr], axis=-1), header=" PPDelay[$1]  Residual[$2] StdExp[$3]") # save the residuals


    # Save the fitted variables into a file
    OutFileName = "variables_gfit_"+Method.lower()+AddInfoToFiles+".res" 
    outf = open(OutFileName, "w")
    Map = VarHand.getMap()

    for q,var in enumerate(Map):
        outf.write("%20s = %15.10f \n" % (var, Res.x[q]) )
    outf.close()

    if ReturnDictionary: 
        return VarHand.getDictOfVars(Res.x)    
    else:
        return Res.x
        
# A function to update the value of the average in the MCMC simulation
# the average is YAverage = (\sum_{i=1}^N y_i) /N
# YAverage is the average value for the step N-1
# ynew is the new point at step #N
# N is the step number        
def UpdateValue(YAverage, ynew, N):
    return  (float(N-1)/float(N)) * YAverage +  ynew/float(N)   
        
# A procedure to perform the MCMC sampling starting from the         
def DoTheMCMCSampling(ListOfData, VarHand, AddConstr=None, 
                      NSteps=10000,            # Number of steps in MCMC simulation
                      NStepsToIgnore=None,     # How many first steps to ignore (overrides PartToIgnore if provided)
                      PartToIgnore=0.1,        # Which part of the MCMC simulation to ignore (is overriden by NStepsToIgnore if not None)
                      AcceptanceProbabilityScale = 1.0, # Scaling coefficient for the Metropolis acceptence probability
                      AddInfoToFiles = None, # (str) something to be added to the files with the histograms, etc
                      NBinsInHist    = 50,   # number of bins in the histograms
                      PrintStatus = True     # bool flag to print the status of the MCMC sampling
                     ):
    if NStepsToIgnore is None:
        NStepsToIgnore = int(PartToIgnore * NSteps)
    
    if AddInfoToFiles is None:
        AddInfoToFiles = ""
    else:
        AddInfoToFiles = "_"+str(AddInfoToFiles)   
        
    #### THIS PART IS A COPY FROM THE DoTheGlobalFit ##############    
    # a cover function to get the Least-Squares value for the fitting
    def PureLSQ(x):
        val = 0.0
        for data in ListOfData:
            data.update(VariablesVector = x)
            val += data.getLSQ()
        return val
    
    # regularization term    
    def LSQConstrVal(x):
        if not AddConstr is None:
            return AddConstr.getRegValue(VarVec = x)
        else:
            return 0.0    
    
    # total LSQ to be minimized
    def TotalLSQ(x):
        return PureLSQ(x) + LSQConstrVal(x)        
    #### THIS PART IS A COPY FROM THE DoTheGlobalFit (END) ########    
    
    XVal   = VarHand.getVarVec()
    LSQVal = TotalLSQ(XVal)
    TRJFileName = "mcmc_sampling"+AddInfoToFiles+".trj"    
    outf=open(TRJFileName, "w")
    outf.write("#%9s" % ("step.num."))
    Map = VarHand.getMap()
    for point in Map:
        outf.write(" %10s" % (point))
    outf.write(" %8s %8s %8s %8s %s\n" % ("Pacc", "Ptrial", "LSQ", "Reg", "Accepted or Not"))
    
    
    Ntrj = 0 # this is the enumeration of the points in the sampled part of the trj
    
    # Here we will store the averaged theoretical fits
    Yavg = [np.zeros(x.getCurrentXtheor(GiveExpPoints = False).shape) for x in ListOfData]
    # Here we will store the standard deviations for the theoretical fits    
    Ystd = [np.zeros(x.getCurrentXtheor(GiveExpPoints = False).shape) for x in ListOfData]

    # Here we will strore the residuals
    YRes = [np.zeros(x.getCurrentXtheor(GiveExpPoints = True).shape) for x in ListOfData]

    
    # here we will store the pure channels average values and standard deviations
    YCompA = [np.zeros((x.getNumOfCh(), x.getCurrentXtheor(GiveExpPoints = False).shape[0])) for x in ListOfData]
    YCompS = [np.zeros((x.getNumOfCh(), x.getCurrentXtheor(GiveExpPoints = False).shape[0])) for x in ListOfData]


    # Here we will store the averaged theoretical values/standard deviations
    Xavg = np.zeros(XVal.shape)
    Xstd = np.zeros(XVal.shape)
    
    # Here we will store the correlation matrix
    XiXj = np.zeros( (len(Xavg),len(Xavg)) ) 
    
    # This is the acceptance rate
    AccRate = 0.0

    for n in range(0,NSteps):
        if n % 1000 == 0 and PrintStatus:
            print("%11i points sampled (current acceptance ratio is %7.4f %%)"  % (n,AccRate*100.0))

        trialX   = VarHand.genTrialVarVec()
        trialLSQ = TotalLSQ(trialX)
        
        Pacc = AcceptanceProbabilityScale * np.exp(LSQVal - trialLSQ)
        
        Ptrial = np.random.uniform(low=0.0, high=1.0, size=None)
        Acc = 0
        if Ptrial <= Pacc:
            XVal   = trialX
            LSQVal = trialLSQ
            VarHand.updateVarVec(VarVec = XVal)
            Acc = 1
        pLSQ = PureLSQ(XVal)
        pReg = LSQConstrVal(XVal)
        
        outf.write(" %9i" % n)
        for x in XVal:
            outf.write(" %10.7f" % (x))
        outf.write(" %8.4e %8.4e %8.4e %8.4e %1i\n" % (Pacc, Ptrial, pLSQ, pReg, Acc))
        outf.flush()
        
        
        if n > NStepsToIgnore:
            if n == NStepsToIgnore + 1 and PrintStatus:
                print("  **** started collecting the stats at the step %i of the MCMC simulation **** " % n)

            Ntrj += 1
            Xavg = UpdateValue(YAverage=Xavg, ynew=XVal,    N=Ntrj)
            Xstd = UpdateValue(YAverage=Xstd, ynew=XVal**2, N=Ntrj) 
            
            XiXj = UpdateValue(YAverage=XiXj, ynew=np.outer(XVal,XVal), N=Ntrj) 
            
            AccRate = UpdateValue(YAverage=AccRate, ynew = float(Acc), N=Ntrj)       
            for q,data in enumerate(ListOfData):
                    
                data.update(VariablesVector = XVal) # update the theoretical curves
                
                Yavg[q] = UpdateValue(YAverage=Yavg[q], ynew=data.getCurrentYtheor(GiveExpPoints = False),    N=Ntrj)
                Ystd[q] = UpdateValue(YAverage=Ystd[q], ynew=data.getCurrentYtheor(GiveExpPoints = False)**2, N=Ntrj)  
                # update the residuals
                YRes[q] = UpdateValue(YAverage=YRes[q], ynew=data.getCurrentYtheor(GiveExpPoints = True),    N=Ntrj)

                for chN in range(0,data.getNumOfCh()):
                    YCompA[q][chN] = UpdateValue(YAverage=YCompA[q][chN], ynew=data.getCurrentYtheorComponent(n=chN,GiveExpPoints = False),       N=Ntrj)   
                    YCompS[q][chN] = UpdateValue(YAverage=YCompS[q][chN], ynew=data.getCurrentYtheorComponent(n=chN,GiveExpPoints = False)**2,    N=Ntrj)    
    
    if PrintStatus:
        print("Sampling of %i points successfully done :). Acceptance rate was %7.4f %%" % (n,AccRate*100.0))

    outf.close()

    # save all the fitted pump-probe dependencies
    for q,data in enumerate(ListOfData):
        OutFileName = data.name+"_mcmc_sampling"+AddInfoToFiles+".res"
        OutFileNameResiduals = data.name+"_residuals_mcmc_sampling"+AddInfoToFiles+".res"
        X = data.getCurrentXtheor(GiveExpPoints = False)
        Y = Yavg[q]
        Z = np.sqrt(Ystd[q] - Y**2) # calculate standard deviations

        XRes = data.getXValues() # get the experimental x-points for the residuals
        YRes[q] = YRes[q] - data.getYValues() # get the residuals
        YErr = data.getErrValues() # errors in y-axis
        Arr2Stack = [X,Y,Z]
        FileHeader = "  PPDelay[$1]  AvYTheorTotal[$2] StdYTheorTotal[$3] "

        ChNames = data.getChNames()
        for n,chn in enumerate(ChNames):
            FileHeader+=" Av(%s)[$%i] Std(%s)[$%i] " % (chn,2*n+4,chn,2*n+5)
            YCompS[q][n] = np.sqrt(YCompS[q][n] - YCompA[q][n]**2)
            Arr2Stack = Arr2Stack + [YCompA[q][n],YCompS[q][n]]

        Arr2Stack = np.stack(Arr2Stack, axis=-1)
        np.savetxt(OutFileName, Arr2Stack, header=FileHeader)
        np.savetxt(OutFileNameResiduals, np.stack([XRes,YRes[q],YErr], axis=-1), header=" PPDelay[$1]  Residual[$2] StdExp[$3]")
        
    # Save fitted values statistical data
    
    OutFileName = "variables_mcmc_sampling"+AddInfoToFiles+".res"    
    outf = open(OutFileName, "w")
    Xstd = np.sqrt(Xstd - Xavg**2)
    XiXj -= np.outer(Xavg,Xavg) 
    XiXj /= np.outer(Xstd,Xstd) # convert covariance into the PCC (Pearson correlation coefficient)
    Map = VarHand.getMap()
    
    res2return = dict()

    # here is the header for the covariance matrix with the names of the variables
    CovHeader = ""

    for q,var in enumerate(Map):
        outf.write("%20s = %15.10f +/- %15.10f\n" % (var, Xavg[q], Xstd[q]) )
        CovHeader += " '%s'[$%i] " % (var,q+1) # we add the variable name and the column number in the square brackets
        res2return.update({var: (Xavg[q], Xstd[q])})
        # also save the distribution of the parameter
        Data = np.loadtxt(TRJFileName, usecols=(q+1))[NStepsToIgnore+1:] # 0th column is the step number
        H,Be = np.histogram(Data, bins = NBinsInHist)
        X = 0.5*(Be[1:] + Be[:-1])
        np.savetxt("variable_"+var+"_sampled_distribution"+AddInfoToFiles+".res", np.stack([X,H],axis=-1) )

    outf.close()
    # save the PCC (Pearson correlation coefficient)
    np.savetxt("pcc_mcmc_sampling.res", XiXj, header=CovHeader)
    
    return res2return



## Function to perform the global fit of the data based on the channels and constraints present
## ListOfData -- is a list of DataSetHandler instances
## VarHand -- is a variables handler (VariablesHandler instance)
## AddConstr is the handler of the additional constraints (AdditionalConstraints instance)
## Method (str) -- is the choice of the optimization algorithm in Scipy.optimize.minimize(method = Method)
##           ('differential_evolution' = 'DiffEvo', 'shgo' = 'SHGO', and 'dual_annealing' = 'DualAnn' ). UPD: "SHGO" is depreciated
## AdditionalArguments -- is the dictionary of additional arguments of the Scipy.optimize for the possible global opt methods
def DoTheLocalFit(ListOfData, VarHand, AddConstr=None, Method = 'Powell', 
                   AdditionalArguments=dict(), ReturnDictionary=False,
                   AddInfoToFiles = None):

    Bounds = VarHand.getBounds()   # these are the boundaries
    IniGuess = VarHand.getVarVec() # the initial guess is the current value stored in VarHand
    
    # a cover function to get the Least-Squares value for the fitting
    def PureLSQ(x):
        val = 0.0
        for data in ListOfData:
            data.update(VariablesVector = x)
            val += data.getLSQ()
        return val
    
    # regularization term    
    def LSQConstrVal(x):
        if not AddConstr is None:
            return AddConstr.getRegValue(VarVec = x)
        else:
            return 0.0    
    
    # total LSQ to be minimized
    def TotalLSQ(x):
        return PureLSQ(x) + LSQConstrVal(x)
    
    Res = spo.minimize(fun = TotalLSQ, x0=IniGuess, bounds=Bounds, method=Method, **AdditionalArguments) 

    # save how the theoretical fitted curves look like
    if AddInfoToFiles is None:
        AddInfoToFiles = ""
    else:
        AddInfoToFiles = "_"+str(AddInfoToFiles)   

    # save all the fitted pump-probe dependencies
    for data in ListOfData:
        MethodName = str(Method)
        OutFileName = data.name+"_lfit_"+MethodName.lower()+AddInfoToFiles+".res"
        OutFileNameResiduals = data.name+"_residuals_lfit_"+MethodName.lower()+AddInfoToFiles+".res"
        X = data.getCurrentXtheor(GiveExpPoints = False)
        Y = data.getCurrentYtheor(GiveExpPoints = False)

        # save the residuals as well
        XRes = data.getXValues() # experimental x-points
        YRes = data.getCurrentYtheor(GiveExpPoints = True) - data.getYValues() # the residual is the "Theory - Experiment"
        YErr = data.getErrValues() # errors in y-axis

        ChNames = data.getChNames()
 
        Arr2Stack = [X,Y]
        # also save how the particular channels look like
        FileHeader = "  PPDelay[$1]  YTheorTotal[$2] "
        for n,chn in enumerate(ChNames):
            FileHeader+=" %s[$%i] " % (chn,n+3)
            Arr2Stack = Arr2Stack + [data.getCurrentYtheorComponent(n=n,GiveExpPoints = False)]
        np.savetxt(OutFileName, np.stack(Arr2Stack, axis=-1), header=FileHeader)

        np.savetxt(OutFileNameResiduals, np.stack([XRes,YRes,YErr], axis=-1), header=" PPDelay[$1]  Residual[$2] StdExp[$3]") # save the residuals

    # Save the fitted variables into a file
    OutFileName = "variables_lfit_"+MethodName.lower()+AddInfoToFiles+".res"   
    outf = open(OutFileName, "w")
    Map = VarHand.getMap()

    for q,var in enumerate(Map):
        outf.write("%20s = %15.10f \n" % (var, Res.x[q]) )
    outf.close()



    if ReturnDictionary: 
        return VarHand.getDictOfVars(Res.x)    
    else:
        return Res.x

# Function to read the data from the files, and then return the 
# three things needed for fitting:
#  1) ListOfData, a list of DataSetHandler instances, storing the datas and the definition of the functions for the fitting
#  2) VarHand, a VariablesHandler class example, which controls the life of the fitted nonlinear variables
#  3) RegInstance (None or AdditionalConstraints class example) -- is the thing controlling the regularization
def ReadTheData(DataDescription,      # data description file name
                ChannelsDescription,  # channels description file name
                VariablesDescription, # variables description file name
                RegularizationDescription = None, # regularization description file name
                xmin=None, xmax=None, Npts=100, xMinMaxPart=0.1   # these are parameters of the DataSetHandler class example initiation
               ):

        # read the files first, producing the 
        Chans = readChannelsDefinitions(ChannelsDescription)
        Datas = readDataSetsDefinitions(DataDescription)
        Vars  = readVariablesDefinitions(VariablesDescription)
        
        # Store only the chosen variables for variation in a special class example
        VarHand = VariablesHandler(DataSets = Datas, Channels = Chans, Variables = Vars) 
        Map     = VarHand.getMap()  # this is the map, that maps the variable name to an index through a list, containing the variables names


        # read regularization parameters, only if requested
        RegInstance   = None
        if not RegularizationDescription is None:
            # Regularization stuff is being stored as a special class example
            RegInstance = AdditionalConstraints(FileName = RegularizationDescription, VariableMap = Map)   

        # now store the data sets + channels that are supposed fit them in a form of special handlers
        ListOfData = []
        # initiate the dataset handlers and store them in a list
        for key in Datas.keys():
            Example = DataSetHandler(DataSet=Datas[key], DataSetName=key, Channels=Chans, VariableMap=Map,  # these are the reqired parameters
                                            xmin=xmin, xmax=xmax, Npts=Npts, xMinMaxPart=xMinMaxPart) # these are optional parameters, that define how the function for plotting will be stored
            ListOfData.append(Example)

        return ListOfData, VarHand, RegInstance





