import pandas as pd
import streamlit as st
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")

#dowload file

#read files
data = pd.read_csv('owid-monkeypox-data.csv')
data = data[['location','iso_code','date','new_cases','total_cases','new_deaths','total_deaths']]

pop = pd.read_csv('API_SP.POP.TOTL_DS2_en_csv_v2_4578059.csv')
#preprocessiong data
all_location = {}
for i in data['iso_code'].unique():
  all_location[i] = data[data['iso_code'] == i].reset_index(drop=True)

popu = pop[['Country Code','2021']].to_dict('index')
pop_dict = {}
for i in popu.values():
  pop_dict[i['Country Code']] = i['2021']

pop_dict['GLP'] = 400000
pop_dict['MTQ'] = 376480
pop_dict['OWID_WRL'] = 7836630792

code = dict(data.groupby('location')['iso_code'].unique())

# SIR model differential equations.
def deriv(x, t, beta, gamma):
    s, i, r = x
    dsdt = -beta * s * i
    didt = beta * s * i - gamma * i
    drdt =  gamma * i
    return [dsdt, didt, drdt]

#plot model
def plotdata(t, s, i,r,R0, e=None):
    # plot the data
    fig = plt.figure(figsize=(12,6))
    ax = [fig.add_subplot(221, axisbelow=True), 
          fig.add_subplot(223),
          fig.add_subplot(122)]

    ax[0].plot(t, s, lw=3, label='Fraction Susceptible')
    ax[0].plot(t, i, lw=3, label='Fraction Infective')
    ax[0].plot(t, r, lw=3, label='Recovered')
    ax[0].set_title('Susceptible and Recovered Populations')
    ax[0].set_xlabel('Time /days')
    ax[0].set_ylabel('Fraction')

    ax[1].plot(t, i, lw=3, label='Infective')
    ax[1].set_title('Infectious Population')
    if e is not None: ax[1].plot(t, e, lw=3, label='Exposed')
    ax[1].set_ylim(0, 1.0)
    ax[1].set_xlabel('Time /days')
    ax[1].set_ylabel('Fraction')

    ax[2].plot(s, i, lw=3, label='s, i trajectory')
    ax[2].plot([1/R0, 1/R0], [0, 1], '--', lw=3, label='di/dt = 0')
    ax[2].plot(s[0], i[0], '.', ms=20, label='Initial Condition')
    ax[2].plot(s[-1], i[-1], '.', ms=20, label='Final Condition')
    ax[2].set_title('State Trajectory')
    ax[2].set_aspect('equal')
    ax[2].set_ylim(0, 1.05)
    ax[2].set_xlim(0, 1.05)
    ax[2].set_xlabel('Susceptible')
    ax[2].set_ylabel('Infectious')

    for a in ax: 
        a.grid(True)
        a.legend()

    plt.tight_layout()
    
    return fig
     

#final model
def SIR(country,t_infective):
  # parameter values
  t_infective = t_infective
  gamma = 1/t_infective
  beta = (all_location[country]['new_cases'].sum()/pop_dict[country])/len(all_location[country]['date'].unique())
  R0 = beta/gamma

  # initial number of infected and recovered individuals
  i_initial = all_location[country]['new_cases'].sum()/pop_dict[country]
  r_initial = 0.00
  s_initial = 1 - i_initial - r_initial
  
  t = np.linspace(0, 100, 1000)
  x_initial = s_initial, i_initial, r_initial
  soln = odeint(deriv, x_initial, t, args=(beta, gamma))
  s, i, r = soln.T
  e = None


  return R0,t_infective,beta,gamma,plotdata(t, s, i,r,R0)

def main():
    st.title("SIR Model for Monkeypox")

    with st.form("questionaire"):
        country = st.selectbox("Country",data['location'].unique())# user's input
        recovery = st.slider("How long Monkeypox recover?", 21, 31, 21)# user's input
        country = code[country][0]
        # clicked==True only when the button is clicked
        clicked = st.form_submit_button("Show Graph")
        if clicked:
        
            # Show SIR
            SIR_param = SIR(country,recovery)
            
            st.pyplot(SIR_param[-1])
            st.success("SIR model parameters for "+str(country)+" is")
            st.success("R0 = "+str(SIR_param[0]))
            st.success("Beta = "+str(SIR_param[2]))
            st.success("Gamma = "+str(SIR_param[3]))

# Run main()
if __name__ == "__main__":
    main()