import numpy as np
import cantera as ct
from assimulo.solvers import CVode
from assimulo.problem import Explicit_Problem
from assimulo.exception import TerminateSimulation

gasconstant = 8314.46261815324

# Defines the rhs
def f(t, y):
    gas.set_unnormalized_mass_fractions(y[1:])
    print(y[0])
    gas.TD = y[0], rho1
    wdot = gas.net_production_rates
    dYdt = wdot * Wk / rho1
    dTdt = - (np.dot(gas.partial_molar_enthalpies - np.ones(gas.n_species) * gasconstant * y[0], wdot) /
              (rho1 * gas.cv))
    yd_v = np.zeros(nsp)
    yd_v[0] = dTdt
    yd_v[1:] =dYdt
    return yd_v

# Defines the Jacobian*vector product
def jacv(t, y, fy, v):
    Jac = np.empty([nsp, nsp])
    dy_perturbed_neg = fy#f(t, y)
    eps_i = np.sqrt(np.finfo(float).eps)
    for i in range(nsp):
        #eps_i = np.sqrt(np.finfo(float).eps)
        y_perturbed = np.copy(y)
        y_perturbed[i] = y_perturbed[i] + eps_i
        dy_perturbed_pos = f(t, y_perturbed)
        Jac[:, i] = (dy_perturbed_pos - dy_perturbed_neg) / (1 * eps_i)
    return np.dot(Jac, v)

def state_events(t, y, sw):
    return np.array([y0[0] + 400 - y[0]])

def handle_event(solver, event_info):
    state_info = event_info[0]  # We are only interested in state events info
    if state_info[0] != 0:      # Check if the first event function have been triggered
        raise TerminateSimulation

mech='gri30.yaml'
gas = ct.Solution(mech)
gas.TP = 900,50*1e5
gas.set_equivalence_ratio(1, 'CH4:1', 'O2:1,N2:3.76')

Wk = gas.molecular_weights
rho1 = gas.density
y0 = np.hstack((gas.T, gas.Y))
nsp =len(y0)

exp_mod = Explicit_Problem(f, y0, name='Example using the Jacobian Vector product')
exp_mod.state_events = state_events  # Sets the state events to the problem
exp_mod.handle_event = handle_event  # Sets the event handling to the problem
exp_mod.jacv = jacv  # Sets the Jacobian
exp_sim = CVode(exp_mod)  # Create a CVode solver

# Set the parameters
exp_sim.iter = 'Newton'  # Default 'FixedPoint'
exp_sim.discr = 'BDF'  # Default 'Adams'
exp_sim.rtol = 1e-4  # Default 1e-6
exp_sim.atol = 1e-10  # Default 1e-6
exp_sim.linear_solver = 'spgmr'  # Change linear solver
#exp_sim.inith=1e-10
exp_sim.maxsteps=1000
# Simulate
t, y = exp_sim.simulate(1, 1000)  # Simulate 5 seconds with 1000 communication points

# Basic tests

import matplotlib
matplotlib.use('Qt5Agg')

import matplotlib.pyplot as plt
plt.plot(t,y[:,0],linestyle="dashed",marker="o") #Plot the solution
plt.xlim(0, 0.05)
plt.ylim(800,1400)
plt.show()