# Declares the model.
using Catalyst
rn = @reaction_network begin
    @parameters σ
    k1, X1 --> X2
    k2, X2 --> X3
    k3, X3 --> 0
end
u0 = [:X1 => 5.0, :X2 => 0.0, :X3 => 0.0]
tend = 1.0
ps = [:k1 => 100.0, :k2 => 20.0, :k3 => 3.0, :σ => 0.1]

# Generate synthetic data.
using Distributions, OrdinaryDiffEq, Plots
oprob = ODEProblem(rn, u0, tend, ps)
sol = solve(oprob)
t_data = range(0, tend; length = 20)
data_X1 = [rand(Normal(X1, max(0,X1)/10.0)) for X1 in sol(t_data; idxs = :X1)]
data_X2 = [rand(Normal(X2, max(0,X2)/10.0)) for X2 in sol(t_data; idxs = :X2)]
data_X3 = [rand(Normal(X3, max(0,X3)/10.0)) for X3 in sol(t_data; idxs = :X3)]

# Plots true solution and data.
plot(sol; label = "")
plot!(t_data, data_X1; seriestype = :scatter, alpha = 0.5, color = 1, label = "")
plot!(t_data, data_X2; seriestype = :scatter, alpha = 0.5, color = 2, label = "")
plot!(t_data, data_X3; seriestype = :scatter, alpha = 0.5, color = 3, label = "")

# Performs parameter fitting using PEtab.
# For observables I have tried with and without the `abs`.
using PEtab, DataFrames, Optim
observables = Dict(
    "obs_X1" => PEtabObservable(:X1, rn.σ * abs(rn.X1)),
    "obs_X2" => PEtabObservable(:X2, rn.σ * abs(rn.X2)),
    "obs_X3" => PEtabObservable(:X3, rn.σ * abs(rn.X3))
)
p_est = [PEtabParameter(ModelingToolkit.getname(p)) for p in parameters(rn)]
measurements = vcat(
    DataFrame(obs_id = "obs_X1", time = t_data, measurement = data_X1),
    DataFrame(obs_id = "obs_X2", time = t_data, measurement = data_X2),
    DataFrame(obs_id = "obs_X3", time = t_data, measurement = data_X3)
)
petab_model = PEtabModel(rn, observables, measurements, p_est; speciemap = u0)
petab_prob = PEtabODEProblem(petab_model)
petab_fit = calibrate_multistart(petab_prob, IPNewton(), 100)
plot(petab_fit, petab_prob)

# Additional evaluation of PEtab problem.
petab_fit.fmin # NaN
petab_prob.nllh(log10.(last.(ps))) # NaN

# Parameter fitting using Optimization.
# Naive approach using just squared error.
using Optimization, OptimizationBBO
p0 = [rand() for _ in parameters(rn)]
oprob_base = ODEProblem(rn, u0, 1.0, Pair.(parameters(rn), p0))
set_p = ModelingToolkit.setp_oop(oprob_base, parameters(rn))
function loss(p, (oprob_base, set_p, t_data, data_X1, data_X2, data_X3))
    oprob = remake(oprob_base; p = set_p(oprob_base, p))
    sol = solve(oprob; verbose = false, maxiters = 10000, saveat = t_data)
    SciMLBase.successful_retcode(sol) || return Inf
    return sum(abs2, data_X1 .- sol[:X1]) +
           sum(abs2, data_X2 .- sol[:X2]) +
           sum(abs2, data_X3 .- sol[:X3])

end
opt_prob = OptimizationProblem(loss, p0, (oprob_base, set_p, t_data, data_X1, data_X2, data_X3);
    lb = fill(1e-3, length(p0)), ub = fill(1e3, length(p0)))
opt_sol = solve(opt_prob, BBO_adaptive_de_rand_1_bin_radiuslimited())

# Plots optimisation solution.
oprob_fitted = remake(oprob_base; p = set_p(oprob_base, opt_sol.u))
sol = solve(oprob_fitted)
plot(sol; label = "")
plot!(t_data, data_X1; seriestype = :scatter, alpha = 0.5, color = 1, label = "")
plot!(t_data, data_X2; seriestype = :scatter, alpha = 0.5, color = 2, label = "")
plot!(t_data, data_X3; seriestype = :scatter, alpha = 0.5, color = 3, label = "")