r"""Solves the Navier-Stokes equations in 2D using a PINN.

The setting is described in https://arxiv.org/pdf/2402.03153

"""

from pathlib import Path
from typing import OrderedDict

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.abstract_space import AbstractApproxSpace
from scimba_torch.approximation_space.nn_space import NNxtSpace
from scimba_torch.domain.meshless_domain.domain_2d import Disk2D, Square2D
from scimba_torch.integration.monte_carlo import DomainSampler, 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.temporal_pde.pinns import (
    NaturalGradientTemporalPinns,
)
from scimba_torch.physical_models.temporal_pde.abstract_temporal_pde import (
    TemporalPDE,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces

# from scimba_torch.utils.scimba_tensors import LabelTensor

torch.manual_seed(0)


class NavierStokes2DStrongForm(TemporalPDE):
    def __init__(self, space: AbstractApproxSpace, **kwargs):
        super().__init__(
            space, residual_size=3, bc_residual_size=10, ic_residual_size=3, **kwargs
        )

        self.residual_size = 3
        self.bc_residual_size = 10
        self.ic_residual_size = 3
        spatial_dim = self.space.spatial_domain.dim
        assert spatial_dim == 2, (
            "This class is designed for 2D Navier-Stokes equations."
        )

        # self.fixed_params = fixed_params
        # self.nu = fixed_params["nu"]  # Viscosity

    def init(self, x, mu):
        # Zero vector for the residuals
        x1, _ = x.get_components()
        zeros = torch.zeros_like(x1)
        return zeros, zeros, zeros

    def rhs(self, w, t, x, mu):
        # Zero vector for the residuals
        x1, _ = x.get_components()
        zeros = torch.zeros_like(x1)
        return zeros, zeros, zeros

    def time_operator(self, w, t, x, mu):
        u, v, p = w.get_components()
        # mu_, k_f = mu.get_components()

        u_t = self.grad(u, t)
        v_t = self.grad(v, t)

        x1, _ = x.get_components()
        zeros = torch.zeros_like(x1)

        return u_t, v_t, zeros

    def space_operator(self, w, t, x, mu):
        u, v, p = w.get_components()
        mu_ = mu.get_components()

        # Incompressibility condition
        u_x, u_y = self.grad(u, x)
        v_x, v_y = self.grad(v, x)
        op1 = u_x + v_y

        # Navier-Stokes equations
        u_xx, _ = self.grad(u_x, x)
        _, u_yy = self.grad(u_y, x)

        v_xx, _ = self.grad(v_x, x)
        _, v_yy = self.grad(v_y, x)

        p_x, p_y = self.grad(p, x)

        op2a = (u * u_x + v * u_y) + p_x - mu_ * (u_xx + u_yy)
        op2b = (u * v_x + v * v_y) + p_y - mu_ * (v_xx + v_yy)

        return op2a, op2b, op1

    def bc_rhs(self, w, t, x, n, mu):
        x1_0, _ = x.get_components(label=0)
        x1_1, _ = x.get_components(label=1)
        x1_2, _ = x.get_components(label=2)
        x1_3, _ = x.get_components(label=3)
        x1_4, _ = x.get_components(label=4)

        return (
            torch.zeros_like(x1_0),
            torch.zeros_like(x1_0),
            torch.zeros_like(x1_1),
            torch.zeros_like(x1_1),
            torch.zeros_like(x1_2),
            torch.zeros_like(x1_2),
            torch.ones_like(x1_3),
            torch.zeros_like(x1_3),
            torch.zeros_like(x1_4),
            torch.zeros_like(x1_4),
        )

    def bc_operator(self, w, t, x, n, mu):
        u, v, p = w.get_components()
        ldown, lright, ltop, lleft, circle = 0, 1, 2, 3, 4

        # Dirichlet conditions
        u_down = w.restrict_to_labels(u, labels=[ldown])
        v_down = w.restrict_to_labels(v, labels=[ldown])
        u_right = w.restrict_to_labels(u, labels=[lright])
        p_right = w.restrict_to_labels(p, labels=[lright])
        u_up = w.restrict_to_labels(u, labels=[ltop])
        v_up = w.restrict_to_labels(v, labels=[ltop])
        u_left = w.restrict_to_labels(u, labels=[lleft])
        v_left = w.restrict_to_labels(v, labels=[lleft])

        u_circle = w.restrict_to_labels(u, labels=[circle])
        v_circle = w.restrict_to_labels(v, labels=[circle])

        return (
            u_down,
            v_down,
            torch.zeros_like(u_right),
            p_right,
            u_up,
            v_up,
            u_left,
            v_left,
            u_circle,
            v_circle,
        )

    def initial_condition(self, x, mu):
        x1, _ = x.get_components()
        zeros = torch.zeros_like(x1)
        return zeros, zeros, zeros

    def functional_time_operator(self, func, t, x, mu, theta):
        u_t = torch.func.jacrev(lambda *args: func(*args)[0], 0)(t, x, mu, theta)
        v_t = torch.func.jacrev(lambda *args: func(*args)[1], 0)(t, x, mu, theta)
        zeros = torch.zeros_like(x[0:1])
        return torch.concatenate((u_t, v_t, zeros), dim=0)

    def functional_space_operator(self, func, t, x, mu, theta):
        f_vals = func(t, x, mu, theta)
        u, v, _ = f_vals[0], f_vals[1], f_vals[2]
        # Incompressibility condition
        grad_u = torch.func.jacrev(lambda *args: func(*args)[0], 1)
        grad_v = torch.func.jacrev(lambda *args: func(*args)[1], 1)
        u_xy = grad_u(t, x, mu, theta)
        v_xy = grad_v(t, x, mu, theta)
        u_x, u_y = u_xy[0:1], u_xy[1:2]
        v_x, v_y = v_xy[0:1], v_xy[1:2]
        op1 = u_x + v_y
        # Navier-Stokes equations
        hessian_u = torch.func.jacrev(grad_u, 1)(t, x, mu, theta)
        hessian_v = torch.func.jacrev(grad_v, 1)(t, x, mu, theta)
        u_xx, u_yy = hessian_u[0, 0], hessian_u[1, 1]
        v_xx, v_yy = hessian_v[0, 0], hessian_v[1, 1]
        p_xy = torch.func.jacrev(lambda *args: func(*args)[2], 1)(t, x, mu, theta)
        p_x, p_y = p_xy[0:1], p_xy[1:2]
        op2a = (u * u_x + v * u_y) + p_x - mu * (u_xx + u_yy)
        op2b = (u * v_x + v * v_y) + p_y - mu * (v_xx + v_yy)
        return torch.concatenate((op2a, op2b, op1), dim=0)

    def functional_operator(self, func, t, x, mu, theta):
        res = self.functional_time_operator(func, t, x, mu, theta)
        res += self.functional_space_operator(func, t, x, mu, theta)
        return res

    def functional_operator_bc_right(self, func, t, x, n, mu, theta):
        p = func(t, x, mu, theta)[2:3]
        return torch.concatenate((0 * p, p), dim=0)

    def functional_operator_bc_other(self, func, t, x, n, mu, theta):
        f_vals = func(t, x, mu, theta)
        return f_vals[0:2]

    def functional_operator_bc(self) -> OrderedDict:
        return OrderedDict(
            [
                ((0, 2, 3, 4), self.functional_operator_bc_other),
                (1, self.functional_operator_bc_right),
            ]
        )

    def functional_operator_ic(self, func, x, mu, theta) -> torch.Tensor:
        t = torch.zeros_like(mu)
        return func(t, x, mu, theta)


def vorticity(space, t, xy, mu):
    xy.x.requires_grad_()
    w = space.evaluate(t, xy, mu)
    u, v, _ = w.get_components()
    dxv, _ = space.grad(v, xy)
    _, dyu = space.grad(u, xy)
    return dxv - dyu


def norm_of_speed(space, t, xy, mu):
    xy.x.requires_grad_()
    w = space.evaluate(t, xy, mu)
    u, v, _ = w.get_components()
    return torch.sqrt(u**2 + v**2)


def run_navier_stokes_2d(
    param=[0.1, 0.1],
    plot_fig=True,
    save_fig=False,
    new_training=False,
    preconditioner="ENG",
):
    # Domain
    t_min, t_max = 0.0, 10.0
    domain_t = [t_min, t_max]
    box = [[-2.5, 7.5], [-2.5, 2.5]]
    center = (0.0, 0.0)
    radius = 1.0
    parameter_domain = [(0.1, 0.1)]

    domain_x = Square2D(box, is_main_domain=True)
    hole = Disk2D(center, radius, is_main_domain=False)
    domain_x.add_hole(hole)

    sampler = TensorizedSampler(
        [
            UniformTimeSampler(domain_t),
            DomainSampler(domain_x),
            UniformParametricSampler(parameter_domain),
        ]
    )

    # Network
    nb_unknowns, nb_parameters = 3, 1
    tlayers = [16, 32, 16]  # Example layer sizes for the MLP
    space = NNxtSpace(
        nb_unknowns,
        nb_parameters,
        GenericMLP,
        domain_x,
        sampler,
        layer_sizes=tlayers,
        activation_type="sine",
    )

    # EDP
    pde = NavierStokes2DStrongForm(space)

    pinns = NaturalGradientTemporalPinns(
        pde,
        in_weights=[1.0, 10.0, 10.0],
        bc_type="weak",
        bc_weight=10.0,
        ic_type="weak",
        ic_weight=10.0,
        one_loss_by_equation=True,
        matrix_regularization=5e-5,
    )

    if preconditioner == "SNG":
        pinns = NaturalGradientTemporalPinns(
            pde,
            in_weights=[1.0, 10.0, 10.0],
            bc_type="weak",
            bc_weight=10.0,
            ic_type="weak",
            ic_weight=10.0,
            one_loss_by_equation=True,
            ng_algo="SNG",
            tol=1e-7,
        )

    ## Load or train the PINN
    filename_extension = "ns_rectangle_" + preconditioner
    default_path = Path(__file__).parent
    default_folder_name = "results"
    if new_training or not pinns.load(
        __file__, filename_extension, default_path, default_folder_name
    ):
        pinns.solve(
            epochs=100,
            n_collocation=8000,
            n_bc_collocation=8000,
            n_ic_collocation=8000,
        )
        pinns.save(__file__, filename_extension, default_path, default_folder_name)

    print(f"Training time: {pinns.training_time:.2f} seconds")
    print(f"Best Loss    : {pinns.best_loss:.2e}")

    additional_scalar_functions = {
        "$\\partial_x v - \\partial_y u$": vorticity,
        "$\\|(u,v)\\|_2$": norm_of_speed,
    }

    plot_abstract_approx_spaces(
        pinns.space,
        domain_x,
        parameter_domain,
        domain_t,
        time_values=[
            t_max,
        ],
        n_visu=128,
        loss=pinns.losses,
        residual=pinns.pde,
        draw_contours=True,
        n_drawn_contours=20,
        parameters_values="mean",
        loss_groups=["bc", "ic"],
        title="Navier Stokes simulation with " + preconditioner + " preconditioning",
        additional_scalar_functions=additional_scalar_functions,
    )

    if save_fig:
        figfilename = "navierstokes.png"
        fig_path = default_path / default_folder_name / figfilename
        plt.savefig(fig_path, bbox_inches="tight")
        print(f"Figure saved at {fig_path}")

    if plot_fig:
        plt.show()

    return pinns


if __name__ == "__main__":
    # Run the Navier-Stokes 2D problem with the specified parameters
    run_navier_stokes_2d(
        param=[0.1, 0.1],
        new_training=True,
        plot_fig=True,
        save_fig=False,
    )