import pybamm
import numpy as np
import matplotlib.pyplot as plt

model_scalar = pybamm.BaseModel()
c = pybamm.Variable("Concentration [mol.m-3]", domain="negative particle")
R = pybamm.Parameter("Particle radius [m]")
D = pybamm.Parameter("Diffusion coefficient [m2.s-1]")
j = pybamm.Parameter("Interfacial current density [A.m-2]")
F = pybamm.Parameter("Faraday constant [C.mol-1]")
c0 = pybamm.Parameter("Initial concentration [mol.m-3]")
N = -D * pybamm.grad(c)
dcdt = -pybamm.div(N)
model_scalar.rhs = {c: dcdt}
lbc = pybamm.Scalar(0)
rbc = -j / F / D
model_scalar.boundary_conditions = {c: {"left": (lbc, "Neumann"), "right": (rbc, "Neumann")}}
model_scalar.initial_conditions = {c: c0}
model_scalar.variables = {
    "Concentration [mol.m-3]": c,
    "Surface concentration [mol.m-3]": pybamm.surf(c),
    "Flux [mol.m-2.s-1]": N,
}
param_scalar = pybamm.ParameterValues(
    {
        "Particle radius [m]": 10e-6,
        "Diffusion coefficient [m2.s-1]": 3.9e-15,
        "Interfacial current density [A.m-2]": 1.4,
        "Faraday constant [C.mol-1]": 96485,
        "Initial concentration [mol.m-3]": 2.5e4,
    }
)
r = pybamm.SpatialVariable(
    "r", domain=["negative particle"], coord_sys="spherical polar"
)
geometry_scalar = {"negative particle": {r: {"min": pybamm.Scalar(0), "max": R}}}
param_scalar.process_model(model_scalar)
param_scalar.process_geometry(geometry_scalar)
submesh_types = {"negative particle": pybamm.Uniform1DSubMesh}
var_pts = {r: 20}
mesh_scalar = pybamm.Mesh(geometry_scalar, submesh_types, var_pts)
spatial_methods = {"negative particle": pybamm.FiniteVolume()}
disc_scalar = pybamm.Discretisation(mesh_scalar, spatial_methods)
disc_scalar.process_model(model_scalar);
model_scalar.rhs[c].render()

model_LUT = pybamm.BaseModel()
inputs = {"Concentration [mol.m-3]": c}
D = pybamm.FunctionParameter("Diffusion coefficient [m2.s-1]", inputs)
N = -D * pybamm.grad(c)
dcdt = -pybamm.div(N)
model_LUT.rhs = {c: dcdt}
rbc = -j / F / pybamm.surf(D)
model_LUT.boundary_conditions = {c: {"left": (lbc, "Neumann"), "right": (rbc, "Neumann")}}
model_LUT.initial_conditions = {c: c0}
model_LUT.variables = {
    "Concentration [mol.m-3]": c,
    "Surface concentration [mol.m-3]": pybamm.surf(c),
    "Flux [mol.m-2.s-1]": N,
}
param_LUT = pybamm.ParameterValues(
    {
        "Particle radius [m]": param_scalar["Particle radius [m]"],
        "Diffusion coefficient [m2.s-1]": pybamm.Interpolant(np.array([0,2*param_scalar["Initial concentration [mol.m-3]"]]),np.array([param_scalar["Diffusion coefficient [m2.s-1]"],param_scalar["Diffusion coefficient [m2.s-1]"]]),c),
        "Interfacial current density [A.m-2]": param_scalar["Interfacial current density [A.m-2]"],
        "Faraday constant [C.mol-1]": param_scalar["Faraday constant [C.mol-1]"],
        "Initial concentration [mol.m-3]": param_scalar["Initial concentration [mol.m-3]"],
    }
)
geometry_LUT = {"negative particle": {r: {"min": pybamm.Scalar(0), "max": R}}}
param_LUT.process_model(model_LUT)
param_LUT.process_geometry(geometry_LUT)
mesh_LUT = pybamm.Mesh(geometry_LUT, submesh_types, var_pts)
disc_LUT = pybamm.Discretisation(mesh_LUT, spatial_methods)
disc_LUT.process_model(model_LUT);
# Uncomment the following line to mitigate simulation inaccuracy of model that uses diffusion coefficient defined by interpolant
# model_LUT.rhs[c].children[1].children[1].children[0].children[1].children[1].children[1].children[0]=pybamm.Scalar(1e-20)
model_LUT.rhs[c].render()

t = [0, 3600]
models=[model_scalar,model_LUT]
solutions = []
for model in models:
    solver = pybamm.ScipySolver()
    solutions.append(solver.solve(model, t))

time=solutions[0].t
c_surf_scalar = solutions[0]["Surface concentration [mol.m-3]"](time)
c_surf_LUT = solutions[1]["Surface concentration [mol.m-3]"](time)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))
ax1.plot(time, c_surf_scalar,color="tab:green")
ax1.plot(time, c_surf_LUT,color="tab:red")
ax1.set_xlabel("Time [s]")
ax1.set_ylabel("Surface concentration [mol.m-3]")
ax1.legend(['scalar', 'LUT'])
ax2.plot(time, c_surf_LUT-c_surf_scalar,color="tab:blue")
ax2.set_xlabel("Time [s]")
ax2.set_ylabel("Error of surface concentration [mol.m-3]",color="tab:blue")
ax2_right = ax2.twinx()
ax2_right.plot(time, (c_surf_LUT-c_surf_scalar)/c_surf_scalar*100,color="tab:orange")
ax2_right.set_ylabel("Relative error of surface concentration [%]",color="tab:orange")
fig.tight_layout()
plt.show()