
# In[] Model for contemporaneous estimation
def stspestmp_contemp(datainlist):
    import numpy as np
    import statsmodels.api as sm        
    import pandas as pd
    import datetime as dt
    
    # Unpack the information in the list
    samp, ref1, ref2, y, xonr, onrvar, datesall  = datainlist
    
    # Function that chooses optimization parameters and fits the model
    def optim_details(mod):

        inititer, initfuntol = 1000, 1e-10
        finfuniter, finfuntol = 10000, 1e-10

        modinit = mod.fit(method = 'minimize',
                          min_method = 'L-BFGS-B',
                          maxiter = inititer,
                          ftol = initfuntol,
                          disp = 0) # Initialize the fit to get new starting values 

        modfit = mod.fit(method = 'minimize',
                          min_method = 'L-BFGS-B',
                          maxiter = finfuniter,
                          ftol = finfuntol,
                          start_params = modinit.params,
                          disp = 0) # Fit it using starting values   
        return modfit

    # Define the RW model with controls (for currencies with ONR variation)
    class StSpMod_wCntrls(sm.tsa.statespace.MLEModel):
    
        def __init__(self, y, x):
            super(StSpMod_wCntrls, self).__init__(endog = y,
                                                  exog = x,
                                                  k_states=1, 
                                                  k_posdef=1,
                                                  initialization = 'diffuse')
    
                
            # Specify the fixed elements of the state space matrices (things that are assumed)
            self['design', 0, 0] = 1.0 # Design matrix is 2x1 of ones
            self['design', 1, 0] = 1.0
            self['transition', 0, 0] = 1.0 # Transition and selection matrices are 1x1 of ones
            self['selection', 0, 0] = 1.0         
            
            if len(np.shape(x)) == 1:
                self.k_exog = 1
            else:
                self.k_exog = np.shape(x)[-1]
    
            self['obs_intercept'] = np.zeros((2,self.nobs))        
            
            # Declare which parameters need to be positive (reference param_names to check )
            self.positive_parameters = slice(0,3) 
    
        @property
        def param_names(self):
            param_names =  [
                'state_var',
                'obs_var1','obs_var2',
                'ref1ref2_spread',
            ]
            for n in range(self.k_exog):
                param_names.extend(['betavar' + str(n) + '_1',
                                   'betavar' + str(n) + '_2'])
            return param_names
        
        @property
        def start_params(self):
            params = [
                0.01,
                0.01,0.01,
                0,
            ]
            for n in range(self.k_exog):
                params.extend([0,0])
            return params
        
        # Constrain the variance parameters to be positive and no smaller than "minvar"
        def transform_params(self, unconstrained):
            minvar = 1e-7
            constrained = unconstrained.copy()
            constrained[self.positive_parameters] = constrained[self.positive_parameters] ** 2 + minvar
            return constrained
    
        def untransform_params(self, constrained):
            minvar = 1e-7
            unconstrained = constrained.copy()
            unconstrained[self.positive_parameters] = (unconstrained[self.positive_parameters] - minvar)** 0.5         
            return unconstrained   
        
        # Define how estimated parameters map to the state space matrices 
        def update(self, params, **kwargs):
            params = super(StSpMod_wCntrls, self).update(params, **kwargs)
            
            self['state_cov', 0, 0] = params[0] # Single state variance in periods 0,2,...,T-1
            
            self['obs_cov', 0, 0] = params[1] # Noise variance in periods 0,2,...,T-1        
            self['obs_cov', 1, 1] = params[2] # Noise variance in periods 1, 3,...,T
            
            self['obs_intercept', 0, :] = params[3] + self.exog @ params[4:len(params):2]      
            self['obs_intercept', 1, :] = self.exog @ params[5:len(params):2]  
            
    # Define the RW model with no controls
    class StSp_Mod_NoCntrls(sm.tsa.statespace.MLEModel):
    
        def __init__(self, y):
            super(StSp_Mod_NoCntrls, self).__init__(endog = y,
                                                    k_states=1, 
                                                    k_posdef=1,
                                                    initialization = 'diffuse')
    
                
            # Specify the fixed elements of the state space matrices (things that are assumed)
            self['design', 0, 0] = 1.0 # Design matrix is 2x1 of ones
            self['design', 1, 0] = 1.0
            self['transition', 0, 0] = 1.0 # Transition and selection matrices are 1x1 of ones
            self['selection', 0, 0] = 1.0         
                
            self['obs_intercept'] = np.zeros((2,self.nobs))        
            
            # Declare which parameters need to be positive (reference param_names to check )
            self.positive_parameters = slice(0,3) 
    
        @property
        def param_names(self):
            param_names =  [
                'state_var', 
                'obs_var1','obs_var2',
                'ref1ref2_spread',
            ]
            return param_names
        
        @property
        def start_params(self):
            params = [
                0.01,
                0.01,0.01,
                0,
            ]
            return params
        
        # Constrain the variance parameters to be positive and no smaller than "minvar"
        def transform_params(self, unconstrained):
            minvar = 1e-7
            constrained = unconstrained.copy()
            constrained[self.positive_parameters] = constrained[self.positive_parameters] ** 2 + minvar
            return constrained
    
        def untransform_params(self, constrained):
            minvar = 1e-7
            unconstrained = constrained.copy()
            unconstrained[self.positive_parameters] = (unconstrained[self.positive_parameters] - minvar)** 0.5         
            return unconstrained   
        
        # Define how estimated parameters map to the state space matrices 
        def update(self, params, **kwargs):
            params = super(StSp_Mod_NoCntrls, self).update(params, **kwargs)
            
            self['state_cov', 0, 0] = params[0] # State variance in periods 0,2,...,T-1
            
            self['obs_cov', 0, 0] = params[1] # Noise variance in periods 0,2,...,T-1        
            self['obs_cov', 1, 1] = params[2] # Noise variance in periods 1, 3,...,T
            
            self['obs_intercept', 0, :] = params[3]

    # Estimate the relevant model (StSpMod_wCntrls if there is variation in ONR, StSpMod_NoCntrls if not)
    if onrvar > 1e-10:
        regtype = StSpMod_wCntrls
        mod = regtype(y, np.reshape(xonr, (len(xonr),1)))
    else:
        regtype = StSp_Mod_NoCntrls
        mod = regtype(y) # Model with policy variables for other ccys
        
    modfit = optim_details(mod)
    
    # Get the smoothed disturbance terms for main estimates (Table 3) and store in a dataFrame
    if samp != 'rolling':
    
        dist_smooth = modfit.smoothed_measurement_disturbance.T

        dfdist = pd.DataFrame({'datefmt' : datesall,
                               'alt' : ref2,
                               'sample' : samp,
                               'obserr_libor' : dist_smooth[:,0],
                               'obserr_arr' : dist_smooth[:,1]})
    else:
        dfdist = 'rolling_na'
    
    wald_test_ns = modfit.wald_test('obs_var1 - obs_var2 = 0')
    
    # Store all relevant results and details for each  model     
    regrslts = pd.DataFrame({'alt' : ref2,
                             'sample' : samp,
                             'date1' : datesall.iloc[0],
                             'dateend' : datesall.iloc[-1],
                             'model' : regtype.__name__,
                             'wald_test' : wald_test_ns,
                             'convergence' : modfit.mle_retvals['converged'],
                             'condnmb' : str("{:e}".format(np.linalg.cond(modfit.cov_params()))),
                             'ressum' : modfit.summary(),
                             'params' : [modfit.params],
                             'tvalues' : [modfit.tvalues],
                             'paramnames' : [modfit.param_names],
                             'llf' : modfit.llf,
                             'distsmooth' : [dfdist]
                              },
                           index = [0])

    print(ref2, samp, datesall.iloc[0], datesall.iloc[-1],
          dt.datetime.now().time().strftime('%H:%M:%S'))
    return regrslts