r"""Solves the 2d Navier Stokes Equations described in

@article{hao2023pinnacle,
  title={PINNacle: A Comprehensive Benchmark of Physics-Informed Neural Networks for Solving PDEs},
  author={Hao, Zhongkai and Yao, Jiachen and Su, Chang and Su, Hang and Wang, Ziao and Lu, Fanzhi and Xia, Zeyu and Zhang, Yichi and Liu, Songming and Lu, Lu and others},
  journal={arXiv preprint arXiv:2306.08827},
  year={2023}
}

See Appendix B1, problem 11, page 17.
"""

from pathlib import Path
from typing import Callable, 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 NNxSpace
from scimba_torch.domain.meshless_domain.domain_2d import Square2D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.elliptic_pde.pinns import (
    NaturalGradientPinnsElliptic,
)
from scimba_torch.physical_models.elliptic_pde.abstract_elliptic_pde import (
    StrongFormEllipticPDE,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces

torch.manual_seed(0)


class NavierStokes2DStrongForm(StrongFormEllipticPDE):
    def __init__(self, space: AbstractApproxSpace, fixed_params, **kwargs):
        super().__init__(
            space, linear=False, residual_size=3, bc_residual_size=8, **kwargs
        )

        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.Re = fixed_params["Re"]  # Reynold's number

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

    def operator(self, w, x, mu):
        u, v, p = w.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 - (1.0 / self.Re) * (u_xx + u_yy)
        op2b = (u * v_x + v * v_y) + p_y - (1.0 / self.Re) * (v_xx + v_yy)

        return op1, op2a, op2b

    def bc_rhs(self, w, 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)

        return (
            torch.zeros_like(x1_0),
            torch.zeros_like(x1_0),
            torch.zeros_like(x1_1),
            torch.zeros_like(x1_1),
            4.0 * x1_2 * (1.0 - x1_2),
            torch.zeros_like(x1_2),
            torch.zeros_like(x1_3),
            torch.zeros_like(x1_3),
        )

    def bc_operator(self, w, x, n, mu):
        u, v, _ = w.get_components()
        ldown, least, ltop, lwest = 0, 1, 2, 3
        u_down = w.restrict_to_labels(u, labels=[ldown])
        v_down = w.restrict_to_labels(v, labels=[ldown])
        u_east = w.restrict_to_labels(u, labels=[least])
        v_east = w.restrict_to_labels(v, labels=[least])
        u_top = w.restrict_to_labels(u, labels=[ltop])
        v_top = w.restrict_to_labels(v, labels=[ltop])
        u_west = w.restrict_to_labels(u, labels=[lwest])
        v_west = w.restrict_to_labels(v, labels=[lwest])

        return u_down, v_down, u_east, v_east, u_top, v_top, u_west, v_west

    def functional_operator_bc_all(self, func, x, n, mu, theta):
        return func(x, mu, theta)[..., :2]

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

    def functional_operator(
        self,
        func: Callable,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        f_vals = func(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], 0)
        grad_v = torch.func.jacrev(lambda *args: func(*args)[1], 0)
        u_xy = grad_u(x, mu, theta)
        v_xy = grad_v(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, 0)(x, mu, theta)
        hessian_v = torch.func.jacrev(grad_v, 0)(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], 0)(x, mu, theta)
        p_x, p_y = p_xy[0:1], p_xy[1:2]
        op2a = (u * u_x + v * u_y) + p_x - (1.0 / self.Re) * (u_xx + u_yy)
        op2b = (u * v_x + v * v_y) + p_y - (1.0 / self.Re) * (v_xx + v_yy)

        return torch.concatenate((op1, op2a, op2b), dim=0)


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


def run_navier_stokes_2d(save_fig=False, new_training=False):
    # Domain
    box = [(0.0, 1.0), (0.0, 1.0)]
    parameter_domain = []

    domain_x = Square2D(box, is_main_domain=True)
    full_bc_domain = domain_x.full_bc_domain()
    for bc in full_bc_domain:
        domain_x.add_bc_domain(bc)

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

    # Network
    nb_unknowns, nb_parameters = 3, 0
    # tlayers = [40, 60, 60, 60, 40]  # Example layer sizes for the MLP
    tlayers = [18, 20, 20, 18]  # Example layer sizes for the MLP
    space = NNxSpace(
        nb_unknowns,
        nb_parameters,
        GenericMLP,
        domain_x,
        sampler,
        layer_sizes=tlayers,
        activation_type="tanh",
    )

    fixed_params = {"Re": 100}
    pde = NavierStokes2DStrongForm(space, fixed_params=fixed_params)

    pinns = NaturalGradientPinnsElliptic(
        pde,
        bc_type="weak",
        # bc_weights=[10.0] * 8,
        one_loss_by_equation=True,
        matrix_regularization=1e-6,
    )

    ## Load or train the PINN
    filename_extension = "ns_square"
    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=1000, n_bc_collocation=2000, verbose=False
        )
        pinns.save(__file__, filename_extension, default_path, default_folder_name)

    # for plotting pinns
    if save_fig:
        additional_scalar_functions = {
            "$\\partial_x v - \\partial_y u$": vorticity,
        }

        plot_abstract_approx_spaces(
            pinns.space,
            domain_x,
            parameter_domain,
            n_visu=300,
            loss=pinns.losses,
            residual=pinns.pde,
            draw_contours=True,
            n_drawn_contours=40,
            loss_groups=["bc"],
            additional_scalar_functions=additional_scalar_functions,
        )

        plt.show()

        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}")

    return pinns


if __name__ == "__main__":
    run_navier_stokes_2d(new_training=True, save_fig=True)