import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import os
import pandas as pd

def func_ode(y,t,qS,qSq,betaV1,betaV2,beta1,beta2,e1,eta1,eta2,etaH,delta2,deltaH,theta2,f1,f2,fH,d):

    Sq,S,E1,E2,H,R,D,V=y

    dSq=qS*S-qSq*Sq
    dS=qSq*Sq-(betaV1*V+betaV2*V+beta1*E1+beta2*E2)*S
    dE1=(betaV1*V+beta1*E1)*S-(e1+eta1)*E1
    dE2=(betaV2*V+beta2*E2)*S+e1*E1-(eta2+delta2+theta2)*E2
    dH=theta2*E2-(etaH+deltaH)*H # theta is for recovered
    dR=eta1*E1+eta2*E2+etaH*H # eta is for recovered
    dD=delta2*E2+deltaH*H # delta is for death
    dV=f1*E1+f2*E2+fH*H-d*V

    dy=[dSq, dS, dE1, dE2, dH, dR, dD, dV]
    return dy

def func_y(t,qS,qSq,betaV1,betaV2,beta1,beta2,e1,eta1,eta2,etaH,delta2,deltaH,theta2,f1,f2,fH,d,y_0):
    """
    Solution to the ODE y'(t) = f(t,y,parameters) with initial condition y(0) = y_0
    """
    y = odeint(func_ode, y_0, t, args=(qS,qSq,betaV1,betaV2,beta1,beta2,e1,eta1,eta2,etaH,delta2,deltaH,theta2,f1,f2,fH,d))
    return y.flatten()

if __name__ == "__main__"# print(numpy_data)
    # row,col=np.shape(numpy_data)
    # for i in range(row):
    #     for j in range(col):
    #         print(type(numpy_data[i,j]))
    # exit()

    ## Parameters (dummy)
    qS, qSq, betaV1, betaV2, beta1, beta2, e1, eta1, eta2, etaH, delta2, deltaH, theta2, f1, f2, fH, d = \
        0, 1e-4, 4e-9, 1e-9, 4e-9, 1e-9, 1/100, 1/21, 1/104, 1/10, 1/200, 1/10400, 1/3.5, 1400, 1000, 1700, 144

    # Used Data
    y_0=numpy_data[0,:].tolist()
    numpy_data=numpy_data[1:60,:]

    ## Reference Time
    number_of_reference_time,_=np.shape(numpy_data)

    ## Solve
    param = qS, qSq, betaV1, betaV2, beta1, beta2, e1, eta1, eta2, etaH, delta2, deltaH, theta2, f1, f2, fH, d, y_0
    t= np.linspace(1,number_of_reference_time,number_of_reference_time)
    popt, cov = curve_fit(func_y, t, numpy_data.flatten(),p0=[param])

    qS,qSq,betaV1,betaV2,beta1,beta2,e1,eta1,eta2,etaH,delta2,deltaH,theta2,f1,f2,fH,d = popt

    ## Check Result
    result=odeint(func_ode,y_0,t,args=(qS,qSq,betaV1,betaV2,beta1,beta2,e1,eta1,eta2,etaH,delta2,deltaH,theta2,f1,f2,fH,d))

    ## Plot
    plt.plot(t, result[:, 0], label='Sq')
    plt.plot(t, result[:, 1], label='S')
    plt.plot(t, result[:, 2], label='E1')
    plt.plot(t, result[:, 3], label='E2')
    plt.plot(t, result[:, 4], label='H')
    plt.plot(t, result[:, 5], label='R')
    plt.plot(t, result[:, 6], label='D')
    plt.plot(t, result[:, 7], label='V')
    plt.legend(loc='best')
    plt.xlabel('t')
    plt.grid()
    plt.show()
    pass