import gradio as gr
from cadCAD.configuration.utils import config_sim
from cadCAD.configuration import Experiment

from cadCAD.engine import ExecutionMode, ExecutionContext
from cadCAD.engine import Executor

from cadCAD import configs
import pandas as pd
import numpy as np

from seir_model import *

def create_and_run_exp(population):
    #Total Population : population
    initial_state = {
        "S": population-10,
        "E": 10,
        "I":0,
        "R":0
    }
    sys_params = {
        "infection_rate":[1],
        "recovery_rate":[1/14], 
        "exposure_rate":[1/3]
    }
        
    partial_state_update_blocks = [
        {
            "policies":{
                "expsoed_growth":p_exposed,
                "infected_growth":p_infected,
                "recovered_growth":p_recovered,
                
            },
            "variables":{
                "S":s_susceptible,
                "E":s_exposed,
                "I":s_infected,
                "R":s_recovered,
            }
         
        }
    
    ]
    
    del configs[:]
    
    timesteps = 100
    sim_config = config_sim({
      "N":1,
      "T":range(timesteps),
      "M":sys_params  
    })
    
    experiment = Experiment()
    experiment.append_configs(
        sim_configs=sim_config,
        initial_state=initial_state,
        partial_state_update_blocks=partial_state_update_blocks
    )

    exec_context = ExecutionContext()
    simulation = Executor(exec_context=exec_context, configs=experiment.configs)
    return simulation

css = """
        .gradio-container {
            font-family: 'IBM Plex Sans', sans-serif;
        }
        .gr-button {
            color: white;
            border-color: black;
            background: black;
        }
        input[type='range'] {
            accent-color: black;
        }
        .dark input[type='range'] {
            accent-color: #dfdfdf;
        }
        .container {
            max-width: 900px;
            margin: auto;
            padding-top: 1.5rem;
        }
        .details:hover {
            text-decoration: underline;
        }
        .gr-button {
            white-space: nowrap;
        }
        .gr-button:focus {
            border-color: rgb(147 197 253 / var(--tw-border-opacity));
            outline: none;
            box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
            --tw-border-opacity: 1;
            --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
            --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
            --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
            --tw-ring-opacity: .5;
        }
        #advanced-btn {
            font-size: .7rem !important;
            line-height: 19px;
            margin-top: 12px;
            margin-bottom: 12px;
            padding: 2px 8px;
            border-radius: 14px !important;
        }
        #advanced-options {
            display: none;
            margin-bottom: 20px;
        }
        .footer {
            margin-bottom: 45px;
            margin-top: 35px;
            text-align: center;
            border-bottom: 1px solid #e5e5e5;
        }
        .footer>p {
            font-size: .8rem;
            display: inline-block;
            padding: 0 10px;
            transform: translateY(10px);
            background: white;
        }
        .dark .footer {
            border-color: #303030;
        }
        .dark .footer>p {
            background: #0b0f19;
        }
"""
  
def plot_seir(population):
    simulation = create_and_run_exp(population)
    raw_result, tensor_fields, sessions = simulation.execute()
    result = pd.DataFrame(raw_result)
    pd.options.plotting.backend = "plotly"
    fig = result.plot(
    kind = "line",
    x = "timestep",
    y= ["S","E","I", "R"])
    fig.update_layout(title = "Population Evolution Over Time",
            xaxis_title="Time (Days)",
            yaxis_title="People")
    return fig

with gr.Blocks(css = css) as demo:
    gr.Markdown("""
                ## Epidemic Simulation
            """)
            
    gr.HTML('''
     <p style="margin-bottom: 10px; font-size: 94%">
                A cadCAD-based simulation of an outbreak of any epidemic using the SEIR compartmental model
              </p>
              ''')
    population_input = gr.Slider(1000, 100000, value=10000, label = "Population")
    simulate_btn = gr.Button('Run Simulation')
    graph = gr.Plot()
    simulate_btn.click(plot_seir, inputs = population_input, outputs = graph)

demo.launch()