# MIT License
# 
# Copyright (c) 2021- Pierre-Yves Taunay 
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np

from sklearn.base import BaseEstimator
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
from sklearn.metrics import mean_squared_error

from lmfit import Model, Parameters

from sympy import lambdify
from sympy.parsing.sympy_parser import parse_expr

# Filter some warnings
import warnings                                                                                     
warnings.simplefilter("ignore", UserWarning)                                                        
warnings.simplefilter("ignore", RuntimeWarning) 


class TheoreticalModelEstimator(BaseEstimator):
    '''
    Models to try: 1/4 - log(PI2) + c0 * PI5 + ...
    ... c1 * (1/PI2**2-1) * PI2**c2 * PI3**c3 * PI4**c4 * PI5**c5 * PI6**c6
    
    The pi-products may/may not be included by setting a mask value to true or false. 
    icdf is a dataframe of initial conditions. The index icidx is used to grab a given initial 
    condition in icdf. cmin and cmax limit the range of constant c0 that appears in front of Pi5.
    '''    
    def __init__(self,icdf, m2=True, m3=True, m4=True, m5=True, m6=True, icidx=0,cmin=1,cmax=10):
        '''
        Estimator constructor. Parameters:
            - icdf: Dataframe of initial conditions
            - icidx: Initial condition to pick
            - cmin, cmax: Bounds of the constant that appears in front of PI5
            - m2,m3,m4,m5,m6: Should those PI products be included in the fit
        '''
        super().__init__()
        self.params_ = None
        self.mse_ = -1
        self.n_parameters = 0
        self.mask_ = []
        self.p0 = np.zeros(7)
        self.icdf = icdf
        
        self.m2 = m2
        self.m3 = m3
        self.m4 = m4
        self.m5 = m5
        self.m6 = m6
        self.icidx = icidx

        self.cmin = cmin
        self.cmax = cmax

    def create_function_string(self):
        fstr = ''
        lambcoeff = []
        lambvar = []
        
        # Start with 1/4 - log(PI2) + c0 * PI5
        fstr = 'log10(1/4 - log(x0) + a0*x3'
        
        lambcoeff = ['a0']
        lambvar = ['x0','x3']
        
        # Build parameters with the initial condition stored in p0
        self.params_ = Parameters()
        self.params_.add('a0', value=self.p0[0], min=self.cmin, max=self.cmax)
        
        self.n_parameters = 1
        if(any(self.mask_)):
            # If any other pi-product is included (i.e. a1 =/= 0), then add on the corresponding term
            fstr += '+ a1 * (x0**(-2)-1)'
            
            # One more parameter: a1 
            self.n_parameters += 1

            lambcoeff.append('a1')
            self.params_.add('a1', value=self.p0[1], min=0)
            
            # Add the pi-products to their respective power to be found
            for kk,v in enumerate(self.mask_):
                if v:
                    aname = 'a' + str(kk+2)
                    mname = 'm' + str(kk+2)
                    xname = 'x' + str(kk)
                    
                    fstr += '*(' + xname + '/' + mname + ')**' + aname
                    self.n_parameters += 1
                    
                    
                    lambcoeff.append(aname)
                    self.params_.add(aname, value=self.p0[kk+2], min=-5,max=5)
                    
                    if kk != 0 and kk != 3:
                        lambvar.append('x' + str(kk))

        # Finish the string: closing parenthesis of log10
        fstr += ')'
        
        # Sort the lambda variables
        lambvar.sort()
        lamb = ['x0','x1','x2','x3','x4',
                'a0','a1','a2','a3','a4','a5','a6',
                'm2','m3','m4','m5','m6']
        
        return lamb, lambvar, lambcoeff, fstr

    def fit(self, X, y):
        '''
        Estimator fit function. Parameters:
            - X: Input parameters
            - y: Input base truth
        Steps:
            - Grab the correct initial condition from the dataframe and stores the parameter
            values in p0
            - Generate the functional string to be lambdified
            - Define the fit
            - Perform the fit
        '''
        ### Check X and Y
        X, y = check_X_y(X, y, accept_sparse=True,copy=True)
          
        ### Create the mask: we pick only those PI-Products
        self.mask_ = [self.m2,self.m3,self.m4,self.m5,self.m6]
        mask = self.mask_ # Necessary for the dataframe query
        
        ### Get the initial conditions by querying the dataframe
        qry = 'm2==@mask[0] and m3==@mask[1] and m4==@mask[2] and m5==@mask[3] and m6==@mask[4]'
        xf = self.icdf.query(qry).reset_index(drop=True)
        xf = xf[['p0','p1','p2','p3','p4','p5','p6']]
        self.p0 = np.array(xf.loc[self.icidx])
        
        ##############################
        ### Build the function to fit
        ##############################
        lamb, lambvar, lambcoeff, fstr = self.create_function_string()

        # Parse the string to create a callable function
        self.func_ = lambdify(lamb,parse_expr(fstr),modules='numpy')

        def fun(x0=1,x1=1,x2=1,x3=1,x4=1,
                a0=4,a1=0,a2=0,a3=0,a4=0,a5=0,a6=0,
                m2=1,m3=1,m4=1,m5=1,m6=1):    
            return self.func_(x0,x1,x2,x3,x4,a0,a1,a2,a3,a4,a5,a6,m2,m3,m4,m5,m6)

        ##############################
        ### Define LMFIT model
        ##############################
        # Define the lmfit model; omit nan's to avoid any exceptions
        self.fit_model_ = Model(fun, 
                                independent_vars=lambvar,
                                param_names=lambcoeff,
                                nan_policy='propagate')
        
        ##############################
        ### Perform the fit
        ##############################
        self.fit_ = self.fit_model_.fit(np.log10(y),
                                        params=self.params_,
                                        x0=X[:,0],
                                        x1=X[:,1],
                                        x2=X[:,2],
                                        x3=X[:,3],
                                        x4=X[:,4],
                                        m2=1.0,
                                        m3=1.0,
                                        m4=1.0,
                                        m5=1.0,
                                        m6=1.0,
                                        method='least_squares')
        
        # The model is now fitted.
        self.is_fitted_ = True
        return self        

    def predict(self, X):
        '''
        Estimator predict function. Use the function self.func_ defined in the "fit" method. 
        We check that the estimator has been fitted (and, therefore, func_ has been defined).
        '''
        X = check_array(X, accept_sparse=True,copy=True)
        check_is_fitted(self, 'is_fitted_')

        # Extract params
        alist = []
        
        for kk in range(7):
            try:
                alist.append(self.fit_.params['a'+str(kk)].value)
            except:
                alist.append(1e10)
         
        mlist = []
        for kk in range(X.shape[1]):
            mlist.append(1.0)
        
        # Extract variables
        x0 = X[:,0]
        x1 = X[:,1]
        x2 = X[:,2]
        x3 = X[:,3]
        x4 = X[:,4]
                
        # Predicted Y using PI2-PI6 and the coefficients
        Yp = self.func_(x0,x1,x2,x3,x4,*alist,*mlist)

        return Yp

    def score(self, X,y):
        '''
        Scoring function. We use the AIC as computed by the estimator itself.
        '''
        X, y = check_X_y(X, y, accept_sparse=True,copy=True)
        y = np.log10(y)
        
        check_is_fitted(self, 'is_fitted_')
                
        nsamples = X.shape[0]

        aic = 1e10
        # Predicted y
        ypred = self.predict(X)

        self.mse_ = mean_squared_error(y,ypred)
        
        aic = self.fit_.aic

        if np.isnan(aic):
            print(self.n_parameters,nsamples,self.mse_)
        
        return -aic
