import sunode
import sunode.wrappers.as_theano

def SIR_sunode(t, y, p):
    return {
        'S': -p.lam * y.S * y.I,
        'I': p.lam * y.S * y.I - p.mu * y.I}
        
  ...
  ...
   
  sir_curves, _, problem, solver, _, _ = sunode.wrappers.as_theano.solve_ivp(
        y0={ # Initial conditions of the ODE
            'S': (S_init, ()),
            'I': (I_init, ()),
        },
        params={
	         # Parameters of the ODE, specify shape
            'lam': (lam, ()),
            'mu': (mu, ()),
            '_dummy': (np.array(1.), ())  # currently, sunode throws an error
        },                                # without this
	        # RHS of the ODE
        rhs=SIR_sunode,
	        # Time points of th solution
        tvals=times,
        t0=times[0],
    )

with pm.Model() as model4:
            sigma = pm.HalfCauchy('sigma', self.likelihood['sigma'], shape=1)
            lam = pm.Lognormal('lambda', self.prior['lam'], self.prior['lambda_std']) # 1.5, 1.5
            mu = pm.Lognormal('mu', self.prior['mu'], self.prior['mu_std'])           # 1.5, 1.5
            res, _, problem, solver, _, _ = sunode.wrappers.as_theano.solve_ivp(
            y0={
                'S': (self.S_init, ()), 'I': (self.I_init, ()),},
            params={
                'lam': (lam, ()), 'mu': (mu, ()), '_dummy': (np.array(1.), ())},
            rhs=self.SIR_sunode,
            tvals=self.time_range,
            t0=self.time_range[0]
            )
            if(likelihood['distribution'] == 'lognormal'):
                I = pm.Lognormal('I', mu=res['I'], sigma=sigma, observed=self.cases_obs_scaled)
            elif(likelihood['distribution'] == 'normal'):
                I = pm.Normal('I', mu=res['I'], sigma=sigma, observed=self.cases_obs_scaled)
            elif(likelihood['distribution'] == 'students-t'):
                I = pm.StudentT( "I",  nu=likelihood['nu'],       # likelihood distribution of the data
                        mu=res['I'],     # likelihood distribution mean, these are the predictions from SIR
                        sigma=sigma,
                        observed=self.cases_obs_scaled
                        )
            R0 = pm.Deterministic('R0',lam/mu)

            trace = pm.sample(self.n_samples, tune=self.n_tune, chains=4, cores=4)
            data = az.from_pymc3(trace=trace)

covid_obj = COVID_data('US', Population=328.2e6)
covid_obj.get_dates(data_begin='10/1/20', data_end='10/28/20')
sir_model = SIR_model_sunode(covid_obj)
likelihood = {'distribution': 'normal', 
              'sigma': 2}
prior = {'lam': 1.5, 
         'mu': 1.5, 
         'lambda_std': 1.5,
         'mu_std': 1.5 }
sir_model.run_SIR_model(n_samples=500, n_tune=500, likelihood=likelihood, prior=prior)