#!/usr/bin/env python
# coding: utf-8

# ## Parameter identification for various models
# 
# To investigate the performance of parameter identification for different electrochemical models we will start with synthetic data from the highest order model in PyBOP (Many-particle DFN) and try to identify the correct parameter values on the reduced order models.
# 
# ### Setting up the Environment
# 
# Before we begin, we need to ensure that we have all the necessary tools. We will install PyBOP and upgrade dependencies:

# In[ ]:


get_ipython().run_line_magic('pip', 'install --upgrade pip ipywidgets pybamm -q')
get_ipython().run_line_magic('pip', 'install pybop -q')


# Next, we import the added packages plus any additional dependencies,

# In[ ]:


import numpy as np
import pybamm

import pybop

go = pybop.PlotlyManager().go
pybop.PlotlyManager().pio.renderers.default = "notebook_connected"


# Let's fix the random seed in order to generate consistent output during development, although this does not need to be done in practice.

# In[ ]:


np.random.seed(8)


# ## Optimising the Parameters

# First, we define the model to be used for the parameter optimisation,

# In[ ]:


parameter_set = pybop.ParameterSet.pybamm("Chen2020")
parameter_set = pybamm.get_size_distribution_parameters(parameter_set)
synth_model = pybop.lithium_ion.DFN(
    parameter_set=parameter_set, options={"particle size": "distribution"}
)


# ### Simulating Forward Model
# 
# We can then simulate the model using the `predict` method, with a default constant current to generate voltage data.

# In[ ]:


n_points = 450
t_eval = np.linspace(0, 1600 + 1000, n_points)
current = np.concatenate(
    [np.ones(200) * parameter_set["Nominal cell capacity [A.h]"], np.zeros(250)]
)
initial_state = {"Initial SoC": 0.5}


# In[ ]:


dataset = pybop.Dataset(
    {
        "Time [s]": t_eval,
        "Current function [A]": current,
    }
)


# In[ ]:


synth_model.build(dataset=dataset, initial_state=initial_state)
synth_model.signal = ["Voltage [V]"]
values = synth_model.simulate(t_eval=t_eval, inputs={})


# ### Adding Noise to Voltage Data
# 
# To make the parameter estimation more realistic, we add Gaussian noise to the data.

# In[ ]:


sigma = 0.001
corrupt_values = values["Voltage [V]"].data + np.random.normal(
    0, sigma, len(values["Voltage [V]"].data)
)
go.Figure(
    data=go.Scatter(x=t_eval, y=corrupt_values, mode="lines"),
    layout=go.Layout(title="Corrupted Voltage", width=800, height=600),
)


# ## Identifying the Parameters

# We will now set up the parameter estimation process by defining the datasets for optimisation and selecting the model parameters we wish to estimate.

# ### Creating a Dataset
# 
# The dataset for optimisation is composed of time, current, and the noisy voltage data:

# In[ ]:


dataset = pybop.Dataset(
    {
        "Time [s]": t_eval,
        "Current function [A]": current,
        "Voltage [V]": corrupt_values,
    }
)


# Next, we define the model parameters for optimisation. Furthermore, PyBOP provides functionality to define a prior for the parameters. The initial parameter values used in the optimisation will be randomly drawn from the prior distribution.

# In[ ]:


parameters = pybop.Parameters(
    pybop.Parameter(
        "Positive electrode thickness [m]",
        prior=pybop.Gaussian(7.56e-05, 0.05e-05),
        bounds=[65e-06, 85e-06],
        true_value=parameter_set["Positive electrode thickness [m]"],
    ),
    pybop.Parameter(
        "Negative electrode thickness [m]",
        prior=pybop.Gaussian(8.52e-05, 0.05e-05),
        bounds=[75e-06, 95e-06],
        true_value=parameter_set["Negative electrode thickness [m]"],
    ),
)


# We can now define the output signal, the problem (which combines the model with the dataset) and construct a cost function which in this example is the `GravimetricEnergyDensity()` used to maximise the gravimetric energy density of the cell.

# In[ ]:


models = [
    pybop.lithium_ion.SPM(parameter_set=parameter_set),
    pybop.lithium_ion.SPMe(parameter_set=parameter_set),
]


# Let's construct PyBOP's optimisation class for each model. This class provides the methods needed to fit the forward model. For this example, we use an evolution strategy (XNES) as the optimiser.

# In[ ]:


optims = []
xs = []
for model in models:
    print(f"Running {model.name}")
    model.set_initial_state(initial_state)
    problem = pybop.FittingProblem(model, parameters, dataset)
    cost = pybop.SumSquaredError(problem)
    optim = pybop.XNES(
        cost, verbose=True, max_iterations=60, max_unchanged_iterations=15
    )
    x, final_cost = optim.run()
    optims.append(optim)
    xs.append(x)


# In[ ]:


for optim, x in zip(optims, xs):
    print(f"| Model: {optim.cost.problem.model.name} | Results: {x} |")


# ## Plotting and Visualisation
# 
# PyBOP provides various plotting utilities to visualise the results of the optimisation.

# ### Comparing System Response
# 
# We can quickly plot the system's response using the estimated parameters compared to the initial parameters:
# 

# In[ ]:


for optim, x in zip(optims, xs):
    pybop.quick_plot(
        optim.cost.problem, problem_inputs=x, title=optim.cost.problem.model.name
    )


# ### Cost Landscape
# 
# Finally, we can visualise the cost landscape and the path taken by the optimiser:
# 

# In[ ]:


for optim in optims:
    pybop.plot2d(optim, steps=10, title=optim.cost.problem.model.name)