import pints

"""
Explanation: Optimisation: First example
This example shows you how to run a global optimisation with Pints.
First, we import pints:
End of explanation
"""


import pints.toy as toy
model = toy.LogisticModel()

"""
Explanation: Next, we need a model: any class that implements the pints.ForwardModel interface.
Usually, you'd write a class for this purpose (that wrapped around whatever simulation package you wanted to use to generate your time series data). But you could also use a pure-Python model.
In the example, we use a logistic model, provided by Pints's toy model module.
End of explanation
"""


real_parameters = [0.015, 500]

"""
Explanation: This model has two parameters: A growth rate (which determines the steepness of the curve) and a carrying capacity (which determines the number the curve converges to).
For the example, we simply pick some nice values:
End of explanation
"""


import numpy as np
times = np.linspace(0, 1000, 100)

"""
Explanation: Finally, we create a list of times (in a real experiment, these would be the times at which the time series was sampled)
End of explanation
"""


values = model.simulate(real_parameters, times)

"""
Explanation: We now have everything we need to run a simulation and generate some toy data:
End of explanation
"""


import matplotlib.pyplot as plt

plt.figure(figsize=(14, 6))
plt.xlabel('Time')
plt.ylabel('Values')
plt.plot(times, values)
plt.show()

"""
Explanation: We can use Matplotlib (or any other plotting package) to have a look at the generated data:
End of explanation
"""


problem = pints.SingleOutputProblem(model, times, values)

"""
Explanation: We now set up an optimisation, to see if we can recover our original parameters from this data.
First, we define a problem (in this case a single-valued time series fitting problem):
End of explanation
"""


error = pints.SumOfSquaresError(problem)

"""
Explanation: We then define an error measure on this problem:
End of explanation
"""


boundaries = pints.RectangularBoundaries([0, 200], [1, 1000])

"""
Explanation: A lot of real problems have physical constraints on the values the parameters can take, so in this example we add them in the form of boundaries:
End of explanation
"""


x0 = np.array([0.5, 500])

"""
Explanation: Finally, we define an initial position to start searching at
End of explanation
"""


found_parameters, found_value = pints.optimise(
    error,
    x0,
    boundaries=boundaries,
    method=pints.XNES,
    )

"""
Explanation: We now run an optimisation, using the xNES method (although we could also have used a different global optimiser, like CMA-ES or PSO):
End of explanation
"""


values2 = model.simulate(found_parameters, times)

plt.figure(figsize=(14, 6))
plt.xlabel('Time')
plt.ylabel('Values')
plt.plot(times, values, 'x', label='Experimental data')
plt.plot(times, values2, label='Simulated data')
plt.legend()
plt.show()

"""
Explanation: We can now look at the quality of the fit by running a simulation with the obtained parameters and plotting it on top of the noisy data:
End of explanation
"""


print('Found solution:          True parameters:' )
for k, x in enumerate(found_parameters):
    print(pints.strfloat(x) + '    ' + pints.strfloat(real_parameters[k]))

"""
Explanation: Because this was a toy model, we can also compare the parameters with the known true parameters:
End of explanation
"""


times = np.linspace(0, 1000, 1000)

values = model.simulate(real_parameters, times)
values += np.random.normal(size=values.shape) * 10

plt.figure(figsize=(14, 6))
plt.xlabel('Time')
plt.ylabel('Values')
plt.plot(times, values)
plt.show()

"""
Explanation: Fitting to noisy data
We now repeat this exercise with noisy data:
End of explanation
"""


problem = pints.SingleOutputProblem(model, times, values)
error = pints.SumOfSquaresError(problem)

found_parameters, found_value = pints.optimise(
    error,
    x0,
    boundaries=boundaries,
    method=pints.XNES,
    )

"""
Explanation: Again, we set up a problem, define an error measure, and run the optimisation:
End of explanation
"""


values2 = model.simulate(found_parameters, times)

plt.figure(figsize=(14, 6))
plt.xlabel('Time')
plt.ylabel('Values')
plt.plot(times, values, label='Noisy experimental data')
plt.plot(times, values2, label='Simulated data')
plt.legend()
plt.show()

"""
Explanation: And we can inspect the result by plotting a simulation:
End of explanation
"""


print('Found solution:          True parameters:' )
for k, x in enumerate(found_parameters):
    print(pints.strfloat(x) + '    ' + pints.strfloat(real_parameters[k]))

"""
Explanation: Looks like a great fit!
Because it's still a toy problem, we can look at the parameter values as well:
End of explanation
"""


print('Error at obtained solution: ')
print(error(found_parameters))

print('Error at true solution: ')
print(error(real_parameters))

"""
Explanation: So it looks like we were close, but not exactly at the right parameters.
Why didn't we find exactly the right solution in the noisy case?
To investigate this, we start by looking at the error measure, evaluated at the true solution and at the obtained solution:
End of explanation
"""


import pints.plot
fig, axes = pints.plot.function_between_points(error, point_1=found_parameters, point_2=real_parameters)
axes.set_ylabel('score')
plt.show()

"""
Explanation: Another way to compare them is to plot the error function on a line drawn between these points:
End of explanation
"""


