%pip install git+https://github.com/pybop-team/PyBOP.git@develop -q 
%pip install pybamm -q

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

synthetic_model = pybamm.lithium_ion.SPM()
params = synthetic_model.default_parameter_values

params.update(
    {
        "Negative electrode active material volume fraction": 0.52,
        "Positive electrode active material volume fraction": 0.63,
    }
)

experiment = pybamm.Experiment(
    [
        (
            "Discharge at 2C for 5 minutes (1 second period)",
            "Rest for 2 minutes (1 second period)",
            "Charge at 1C for 5 minutes (1 second period)",
            "Rest for 2 minutes (1 second period)",
        ),
    ]
    * 2
)
sim = pybamm.Simulation(synthetic_model, experiment=experiment, parameter_values=params)
synthetic_sol = sim.solve()

sim.plot()

corrupt_V = synthetic_sol["Terminal voltage [V]"].data
corrupt_V += np.random.normal(0,0.005,len(corrupt_V))

model = pybop.lithium_ion.SPM()
observations = [
            pybop.Observed("Time [s]", synthetic_sol["Time [s]"].data),
            pybop.Observed("Current function [A]", synthetic_sol["Current [A]"].data),
            pybop.Observed("Voltage [V]", corrupt_V),
        ]

fit_params = [
    pybop.Parameter(
        "Negative electrode active material volume fraction",
        prior=pybop.Gaussian(0.5, 0.02),
        bounds=[0.375, 0.625],
    ),
    pybop.Parameter(
        "Positive electrode active material volume fraction",
        prior=pybop.Gaussian(0.65, 0.02),
        bounds=[0.525, 0.75],
    ),
]

parameterisation = pybop.Parameterisation(
    model, observations=observations, fit_parameters=fit_params
)

results, last_optim, num_evals = parameterisation.rmse(
    signal="Voltage [V]", method="nlopt"
)

params.update(
        {"Negative electrode active material volume fraction": results[0], 
        "Positive electrode active material volume fraction": results[1]}
        )
optsol = sim.solve()["Terminal voltage [V]"].data

plt.plot(corrupt_V, label='Groundtruth')
plt.plot(optsol, label='Estimated')
plt.xlabel('Time (s)')
plt.ylabel('Voltage (V)')
plt.legend()