import os
import numpy as np
import readdy

# set up the system

system = readdy.ReactionDiffusionSystem(
    box_size=[4.4964, 4.4964, 4.4964],
    unit_system={"length_unit": "micrometer", "time_unit": "second"}
)

system.add_species("E", diffusion_constant=10.)
system.add_species("S", diffusion_constant=10.)
system.add_species("ES", diffusion_constant=10.)
system.add_species("P", diffusion_constant=10.)

# mutual interaction radius is 0.03
# rate constant is lambda from Erban-Chapman, 2009.
# solve it using k_macro = 4*pi*D*R*(1-tanh(kappa*R)/(kappa*R)) and kappa = sqrt(lambda/D)
# where D is mutual diffusion coefficient and R is mutual interaction radius
# and then numerically solve for lambda

system.reactions.add("fwd: E +(0.03) S -> ES", rate=86.78638438)
system.reactions.add("back: ES -> E +(0.03) S", rate=1.)
system.reactions.add("prod: ES -> E +(0.03) P", rate=1.)

# simulate the system

simulation = system.simulation(kernel="SingleCPU")
simulation.output_file = "out.h5"
simulation.reaction_handler = "UncontrolledApproximation"

n_particles_e = 909
n_particles_s = 9091
edge_length = system.box_size[0]
initial_positions_e = np.random.random(size=(n_particles_e, 3)) * edge_length - .5*edge_length
initial_positions_s = np.random.random(size=(n_particles_s, 3)) * edge_length - .5*edge_length
simulation.add_particles("E", initial_positions_e)
simulation.add_particles("S", initial_positions_s)

simulation.observe.number_of_particles(stride=1, types=["E", "S", "ES", "P"])

if os.path.exists(simulation.output_file):
    os.remove(simulation.output_file)

dt = 1e-3
n_steps = int(10. / dt)

simulation.run(n_steps=n_steps, timestep=dt)

# obtain simulation output

traj = readdy.Trajectory(simulation.output_file)
time, counts = traj.read_observable_number_of_particles()

np.savetxt('ReaDDyBenchmark_out.txt', counts)

concentrations = counts / system.box_volume.magnitude
time = time * dt

from scipy.integrate import odeint

def f(x, t0, kf, kr, kcat):
    """
    x: state vector with concentrations of E, S, ES, P
    """
    return np.array([
        -kf * x[0] * x[1] + (kr+kcat)*x[2],
        -kf * x[0] * x[1] + kr * x[2],
        kf * x[0] * x[1]- (kr+kcat)*x[2],
        kcat*x[2]
    ])

init_state = np.array([n_particles_e, n_particles_s, 0., 0.]) / system.box_volume.magnitude
ode_time = np.linspace(0.,10,100000)
ode_result = odeint(f, y0=init_state, t=ode_time, args=(0.98e-2, 1., 1.))

perf = simulation._simulation.performance_root()
print(perf)

import matplotlib.pyplot as plt

stride = 100
plt.plot(time[::stride], concentrations[:,0][::stride], "-", label='E')
plt.plot(time[::stride], concentrations[:,1][::stride], "-", label='S')
plt.plot(time[::stride], concentrations[:,2][::stride], "-", label='ES')
plt.plot(time[::stride], concentrations[:,3][::stride], "-", label='P')
plt.plot(ode_time, ode_result[:,0], "--", color="grey", label="LMA solution")
plt.plot(ode_time, ode_result[:,1], "--", color="grey")
plt.plot(ode_time, ode_result[:,2], "--", color="grey")
plt.plot(ode_time, ode_result[:,3], "--", color="grey")
plt.legend()
plt.xlabel(r"time in $\mathrm{s}$")
plt.ylabel(r"concentration in $\mathrm{\mu m}^{-3}$")
plt.show()