from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
from scipy.integrate import solve_ivp
import numpy as np
from tqdm import tqdm
from scipy.optimize import curve_fit
from functools import partial
import multiprocessing as mp
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import MultiTaskLassoCV , Lasso
from sklearn.multioutput import MultiOutputRegressor
import xgboost as xgb
    
def mae(pred, true):
    return np.mean(np.abs(pred - true))


class SIR_fitter():
    ''' Class that use the features to extract SIR parameters. Fitted parameters will be keyed by country and date. 
    infected_days: days previous to day0 to sum the number of currently infected.
    semi_fit_days: number of days before and after the actual day to fit the SIR parameters on.
    beta_i: initial value for fitting beta.
    gamma_i: initial value for fitting gamma.
    '''
    def __init__(self, moving_average=False, 
                 infection_days=7, semi_fit_days=7,
                 beta_i=0.6, gamma_i=1/7,nprocs=4):
        
        self.infection_days=infection_days
        self.semi_fit=semi_fit_days
        self.fit_days=semi_fit_days*2+1
        self.time_integ=np.linspace(-self.semi_fit,self.semi_fit,self.fit_days)
        #self.time_integ=np.linspace(0,1,2)  # questa
        self.beta_i=beta_i
        self.gamma_i=gamma_i
        self.nprocs=nprocs
        self.moving_average=moving_average
        self.tempone=np.linspace(0,10000,10001)
        
    def fit_country(self,df_country):
        COL = ['NewCases'] if not self.moving_average else ['MA']
        gdf = df_country
        gdf['beta']=np.nan
        gdf['gamma']=np.nan
        gdf['predicted_cases']=gdf[COL].copy()
        all_case_data = np.array(gdf[COL])
        target_cases = all_case_data.copy()
        nb_total_days = len(gdf)
        for d,(idx,row) in enumerate(gdf.iterrows()):
            if d>self.semi_fit and d<(nb_total_days-self.semi_fit):
                N=row.Population
                # X_cases = all_case_data[d - self.semi_fit:d+self.semi_fit+1].reshape(-1)
                X_true_cases = target_cases[d - self.semi_fit:d+self.semi_fit+1].reshape(-1)
                I0=all_case_data[d-self.semi_fit- self.infection_days: 
                                  d-self.semi_fit+1].sum()
                Ic0=all_case_data[0:(d-self.semi_fit)+1].sum()
                R0=Ic0-I0
                S0=N-I0-R0
                x0=(S0,I0,Ic0,R0)
                time_integ=self.tempone[d-self.semi_fit:d+self.semi_fit+1].reshape(-1)
                if I0<0:
                    raise ValueError('Infected was {} for population {}'.format(x0[1],N))
                elif I0<1:
                    popt=np.array([np.nan,np.nan])
                else:
                    fintegranda=partial(self.__SIR_integrate,time_integ,x0,N)
                    popt, pcov = curve_fit(fintegranda, time_integ, 
                            X_true_cases[1:],
                            p0=[self.beta_i,self.gamma_i],maxfev=5000,bounds=([0.,0.],
                                                            [np.inf,1.]))
                    gdf.loc[idx,'beta']=popt[0]
                    gdf.loc[idx,'gamma']=popt[1]
                
                    if d < nb_total_days-self.semi_fit-1:
                        fut_integ=self.tempone[d-self.semi_fit:d+self.semi_fit+2]
                        all_case_data[d+self.semi_fit+1]=fintegranda(fut_integ,
                                                                     popt[0],popt[1])[-1]
                        gdf.loc[idx,'predicted_cases']=all_case_data[d+self.semi_fit]
        return gdf
        
        
    def __SIR_ode(self,t,x0, N, beta, gamma):
        S, I, Ic, R = x0
        dS = -beta * S * I / N
        dI = beta * S * I / N - gamma * I
        # Computes the cumulative new cases during the integration period
        dIc= -dS
        dR = gamma * I
        return dS, dI, dIc,dR
    
    def __SIR_integrate(self,ttotp,x0,N,ti,beta,gamma):
        ''' Argument ti not used but needed by curve_fit '''
        sol=solve_ivp(self.__SIR_ode,[ttotp[0],ttotp[-1]],x0,args=(N,beta,gamma),t_eval=ttotp)
        #lung=len(sol.y[0])
        # The only variable to predict is "NewCases", i.d. the difference of the cumulative Ic
        return np.diff(sol.y[2])#.reshape(lung,1).flatten()
    
   
        
   
    def row_predict(self,GeoID,date,pars=None):
        '''
        Given the SIR parameters, predicts the new cases. The last one is the actual prediction for the final MAE
        '''
        row=self.df.loc[(self.df.GeoID==GeoID)&(self.df.Date==pd.to_datetime(date)),:].iloc[0,:]
        N,x0=self.row_initial_conditions(row)
        if np.isnan(pars[0]):
            return np.repeat(row.CasesDay0, self.fit_days)[1:]
        if pars is not None:
            beta=pars[0]
            gamma=pars[1]
        else:
            beta=self.df_pars.loc[(self.df_pars.GeoID==GeoID) & 
                                  (self.df_pars.Date==pd.to_datetime(date)),'beta'].iloc[0]
            gamma=self.df_pars.loc[(self.df_pars.GeoID==GeoID) & 
                                  (self.df_pars.Date==pd.to_datetime(date)),'gamma'].iloc[0]
        Ipred=self.__SIR_integrate(self.time_integ,x0,N,self.time_integ,beta,gamma)
        return Ipred
    
    def fit(self,df,save_to=None):
        '''
        Fit SIR parameters on all the data. 
        save_to: path to save the results in pickle format. Results are saved a Pandas DataFrame having columns: GeoID,Date,beta,gamma 
        '''
        if self.semi_fit<3:
            raise ValueError('ValueError: semi_fit_days should be higher than 2')
        geo_ids = df.GeoID.unique()
        COL = ['NewCases'] if not self.moving_average else ['MA']
        df=df[['GeoID','Date','Population']+COL]
        self.df_chunks=[df[df.GeoID==g].sort_values('Date') 
                        for g in geo_ids]
        nchunks=len(self.df_chunks)
        pool=mp.Pool(self.nprocs)
        outputs=list(tqdm(pool.imap(self.fit_country,self.df_chunks),total=nchunks))
        pool.close()
        pool.join()
        self.df_pars=pd.concat(outputs)
        self.df_pars.sort_values(['GeoID','Date'],inplace=True)
        if save_to is not None:
            with open(save_to,'wb') as f:
                pickle.dump(self.df_pars,f)
                
        # Return the classifier
        return self

class SIR_predictor(BaseEstimator, RegressorMixin, SIR_fitter):
    def __init__(self, df=None,moving_average=True, 
                 infection_days=7, semi_fit=3,
                 beta_i=0.6, gamma_i=1/7,lookback_days=30,
                 MLmodel= 'MultiOutputRegressor(xgb.XGBRegressor())', 
                 paral_predict=False,pre_computed=None,nprocs=4):
        self.df=df
        self.moving_average=moving_average
        self.infection_days=infection_days
        self.semi_fit=semi_fit
        self.beta_i=beta_i
        self.gamma_i=gamma_i
        self.lookback_days=lookback_days
        self.nprocs=nprocs
        self.paral_predict=paral_predict
        self.MLmodel=MLmodel
        self.pre_computed=pre_computed
        
    
    def SIR_ode(self,t,x0, N, beta, gamma):
        return self._SIR_fitter__SIR_ode(t,x0, N, beta, gamma)
    def SIR_integrate(self,ttotp,x0,N,ti,beta,gamma):
        return self._SIR_fitter__SIR_integrate(ttotp,x0,N,ti,beta,gamma)
    def fit(self,X,y):
        check_X_y(X,y)
        if self.pre_computed is None:
            if 'MA' in self.df.columns:
                self.df.loc[self.df.MA<0,'MA']=0.
            if 'NewCases' in self.df.columns:
                self.df.loc[self.df.NewCases<0,'NewCases']=0.
            self.df.loc[self.df.ConfirmedCases<0,'ConfirmedCases']=0.

            self.SFmodel=SIR_fitter(self.moving_average, 
                     self.infection_days, self.semi_fit,
                     self.beta_i, self.gamma_i,self.nprocs)
            print('Fitting SIR parameters...')
            self.SFmodel.fit(self.df)
        else:
            print('Already computed SIR parameters')
            self.SFmodel=SIR_fitter(self.moving_average, 
                     self.infection_days, self.semi_fit,
                     self.beta_i, self.gamma_i,self.nprocs)
            self.SFmodel.df_pars=self.pre_computed
        #self.SFmodel.df_pars=self.df.copy()
        #self.SFmodel.df_pars['beta']=0.6
        #self.SFmodel.df_pars['gamma']=1/7
        self.df=self.df.merge(self.SFmodel.df_pars[['GeoID','Date','beta','gamma']],
                 how='left',on=['GeoID','Date'],left_index=True).dropna(subset=['beta','gamma'])
        #print(self.df.loc[X[:,-1]])
        
        #Predict SIR parameters instead of cases
        #Use only row in training set AND with SIR pars not being nans
        idx=np.array(list(set(X[:,-1]).intersection(set(self.df.index))),
                     dtype=int).reshape(-1)
        #print(len(idx),idx[:10])
        self.y_pars=np.array(self.df.loc[idx, ['beta','gamma']])
        #print(self.y_pars.shape)
        self.X_pars=X[np.in1d(X[:,-1],idx)]
        #print(self.X_pars.shape)
        # remove last column (df.index)
        # remove first lookback_days columns: not using cases to predict parameters
        self.X_pars=self.X_pars[:,self.lookback_days+1:-1]
        #print('ML training:',self.X_pars.shape)
        #exit()
        #self.X_pars=self.X_pars[:,:-1]
        
        self.MLmodel=eval(str(self.MLmodel))
        #print('\n ML model: ',type(self.MLmodel))
        self.MLmodel.fit(self.X_pars,self.y_pars)
        self.TMAE=mae(self.MLmodel.predict(self.X_pars),self.y_pars)
        print('Training MAE on SIR params:', self.TMAE)
        return self
    
    def predict_pars(self,X):
        return self.MLmodel.predict(X[:,self.lookback_days+1:-1])
        #except Exception as e:
        #    print(e)
        #    print(X[:,self.lookback_days+1],X[:,-1])
        #    print(X[:,self.lookback_days+1:-1].shape)
        #    exit()
        #return self.MLmodel.predict(X[:,:-1])
    
    def predict_chunk(self,X_chunk):
        y_chunk=[]
        for i in range(X_chunk.shape[0]):
            pars=self.predict_pars(X_chunk[i,:].reshape(1,-1))
            beta=pars[0][0]
            gamma=pars[0][1]
            time_integ=[0,1]
            
            N=X_chunk[i,self.lookback_days+1]
            #I0=X_chunk[i,self.lookback_days-1].sum()
            #Ic0=X_chunk[i,self.lookback_days]
            #R0=Ic0-I0
            
            I0=X_chunk[i,self.lookback_days-self.infection_days:self.lookback_days].sum()
            Ic0=X_chunk[i,self.lookback_days]
            R0=Ic0-I0
            
            S0=N-I0-R0
            x0=(S0,I0,Ic0,R0)
            #print(x0)
            
            #I=solve_ivp(self.SIR_ode,[time_integ[0],time_integ[-1]],
            #            x0,args=(N,beta,gamma),t_eval=time_integ).y[1][1]
            I=np.diff(solve_ivp(self.SIR_ode,[time_integ[0],time_integ[-1]],
                        x0,args=(N,beta,gamma),t_eval=time_integ).y[2])[-1]
            y_chunk.append(I)
        return y_chunk
    
    def predict(self,X):
        if self.paral_predict:
            nchunks=self.nprocs*10
            X_chunks=np.array_split(X,nchunks)
            pool=mp.Pool(self.nprocs)
            y_chunks=list(pool.imap(self.predict_chunk,X_chunks))
            pool.close()
            pool.join()
            y=[item for sublist in y_chunks for item in sublist]
            return y
        else:
            y=self.predict_chunk(X)
            return y
    
    def score(self,X_test,y_test):
        #check_X_y(X_test,y_test)
        #y_pred=self.predict(X_test)
        # Take negative to mean "score"
        #cases_mae=mae(y_pred,y_test)
        #print('Training MAE on cases:',cases_mae)
        #return -cases_mae
        return -self.TMAE
        
#Questa non ha il fit, non è usabile
class SIRRegressor(BaseEstimator, RegressorMixin):
    ''' Model that use the features to extract SIR parameters and compute predictions. 
    single_pred_days: number of next days to predict. Predictions could be of more than one day, so to have a multi-variate regression (NOT TESTED YET)'''
    def __init__(self, single_pred_days=1,lookback_days=30,infection_days=15):
        
        #self.demo_param = demo_param
        self.params=None
        self.single_pred_days=single_pred_days
        self.lookback_days=lookback_days
        self.infection_days=infection_days
        
    def __SIR_ode(self,t,x0, N, beta, gamma):
        S, I, Ic, R = x0
        dS = -beta * S * I / N
        dI = beta * S * I / N - gamma * I
        # Computes the cumulative new cases during the integration period
        dIc= -dS
        dR = gamma * I
        return dS, dI, dIc,dR
    
    def __SIR_integrate(self,ttotp,x0,N,beta,gamma):
        sol=solve_ivp(self.__SIR_ode,[ttotp[0],ttotp[-1]],x0,args=(N,beta,gamma),t_eval=ttotp)
        #lung=len(sol.y[0])
        # The only variable to predict is "NewCases", i.d. the difference of the cumulative Ic
        return np.diff(sol.y[2],prepend=x0[2])#.reshape(lung,1).flatten()
        
        
        
    def fit(self, X, y):
        # Check that X and y have correct shape
        X, y = check_X_y(X, y)
        self.X_ = X
        self.y_ = y
        
        #Assign random parameters, to be used for scalar product with the features.
        # We sample two rows: one to get beta and one to get gamma
        self.params=np.random.uniform(0,1,(2,X.shape[1]))
        # Return the classifier
        return self
    
    def predict(self, X):
        # Check is fit had been called
        check_is_fitted(self)
        # Input validation
        X = check_array(X)
        
        # Normalization is to have small values also with random parameters (to be removed later maybe)
        X_normed = (X - X.min(0)) / X.ptp(0)
        X_normed=np.nan_to_num(X_normed)
        
        y_pred=[]
        time_integ=np.linspace(0,self.single_pred_days,self.single_pred_days+1)
        for i in tqdm(range(X.shape[0]),total=X.shape[0]):
            # Apply ML model parameters to get SIR parameters
            # Division for X.shape[1]  should ensure values lower than 1, but should be removed when there will be a real training
            beta=self.params[0].dot(X_normed[i])/X.shape[1]
            gamma=self.params[1].dot(X_normed[i])/X.shape[1]
            
            # Total population
            N=X[i,self.lookback_days+1]
            
            # Currently infected individuals
            I0=X[i,self.lookback_days-self.infection_days:self.lookback_days-1].sum()
            
            # Recovered individuals (taken as current total confirmed cases)
            R0=X[i,self.lookback_days]
            
            # Susceptible individuals
            S0=N-I0-R0
            
            # Initial condition of integration
            x0=(S0,I0,R0,R0)
            Ipred=self.__SIR_integrate(time_integ,x0,N,beta,gamma)[-1]
            y_pred.append(Ipred)
        
        return np.array(y_pred)
    


class SIR_parfinder():
    ''' DO NOT USE: Class that use the features to extract SIR parameters. Fitted parameters will be keyed by country and date. 
    single_pred_days: number of next days to predict. Predictions could be of more than one day, so to have a multi-variate regression (NOT TESTED YET)
    lookback_days: past days to use. Must be higher than fit_days+infected_days.
    infected_days: days previous to day0 to sum the number of currently infected.
    fit_days: number of days before the last day to fit the SIR parameters on (only for labelization).
    beta_i: initial value for fitting beta (only for labelization).
    gamma_i: initial value for fitting gamma (only for labelization).
    '''
    def __init__(self, df,moving_average=False, 
                 infection_days=7, semi_fit_days=7,
                 beta_i=0.6, gamma_i=1/7,nprocs=4):
        
        self.infection_days=infection_days
        self.semi_fit=semi_fit_days
        self.fit_days=semi_fit_days*2+1
        self.time_integ=np.linspace(-self.semi_fit,self.semi_fit,self.fit_days)
        self.beta_i=beta_i
        self.gamma_i=gamma_i
        self.nprocs=nprocs
        self.df=self.__add_lookback(df,moving_average)
    
    def __add_lookback(self,df,moving_average):
        COL = ['NewCases'] if not moving_average else ['MA']
        X_cols = df.columns
        y_col = COL
        geo_ids = df.GeoID.unique()
        fit_ids=['CasesDay{}'.format(d) for d in range(-self.semi_fit,self.semi_fit+1)]
        self.fit_ids=fit_ids
        #print(fit_ids)
        df=df[['GeoID','Date','Population']+COL]
        df=pd.concat([df,pd.DataFrame(columns=fit_ids+['I0','R0'])])
        print('Adding lookback days to the dataframe...')
        for g in tqdm(geo_ids):
            gdf = df[df.GeoID == g]
            all_case_data = np.array(gdf[COL])
            nb_total_days = len(gdf)
            for d,(idx,row) in enumerate(gdf.iterrows()):
                if d>self.semi_fit and d<(nb_total_days-self.semi_fit):
                    X_cases = all_case_data[d - self.semi_fit:d+self.semi_fit+1].reshape(-1)
                    try:
                        df.loc[idx,fit_ids]=X_cases
                        df.loc[idx,'I0']=all_case_data[d-self.semi_fit- self.infection_days: 
                                                       d-self.semi_fit+1].sum()
                        df.loc[idx,'R0']=all_case_data[0: 
                                                       (d-self.semi_fit-
                                                       self.infection_days)].sum()
                    except ValueError:
                        print(row.GeoID)
                        print(row.Date)
                        print(df.loc[idx-self.semi_fit,fit_ids].shape)
                        print(X_cases.shape)
                        raise ValueError('Mismatch in shapes for this entry, check the code...')
                        
        return df.dropna()
        
        
    def __SIR_ode(self,t,x0, N, beta, gamma):
        S, I, Ic, R = x0
        dS = -beta * S * I / N
        dI = beta * S * I / N - gamma * I
        # Computes the cumulative new cases during the integration period
        dIc= -dS
        dR = gamma * I
        return dS, dI, dIc,dR
    
    def __SIR_integrate(self,ttotp,x0,N,ti,beta,gamma):
        ''' Argument ti not used but needed by curve_fit '''
        sol=solve_ivp(self.__SIR_ode,[ttotp[0],ttotp[-1]],x0,args=(N,beta,gamma),t_eval=ttotp)
        #lung=len(sol.y[0])
        # The only variable to predict is "NewCases", i.d. the difference of the cumulative Ic
        return np.diff(sol.y[2])#.reshape(lung,1).flatten()
    
    def labelize_chunk(self,df_chunk):
        #df_chunk=self.df_chunks[i]
        pars_df=df_chunk[['GeoID','Date']].copy()
        pars_df['beta']=np.nan
        pars_df['gamma']=np.nan
        for j,(idx,row) in enumerate(df_chunk.iterrows()):
            pars_df.iloc[j,:]=[row.GeoID,row.Date]+list(self.row_fit(row))
        return pars_df        
        
    def row_initial_conditions(self,row):
        '''
        Returns the initial condition from the the SIR integration of this row.
        '''
        N=row.Population           
        # Currently infected individuals (sum of new cases on the previous infection_days before the first fit day)
        I0=row.I0          
        # Recovered individuals (taken as current total confirmed cases)
        R0=row.R0     
        # Susceptible individuals
        S0=N-I0-R0            
        # Initial condition of integration
        x0=(S0,I0,R0,R0)
        return N,x0
    
    def row_observed_cases(self,row):
        '''
        Return an array with length fit_days with the real observed new_cases of that row.
        '''
        return row[self.fit_ids[1:]].values
    
    def row_fit(self,row):
        '''
        Fit SIR parameters for one observation
        '''
        N,x0=self.row_initial_conditions(row)
        if x0[1]<0:
            raise ValueError('Infected was {} for popolation {}'.format(x0[1],X_i[self.lookback_days+1]))
        elif x0[1]<1:
            popt=np.array([np.nan,np.nan])
        else:
            fintegranda=partial(self.__SIR_integrate,self.time_integ,x0,N)
            popt, pcov = curve_fit(fintegranda, self.time_integ, 
                       self.row_observed_cases(row),
                       p0=[self.beta_i,self.gamma_i],maxfev=5000,bounds=([0.,0.],
                                                           [np.inf,1.]))
        return popt.reshape(-1)
        
    def row_predict(self,GeoID,date,pars=None):
        '''
        Given the SIR parameters, predicts the new cases. The last one is the actual prediction for the final MAE
        '''
        row=self.df.loc[(self.df.GeoID==GeoID)&(self.df.Date==pd.to_datetime(date)),:].iloc[0,:]
        N,x0=self.row_initial_conditions(row)
        if np.isnan(pars[0]):
            return np.repeat(row.CasesDay0, self.fit_days)[1:]
        if pars is not None:
            beta=pars[0]
            gamma=pars[1]
        else:
            beta=self.df_pars.loc[(self.df_pars.GeoID==GeoID) & 
                                  (self.df_pars.Date==pd.to_datetime(date)),'beta'].iloc[0]
            gamma=self.df_pars.loc[(self.df_pars.GeoID==GeoID) & 
                                  (self.df_pars.Date==pd.to_datetime(date)),'gamma'].iloc[0]
        Ipred=self.__SIR_integrate(self.time_integ,x0,N,self.time_integ,beta,gamma)
        return Ipred
    
    def fit(self,save_to=None):
        '''
        Fit SIR parameters on all the data. 
        save_to: path to save the results in pickle format. Results are saved a Pandas DataFrame having columns: GeoID,Date,beta,gamma 
        '''
        if self.semi_fit<3:
            raise ValueError('ValueError: semi_fit_days should be higher than 2')
        
        nchunks=self.nprocs*10
        self.df_chunks = np.array_split(self.df,nchunks)
        #print(type(self.df_chunks))
        empty_cunks=[df for df in self.df_chunks if df.shape[0]==0]
        if len(empty_cunks):
            print('{} empty chunks'.format(len(empty_chunks)))
        #nchunks=len(self.X_chunks)
        pool=mp.Pool(self.nprocs)
        outputs=list(tqdm(pool.imap(self.labelize_chunk,self.df_chunks),total=nchunks))
        pool.close()
        pool.join()
        self.df_pars=pd.concat(outputs)
        self.df_pars.sort_values(['GeoID','Date'],inplace=True)
        if save_to is not None:
            with open(save_to,'wb') as f:
                pickle.dump(self.df_pars,f)
                
        # Return the classifier
        return self
    

scipy.optimize.curve_fit scipy.integrate.solve_ivp