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.

Two different NNs are used, one for uv, one for the pression.
"""

import timeit

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

from scimba_jax.domains.meshless_domains.base import SurfacicDomain, VolumetricDomain
from scimba_jax.domains.meshless_domains.domains_2d import Square2D
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DataSampler,
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo_parameters import (
    UniformParametricSampler,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import (
    ParamFieldFunction,
    ParamScalarFunction,
    ParamVecFunction,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.abstract_physical_model import (
    AbstractPhysicalModel,
)
from scimba_jax.physical_models.abstract_residuals import (
    NDARRAYS_FUNC_TYPE,
    BoundaryResidual,
    DataResidual,
    InteriorResidual,
)
from scimba_jax.plots.plots_nd import (
    plot_abstract_approx_space,
)

N_COLLOC = 1000
N_BC_COLLOC = 2000
N_EPOCHS = 200
N_EPOCHS_ADAM = 1000
N_EPOCHS_SSBFGS = 1000


# The interior residual
class StationaryNavierStokes(InteriorResidual):
    def __init__(self, domain, f_rhs, Re):
        super().__init__(domain=domain, size=3, model_type="x_mu", f_rhs=f_rhs)
        self.Re = Re

    def construct_residual(self, *vars):
        uv, p = vars[0], vars[1]
        assert isinstance(uv, ParamFieldFunction)
        assert isinstance(p, ParamScalarFunction)

        # Incompressibility condition
        op1 = uv.divergence_x()

        # Navier-Stokes equations
        uv_dot_uv_xy = uv @ uv.jacobian_x()
        grad_p = p.gradient_x()
        lapl = uv.laplacian_x()
        op2 = uv_dot_uv_xy + grad_p - (1.0 / self.Re) * lapl
        return ParamVecFunction.cat((op1, op2))


# The boundary residual: Dirichlet on 2 first components
class RestrainedDirichletResidual(BoundaryResidual):
    def __init__(
        self,
        domain: dict[str, SurfacicDomain],
        size: int = 2,
        model_type: str = "x_mu",
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(
            domain=domain,
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
        )

    # def construct_residual(self, *vars):
    #     # Here assume that variables are functions of normals
    #     uvp = vars[0]
    #     u = uvp.component(0)
    #     v = uvp.component(1)
    #     return ParamVecFunction.cat((u, v))

    def construct_residual(self, *vars):
        # Here assume that variables are functions of normals
        uv = vars[0]
        assert isinstance(uv, ParamFieldFunction)
        return uv


# A data loss for the condition p(0,0) = 0
class RestrainedCollocDataResidual(DataResidual):
    def __init__(self, size: int = 1, model_type: str = "x_mu"):
        super().__init__(size, model_type)

    # def construct_residual(self, *vars):
    #     uvp = vars[0]
    #     p = uvp.component(2)
    #     return p

    def construct_residual(self, *vars):
        p = vars[1]
        assert isinstance(p, ParamScalarFunction)
        return p


# The model
class StationaryNavierStokesEquation(AbstractPhysicalModel):
    def __init__(
        self,
        main_domain: VolumetricDomain,
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
        f_bc_rhs: dict[str, NDARRAYS_FUNC_TYPE] = {},
        Re: float = 100,
    ):
        super().__init__(main_domain=main_domain)
        self.physical_residuals = {
            self.main_domain.get_label(): StationaryNavierStokes(
                domain=main_domain,
                f_rhs=f_rhs,
                Re=Re,
            ),
        }
        for label in self.boundaries:
            self.physical_residuals[label] = RestrainedDirichletResidual(
                domain=self.boundaries[label],
                size=2,
                model_type="x_mu",
                f_rhs=f_bc_rhs[label],
            )


f_bc_rhs = {
    "bc SEW": lambda *args: jnp.array([0.0, 0.0]),
    "bc N": lambda *args: jnp.array([4 * args[0][0] * (1.0 - args[0][0]), 0.0]),
}

key = jax.random.PRNGKey(0)

domain_x = [(0.0, 1.0), (0.0, 1.0)]
domain_mu = []

dx = Square2D(domain_x, is_main_domain=True)

dx.set_boundaries_dict(
    {
        "bc SEW": ["bc south", "bc east", "bc west"],
        "bc N": ["bc north"],
    }
)

# For the constraint p(0,0) = 0, we use a data loss
# on the point (0,0)
x_data = jnp.zeros((1, 2))
mu_data = jnp.zeros((1, 0))
y_data = jnp.zeros((1, 1))
data_sampler = DataSampler((x_data, mu_data, y_data))

sampler = TensorizedSampler(
    [DomainSampler(dx), UniformParametricSampler(domain_mu)],
    bc=True,
    data_samplers={"data": data_sampler},
)

# sample the domain
key, sample_dict = sampler.sample(key, N_COLLOC)

# create an MLP with 2 inputs and 3 outputs for u,v,p
nn_uv = MLP(in_size=2, out_size=2, hidden_sizes=[18, 20, 18], key=key)
nn_p = MLP(in_size=2, out_size=1, hidden_sizes=[18, 20, 18], key=key)

# create the approximation space
space = ApproximationSpace(
    {"x": 2},
    [(nn_uv, "field", 2), (nn_p, "scalar", 1)],
    model_type="x_mu",
)

# create the model
# add a data loss for the constraint p(0,0) = 0
model = StationaryNavierStokesEquation(dx, f_bc_rhs=f_bc_rhs)
model.add_data_residual("data", RestrainedCollocDataResidual(size=1, model_type="x_mu"))

print("\n\n")
print("@@@@@@@@@@@@@@@ create a pinn with weak bc @@@@@@@@@@@@@@@@@@@@@")

pinn = Projector(model, space, sampler)

loss = pinn.evaluate_loss(space, sample_dict)
print("initial loss: ", loss)
losses = pinn.evaluate_losses(space, sample_dict)
print("initial losses: ", losses)

print("\n\n")
print("@@@@@@@@@@@@@@@ train with Energy Natural Gradient @@@@@@@@@@@@@@@@@@@@@")

start = timeit.default_timer()
new_loss, nspace, key, loss_history = pinn.project(
    space, key, N_EPOCHS, N_COLLOC, N_BC_COLLOC, n_dl=1
)
pinn.losses.loss_history = loss_history
pinn.best_loss = new_loss
pinn.space = nspace
end = timeit.default_timer()

print("best loss: ", new_loss)
print("time for %d epochs: " % N_EPOCHS, end - start)

plot_abstract_approx_space(
    pinn.space,  # the approximation space
    dx,  # the spatial domain
    domain_mu,  # the parameter's domain
    components=[0, 1, 2],
    loss=pinn.losses,  # for plot of the loss: the losses
    residual=pinn.model,
    draw_contours=True,
    n_drawn_contours=20,
    loss_groups=["bc", "interior"],
)
plt.show()

print("\n\n")
print("@@@@@@@@@@@@@@@ train with Adam @@@@@@@@@@@@@@@@@@@@@@")

# create an MLP with 2 inputs and 3 outputs for u,v,p
nn_uv = MLP(in_size=2, out_size=2, hidden_sizes=[18, 20, 18], key=key)
nn_p = MLP(in_size=2, out_size=1, hidden_sizes=[18, 20, 18], key=key)

# create the approximation space
space2 = ApproximationSpace(
    {"x": 2},
    [(nn_uv, "field", 2), (nn_p, "scalar", 1)],
    model_type="x_mu",
)
model = StationaryNavierStokesEquation(dx, f_bc_rhs=f_bc_rhs)
model.add_data_residual("data", RestrainedCollocDataResidual(size=1, model_type="x_mu"))

pinn2 = Projector(model, space2, sampler, optimizer="Adam", learning_rate=1e-4)

start = timeit.default_timer()
new_loss, nspace, key, loss_history = pinn2.project(
    space2, key, N_EPOCHS_ADAM, N_COLLOC, N_BC_COLLOC, n_dl=1
)
pinn2.losses.loss_history = loss_history
pinn2.best_loss = new_loss
pinn2.space = nspace
end = timeit.default_timer()

print("best loss: ", new_loss)
print("time for %d epochs: " % N_EPOCHS_ADAM, end - start)

plot_abstract_approx_space(
    pinn2.space,  # the approximation space
    dx,  # the spatial domain
    domain_mu,  # the parameter's domain
    components=[0, 1, 2],
    loss=pinn2.losses,  # for plot of the loss: the losses
    residual=pinn2.model,
    draw_contours=True,
    n_drawn_contours=20,
    loss_groups=["bc", "interior"],
)
plt.show()