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


# Load the OCP charge data with decreasing stoichiometry
ocp_data = {
    "Stoichiometry": np.asarray([1, 0.8747, 0.84609, 0.80795, 0.76981, 0.73168]),
    "Voltage [V]": np.asarray([3.5, 3.5558, 3.5802, 3.622, 3.6626, 3.6972]),
}

for direction in ["decreasing", "increasing"]:

    if direction == "decreasing":
        def ocp(sto):
            return pybamm.Interpolant(ocp_data["Stoichiometry"], ocp_data["Voltage [V]"], sto)
    else:
        def ocp(sto):
            return pybamm.Interpolant(np.flip(ocp_data["Stoichiometry"]), np.flip(ocp_data["Voltage [V]"]), sto)

    # Define parameters
    parameter_set = {
        'Nominal cell capacity [A.h]': 0.003,
        'Current function [A]': 0.003,
        'Initial stoichiometry': 0.874707,
        'Electrode OCP [V]': ocp,
        'Theoretical electrode capacity [A.s]': 17.621737,
        'Particle diffusion time scale [s]': 35344.0,
        'Series resistance [Ohm]': 1,
    }

    # Define a charge pulse and relaxation
    time_span = np.linspace(0, 7200, 121)  # 2 hours, 1 point per minute
    current = 0 * time_span
    current[11:20] = -0.00035  # 10 minute charging pulse
    parameter_set["Current function [A]"] = pybamm.Interpolant(time_span, current, pybamm.t)

    # Set up model
    model = pybop.lithium_ion.SPDiffusion(parameter_set=parameter_set, electrode="positive")

    # Set up and run simulation
    sim = pybamm.Simulation(
        model=model.pybamm_model,
        parameter_values=model._unprocessed_parameter_set,
        solver=pybamm.CasadiSolver(),
    )
    sol = sim.solve(t_eval=time_span)
    sol.plot(
        [
            {"Particle stoichiometry", "Particle surface stoichiometry"},
            "Current [A]",
            {"Voltage [V]", "Open-circuit voltage [V]"},
        ],
    )
    plt.show()