# General Purpose
import numpy as np
from matplotlib import pyplot as plt
from scipy.integrate import odeint
from scipy.optimize import fsolve

# Jupyter Specifics
from IPython.display import HTML, display
from ipywidgets.widgets import FloatSlider, Layout, HBox, VBox, interactive_output, Button, Label

%matplotlib inline

style = {'description_width': '50px'}
slider_layout = Layout(width='90%')

# -*- coding: utf-8 -*-
#quelques couleurs
rouge_A,vert1_A,vert2_A,vert3_A,vert4_A,gris1_A='#C60C2E','#005157','#627D77','#9EB28F','#C5E5A4','#595A5C'
coule=[rouge_A,vert1_A,vert2_A,vert3_A,vert4_A,gris1_A]

time_meas=np.loadtxt("time")
traffic=np.loadtxt("traffic")
PM10=np.loadtxt("PM10")
PM10=PM10*100

box_layout = Layout(display='flex',
                    flex_flow='row',
                    align_items='center',
                    width='100%')

box_layout_v = Layout(display='flex',
                    flex_flow='column',
                    align_items='flex-start',
                    width='30%')

style = {'description_width': '15px'}

slayout = Layout(width='300px')


def fc_crank_nicolson(Cp,C,dt,alpha,beta,tau,delta,N,Ce):
    term_C=dt*(alpha*N**2+(beta*N+tau)*(Ce-C)-delta*C)
    term_Cp=dt*(alpha*N**2+(beta*N+tau)*(Ce-Cp)-delta*Cp)
    return -Cp + C + 0.5*term_C + 0.5*term_Cp

def main(alpha1,beta,delta,tau):
    # Crank-Nicolson scheme for air polluton equation in the underground
 
    dt=0.25 # 1/h
    sim_time=168

    # alpha=0.19
    # beta=0.15
    # delta=0.15#	1/h
    # tau=1# vol/h

    fig, ax = plt.subplots(figsize=(15, 10))
    
    Ce=15 # constant outdoor concentration
    C=PM10[0] # initial concentration
    t,i=min(time_meas),0 
    time,concentration=[],[]
    while t <= max(time_meas):
        time.append(t)
        concentration.append(C)
        N=traffic[i]
        # solve for C+, concentration at the next time step
        Cp=fsolve(fc_crank_nicolson, C, args=(C,dt,alpha1,beta,tau,delta,N,Ce))
        C=Cp
        t+=dt
        i+=1

    plt.plot(time,concentration,color=coule[-1],linestyle="--",alpha=0.5,marker='',label='model')
    plt.plot(time,PM10,color=coule[1],linestyle="",alpha=0.25,marker='o',markersize='2.5',label='scaled measurements')
    plt.xlabel(r"Time [h]",fontsize=14,fontweight='bold')
    plt.ylabel(r"Underground $PM_{10}$ [µg/m$^3$]",fontsize=14,fontweight='bold')
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)
    plt.legend(fontsize=14)
    plt.show()

alpha1=FloatSlider(min=0, max=1, step=0.1, value=0.19, description=r'$\alpha$', style=style, layout=slider_layout)
beta=FloatSlider(min=0, max=1, step=0.1, value=0.15, description=r'$\beta$', style=style, layout=slider_layout)
delta=FloatSlider(min=0, max=1, step=0.1, value=0.15, description=r'$\delta$', style=style, layout=slider_layout)
tau=FloatSlider(min=0, max=10, step=0.5, value=1, description=r'$\tau$', style=style, layout=slider_layout)

out=interactive_output(main,{'alpha1': alpha1, 'beta': beta, 'delta': delta,'tau':tau});

button = Button(description='Reset')

def on_button_clicked(b):
        alpha1.value, beta.value, delta.value, tau.value=0.19, 0.15, 0.15,1
            
button.on_click(on_button_clicked)

ui=VBox([Label('Parameters:'),alpha1, beta, delta,tau, button], layout=box_layout_v)
HBox([ui, out])