r"""Solves the electrostatics equation in 2D on a bi-material domain using PINNs.

.. math::

    -\sigma (\partial_{xx} u_1 - \partial_{yy} u_1) = 0 in \Omega

The domain is defined as a bi-material region (a rectangle within a square) with different material properties (:math:`\sigma = 1` in the outer domain and :math:`\sigma = 5` in the inner domain); the sixteen boundary and interface conditions are all enforced weakly.

The neural network used is a simple MLP (Multilayer Perceptron); the optimization is performed using either Natural Gradient Descent, or SS-BFGS.
"""

import timeit

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

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 (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo_parameters import (
    UniformParametricSampler,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_field import (
    ParamFieldFunction,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import (
    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 (
    BoundaryResidual,
    InteriorResidual,
)
from scimba_jax.plots.plots_nd import (
    plot_abstract_approx_space,
    plot_abstract_approx_spaces,
)


class BiMaterialLaplacian(InteriorResidual):
    def __init__(self, fact, component, domain, f_rhs):
        super().__init__(domain=domain, size=1, model_type="x_mu", f_rhs=f_rhs)
        self.fact = fact
        self.component = component

    def construct_residual(self, *vars):
        rho = vars[0]
        assert isinstance(rho, ParamVecFunction)
        # split outputs
        rho_c = rho.component(self.component)
        # compute laplacian
        lap_c = rho_c.laplacian_x()
        # scale by factor
        return -self.fact * lap_c


class NeumannResidual(BoundaryResidual):
    def __init__(
        self, index, domain, size=1, model_type="x_mu", f_rhs=None, time_domain=None
    ):
        super().__init__(
            domain=domain,
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
            time_domain=time_domain,
        )
        self.index = index

    def construct_residual(self, *vars):
        rho = vars[0]
        assert isinstance(rho, ParamVecFunction)
        rho_index = rho.component(self.index)
        # rho_with_n = rho_index.with_normals()
        grad_x_rho = rho_index.gradient_x()
        n_model = ParamFieldFunction(
            rho.dims,
            lambda *args: args[rho.argnums["n"]],
            rho.f_type,
        )
        grad_n_rho = grad_x_rho.dot(n_model)
        return grad_n_rho


class DirichletResidual(BoundaryResidual):
    def __init__(
        self, index, domain, size=1, model_type="x_mu", f_rhs=None, time_domain=None
    ):
        super().__init__(
            domain=domain,
            size=size,
            model_type=model_type,
            f_rhs=f_rhs,
            time_domain=time_domain,
        )
        self.index = index

    def construct_residual(self, *vars):
        rho = vars[0]
        assert isinstance(rho, ParamVecFunction)
        rho_index = rho.component(self.index)
        return rho_index


class InterfaceResidual(BoundaryResidual):
    def __init__(self, fact, index, domain, f_rhs):
        super().__init__(domain=domain, size=2, model_type="x_mu", f_rhs=f_rhs)

        self.fact = fact
        self.index = index

    def construct_residual(self, *vars):
        rho = vars[0]
        assert isinstance(rho, ParamVecFunction)
        # split outputs
        rho_i = rho.component(0)
        rho_o = rho.component(1)
        # compute partial derivative with respect to symbol
        rho_i_index = rho_i.partial_derivative_x(component_index=self.index)
        rho_o_index = rho_o.partial_derivative_x(component_index=self.index)
        # scale by factor
        equality = rho_i - rho_o
        discontinuity = rho_i_index - self.fact * rho_o_index
        return ParamVecFunction.cat([equality, discontinuity])


class BiMaterialElectrostatic(AbstractPhysicalModel):
    """BiMaterialElectrostatic model."""

    def __init__(self, main_domain, f_rhs, f_bc_rhs):
        super().__init__(main_domain=main_domain)

        self.e_r = 3.3
        # inner residuals
        self.physical_residuals = {
            "outer": BiMaterialLaplacian(
                fact=1.0, component=0, domain=main_domain, f_rhs=f_rhs["outer"]
            ),
            "inner": BiMaterialLaplacian(
                fact=5.0, component=1, domain=main_domain, f_rhs=f_rhs["inner"]
            ),
        }
        # boundary residuals
        self.physical_residuals["bc SN"] = NeumannResidual(
            index=0,
            domain=self.boundaries["bc SN"],
            size=1,
            model_type="x_mu",
            f_rhs=f_bc_rhs["bc SN"],
        )
        self.physical_residuals["bc E"] = DirichletResidual(
            index=0,
            domain=self.boundaries["bc E"],
            size=1,
            model_type="x_mu",
            f_rhs=f_bc_rhs["bc E"],
        )
        self.physical_residuals["bc W"] = DirichletResidual(
            index=0,
            domain=self.boundaries["bc W"],
            size=1,
            model_type="x_mu",
            f_rhs=f_bc_rhs["bc W"],
        )
        self.physical_residuals["bc inner_SN"] = InterfaceResidual(
            fact=self.e_r,
            index=1,
            domain=self.boundaries["bc inner_SN"],
            f_rhs=f_bc_rhs["bc inner_SN"],
        )
        self.physical_residuals["bc inner_WE"] = InterfaceResidual(
            fact=self.e_r,
            index=0,
            domain=self.boundaries["bc inner_WE"],
            f_rhs=f_bc_rhs["bc inner_WE"],
        )


def test_f(x, mu):
    # jax.debug.print("x.shape: {x}", x=x.shape)
    print("x.shape: ", x.shape)
    return jnp.zeros_like(x[:, 0])


f_rhs = {
    "inner": lambda *args: 0.0,
    "outer": lambda *args: 0.0,
}
f_bc_rhs = {
    "bc SN": lambda *args: 0.0,
    "bc E": lambda *args: -1.0,
    "bc W": lambda *args: 1.0,
    "bc inner_SN": lambda *args: 0.0,
    "bc inner_WE": lambda *args: 0.0,
}

domain_x = [(-1.0, 1.0), (-1.0, 1.0)]
domain_ix = [(-0.5, 0.5), (-0.25, 0.25)]
domain_mu = []

big_square = Square2D(domain_x, is_main_domain=True, label_str="outer")
small_square = Square2D(domain_ix, is_main_domain=False, label_str="inner")

big_square.add_subdomain(small_square)

big_square.set_boundaries_dict(
    {
        "bc SN": ["bc south", "bc north"],
        "bc E": ["bc east"],
        "bc W": ["bc west"],
        "bc inner_SN": ["bc inner south", "bc inner north"],
        "bc inner_WE": ["bc inner west", "bc inner east"],
    }
)
# bnd = big_square.get_all_boundaries()
# for label in bnd:
#     s = ""
#     for lab in bnd[label]:
#         s += lab + ", "
#     print(label, ": ", s)

sampler = TensorizedSampler(
    [DomainSampler(big_square), UniformParametricSampler(domain_mu)], bc=True
)

key = jax.random.PRNGKey(0)
# N_COLLOC = 8
# N_BC_COLLOC = 8
# key, sample = sampler.sample(key, N_COLLOC, N_BC_COLLOC)
# for label in sample:
#     print(label, ": %d pnts" % len(sample[label][0]))

nn = MLP(in_size=2, out_size=2, hidden_sizes=[32, 32], key=key)
space = ApproximationSpace({"x": 2}, [(nn, "vec", 2)], model_type="x_mu")

model = BiMaterialElectrostatic(big_square, f_rhs, f_bc_rhs)

bc_weight = 30.0
weights = {
    "outer": [1.0],
    "inner": [1.0],
    "bc SN": [bc_weight * 2.0],
    "bc E": [bc_weight * 3.0],
    "bc W": [bc_weight * 3.0],
    "bc inner_SN": [bc_weight * 5.0, bc_weight],
    "bc inner_WE": [bc_weight * 4.0, bc_weight],
}
# ev = Evaluator(model, space)

N_EPOCHS_ENG = 400
N_EPOCHS_SSBFGS = 2000
N_COLLOC = 2000
N_BC_COLLOC = 6000
key, sample_dict = sampler.sample(key, N_COLLOC, N_BC_COLLOC)


print("\n\n")
print("@@@@@@@@@@@@@@@ test SS-BFGS @@@@@@@@@@@@@@@@@@@@@@")

pinn = Projector(
    model,
    space,
    sampler,
    optimizer="SS-BFGS",
    # weights=weights
)
loss0 = pinn.evaluate_loss(space, sample_dict)
print("initial loss: ", loss0)

start = timeit.default_timer()
new_loss, nspace, key, loss_history = pinn.project(
    space, key, N_EPOCHS_SSBFGS, N_COLLOC, N_BC_COLLOC
)
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_SSBFGS, end - start)

plot_abstract_approx_space(
    pinn.space,
    big_square,
    domain_mu,
    loss=pinn.losses,
    residual=pinn.model,
    derivatives=["ux"],
    components=[{"outer": 0, "inner": 1}],
    draw_contours=True,
    n_drawn_contours=20,
    title="Solving Bimaterial Electrostatic with SS-BFGS",
    loss_groups=["bc"],
)
plt.show()

print("\n\n")
print("@@@@@@@@@@@@@@@ test ENG @@@@@@@@@@@@@@@@@@@@@@")

nn = MLP(in_size=2, out_size=2, hidden_sizes=[16, 16, 16], key=key)
space2 = ApproximationSpace({"x": 2}, [(nn, "vec", 2)], model_type="x_mu")

pinn2 = Projector(
    model,
    space2,
    sampler,
    # weights=weights,
    matrix_regularization=5e-4,
)
loss0 = pinn2.evaluate_loss(space2, sample_dict)
print("initial loss: ", loss0)

start = timeit.default_timer()
new_loss, nspace, key, loss_history = pinn2.project(
    space2, key, N_EPOCHS_ENG, N_COLLOC, N_BC_COLLOC
)
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_ENG, end - start)

# pinn2.plot(
#     components=[{"outer": 0, "inner": 1}],
# )
# plt.show()
#
plot_abstract_approx_spaces(
    (pinn.space, pinn2.space),
    big_square,
    domain_mu,
    loss=(pinn.losses, pinn2.losses),
    residual=(pinn.model, pinn2.model),
    derivatives=["ux"],
    components=[{"outer": 0, "inner": 1}],
    draw_contours=True,
    n_drawn_contours=20,
    title="Solving Bimaterial Electrostatic",
    titles=("SS-BFGS", "ENG"),
    loss_groups=(["bc"], ["bc"]),
)
plt.show()