r"""Solves a simple ODE using a PINN.

.. math::

    \frac{du}{dt} = \mu u, \quad t \in (0, T), \quad u(0) = 1

where :math:`u: (0, T) \times (\mu_{\min}, \mu_{\max}) \to \mathbb{R}` is the unknown function, with :math:`(0, T) \subset \mathbb{R}` the time domain and :math:`(\mu_{\min}, \mu_{\max}) \subset \mathbb{R}` the parameter domain.

The exact solution is

.. math::

    u(t, \mu) = e^{\mu t}.

Three training strategies are compared: ssBFGS with weak IC, energy natural gradient with weak IC, and energy natural gradient with strong IC.
"""

# %%

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.nn_space import NNtSpace
from scimba_torch.integration.monte_carlo import TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.integration.monte_carlo_time import UniformTimeSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import (
    GenericMLP,
)
from scimba_torch.numerical_solvers.ode.pinns import (
    NaturalGradientODEPinns,
    ODEPinns,
)
from scimba_torch.physical_models.ode.simple_ode import SimpleODE
from scimba_torch.utils.scimba_tensors import LabelTensor

# %%


def f_ini(mu):
    t = LabelTensor(torch.zeros_like(mu.x))
    return exact(t, mu)


def exact(t, mu):
    t_ = t.get_components()
    mu_ = mu.get_components()
    return torch.exp(mu_ * t_)


t_min, t_max = 0.0, 1.0
mu_min, mu_max = 1.0, 2.0

sampler = TensorizedSampler(
    [
        UniformTimeSampler((t_min, t_max)),
        UniformParametricSampler([(mu_min, mu_max)]),
    ]
)


def post_processing(
    inputs: torch.Tensor, t: LabelTensor, mu: LabelTensor
) -> torch.Tensor:
    t_ = t.get_components()
    return 1 + t_ * inputs


def functional_post_processing(
    func, t: torch.Tensor, mu: torch.Tensor, theta: torch.Tensor
) -> torch.Tensor:
    return 1 + t[0] * func(t, mu, theta)


def create_ode():
    space = NNtSpace(1, 1, GenericMLP, sampler, layer_sizes=[12, 12])
    return SimpleODE(space, init=f_ini)


def create_ode_pp():
    space = NNtSpace(
        1, 1, GenericMLP, sampler, layer_sizes=[12, 12], post_processing=post_processing
    )
    return SimpleODE(space, init=f_ini)


def plot_pinn(pinn, title: str, mu_val: float = (mu_min + mu_max) / 2):
    """Plot the results of the PINN."""
    with torch.no_grad():
        fig, ax = plt.subplots(1, 3, figsize=(15, 5))
        t = LabelTensor(torch.linspace(t_min, t_max, 1000)[:, None])
        mu = LabelTensor(mu_val * torch.ones_like(t.x))
        u = pinn.evaluate(t, mu).w
        pinn.losses.plot(ax[0])
        ax[1].plot(t.x, u, label="approximation")
        ax[1].plot(t.x, torch.exp(mu.x * t.x), label="exact")
        ax[1].set_title(f"{title}")
        ax[1].legend()
        error = u - torch.exp(mu.x * t.x)
        ax[2].plot(t.x, error, label="error")
        ax[2].set_title(f"Error, L2 = {torch.sqrt(torch.mean(error**2)):.2e}")
        ax[2].legend()
        plt.show()


# %%

pinn_ssbfgs = ODEPinns(create_ode(), ic_type="weak", optimizers="ssbfgs", ic_weight=100)

resume_solve = True
if resume_solve or not pinn_ssbfgs.load(__file__, "simple_ode_ssbfgs"):
    pinn_ssbfgs.solve(epochs=200, n_collocation=2000, n_ic_collocation=500)
    pinn_ssbfgs.save(__file__, "simple_ode_ssbfgs")

plot_pinn(pinn_ssbfgs, "ssBFGS with weak IC")

# %%

pinn_eng_weak = NaturalGradientODEPinns(
    create_ode(), ic_type="weak", ic_weight=100, matrix_regularization=1e-5
)

resume_solve = True
if resume_solve or not pinn_eng_weak.load(__file__, "simple_ode_eng"):
    pinn_eng_weak.solve(epochs=200, n_collocation=2000, n_ic_collocation=500)
    pinn_eng_weak.save(__file__, "simple_ode_eng")

plot_pinn(pinn_eng_weak, "ENG with weak IC")

# %%

pinn_eng_strong = NaturalGradientODEPinns(
    create_ode_pp(), functional_post_processing=functional_post_processing
)

resume_solve = True
if resume_solve or not pinn_eng_strong.load(__file__, "simple_ode_eng_strong"):
    pinn_eng_strong.solve(epochs=100, n_collocation=1000)
    pinn_eng_strong.save(__file__, "simple_ode_eng_strong")

plot_pinn(pinn_eng_strong, "ENG with strong IC")

# %%