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.
"""

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):
        uvp = vars[0]
        u = uvp.component(0)
        v = uvp.component(1)
        p = uvp.component(2)
        assert isinstance(u, ParamScalarFunction)
        assert isinstance(v, ParamScalarFunction)
        assert isinstance(p, ParamScalarFunction)
        uv = ParamFieldFunction.cat((u, v), "x")

        # 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 = 1,
        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))


# 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


# 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)

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

nn = MLP(in_size=2, out_size=3, hidden_sizes=[18, 20, 20, 18], key=key)
space = ApproximationSpace({"x": 2}, [(nn, "vec", 3)], 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"))

#
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 @@@@@@@@@@@@@@@@@@@@@@")

nn = MLP(in_size=2, out_size=3, hidden_sizes=[18, 20, 20, 18], key=key)
space2 = ApproximationSpace({"x": 2}, [(nn, "vec", 3)], 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()