import pybop
import numpy as np

# Define model
parameter_set = pybop.ParameterSet.pybamm("Chen2020")
model = pybop.lithium_ion.SPM(parameter_set=parameter_set)

# Fitting parameters
parameters = [
    pybop.Parameter(
        "Negative particle radius [m]",
        prior=pybop.Gaussian(6e-06, 0.1e-6),
        bounds=[1e-6, 9e-6],
        true_value=parameter_set["Negative particle radius [m]"],
    ),
    pybop.Parameter(
        "Positive particle radius [m]",
        prior=pybop.Gaussian(4.5e-06, 0.1e-6),
        bounds=[1e-6, 9e-6],
        true_value=parameter_set["Positive particle radius [m]"],
    ),
]

# Generate data
sigma = 0.001
init_soc = 1  # start from full charge
experiment = pybop.Experiment(
    [
        "Discharge at 3A until 3.8 V (2 seconds period)",
    ],
)

values = model.predict(experiment=experiment, init_soc=init_soc)
corrupt_values = values["Voltage [V]"].data + np.random.normal(
    0, sigma, len(values["Time [s]"].data)
)
# Form dataset
dataset = pybop.Dataset(
    {
        "Time [s]": values["Time [s]"].data,
        "Current function [A]": values["Current [A]"].data,
        "Voltage [V]": corrupt_values,
    }
)

# Generate problem, cost function, and optimisation class
problem = pybop.FittingProblem(model, parameters, dataset, init_soc=init_soc)
cost = pybop.SumSquaredError(problem)
optim = pybop.Optimisation(cost, optimiser=pybop.CMAES)
optim.set_max_iterations(100)

# Run the optimisation
x, final_cost = optim.run()
print(
    "True parameters:",
    [
        parameters[0].true_value,
        parameters[1].true_value,
    ],
)
print("Estimated parameters:", x)

# Plot the timeseries output
pybop.quick_plot(x, cost, title="Optimised Comparison")

# Plot convergence
pybop.plot_convergence(optim)

# Plot the parameter traces
pybop.plot_parameters(optim)

# Plot the cost landscape
pybop.plot_cost2d(cost, steps=15)

# Plot the cost landscape with optimisation path and updated bounds
pybop.plot_cost2d(cost, optim=optim, steps=15)