"""Solves the radiative transfer equation with a discrete PINN."""

import time

import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import numpy as np
import torch

from scimba_torch.approximation_space.nn_space import NNxvSpace
from scimba_torch.domain.meshless_domain.domain_2d import Circle2D, Square2D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import (
    UniformParametricSampler,
    UniformVelocitySampler,
)
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.temporal_pde.discrete_pinn import DiscretePINN
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
    TimeDiscreteNaturalGradientProjector,
)
from scimba_torch.physical_models.kinetic_pde.radiative_transfer import (
    RadiativeTransfer,
)
from scimba_torch.physical_models.temporal_pde.abstract_temporal_pde import TemporalPDE
from scimba_torch.utils.scimba_tensors import LabelTensor


def f_rhs(w, x, v, mu):
    x_x, x_y = x.get_components()
    return 0 * x_x


def f_bc(w, x, v, n, mu):
    x1 = x.get_components()
    return 0 * x1


def f_ini(x, v, mu):
    return exact(x, v, LabelTensor(torch.zeros_like(x.x)))


def exact(x, v, t):
    x_x, x_y = x.get_components()
    v_x, v_y = v.get_components()
    t_1, t_2 = t.get_components()
    # return torch.exp(-((x_x - v_x * 0.001) ** 2 * 10)) * torch.exp(-((x_y - v_y * 0.001) ** 2 * 10)) * torch.exp(-((v_x - 1.0) ** 2 * 1)) * torch.exp(-((v_y) ** 2 * 1))
    # return torch.exp(-((x_x - v_x * t_1) ** 2 * 10)) * torch.exp(-((x_y - v_y * t_1) ** 2 * 10)) * torch.exp(-((v_x - 1.0) ** 2 * 1)) * torch.exp(-((v_y) ** 2 * 1))
    # return torch.exp(-((x_x) ** 2 * 10)) * torch.exp(-((x_y) ** 2 * 10)) * torch.exp(-((v_x - 1.0) ** 2 * 0.1)) * torch.exp(-((v_y) ** 2 * 0.1))
    return torch.exp(-((x_x - v_x * t_1 + 0.7) ** 2 * 10)) * torch.exp(
        -((x_y - v_y * t_1) ** 2 * 10)
    ) * torch.exp(-((v_x - 1.0) ** 2 * 0.1)) * torch.exp(
        -((v_y) ** 2 * 0.1)
    ) + torch.exp(-((x_x - v_x * t_1) ** 2 * 10)) * torch.exp(
        -((x_y - v_y * t_2 + 0.7) ** 2 * 10)
    ) * torch.exp(-((v_x) ** 2 * 0.1)) * torch.exp(-((v_y - 1.0) ** 2 * 0.1))


# %%

start_time = time.time()

torch.random.manual_seed(0)

domain_x = Square2D([(-1, 1), (-1, 1)], is_main_domain=True)

domain_v = Circle2D(torch.Tensor([0, 0]), 1.0)

sampler = TensorizedSampler(
    [
        DomainSampler(domain_x),
        UniformVelocitySampler(domain_v),
        UniformParametricSampler([]),
    ]
)

####################### Pinn discret #################
space = NNxvSpace(1, 0, GenericMLP, domain_x, domain_v, sampler, layer_sizes=[40, 40])
pde = RadiativeTransfer(space, TemporalPDE, init=f_ini, f=f_rhs, g=f_bc)

# TEST MÊME PARAM QUE 2_sources_1 MAIS AVEC PLUS GROS NN

projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)

scheme = DiscretePINN(pde, projector, scheme="rk4")
scheme.initialization(
    epochs=200,
    verbose=True,
    n_collocation=20000,
    type_linesearch="armijo",
    data_linesearch={"n_step_max": 10, "alpha": 0.01},
)
scheme.projector.save("ini_radiative")

end_init = time.time()
# %%
scheme.projector.load("ini_radiative")
N_sample = 50000
N_sample_axis = int(np.sqrt(N_sample))

x, v, mu = sampler.sample(N_sample_axis**2)

x_x = torch.linspace(-1, 1, N_sample_axis)
x_y = torch.linspace(-1, 1, N_sample_axis)
X, Y = torch.meshgrid(x_x, x_y)
x = LabelTensor(torch.stack((X.flatten(), Y.flatten())).T)

# v_x, v_y = v.get_components()
# i = torch.randint(0, N_sample_axis**2, (1,))

# v_plot1 = [0.7, 0.51]
# v_plot2 = [1.0, 0.0]
# v_plot3 = [0.7, -0.51]

v_plot1 = [0.0, 1.0]
v_plot2 = [1.0, 0.0]
v_plot3 = [np.sqrt(2) / 2, np.sqrt(2) / 2]

v1 = LabelTensor(torch.tensor(v_plot1).unsqueeze(0).tile((N_sample_axis**2, 1)))
v2 = LabelTensor(torch.tensor(v_plot2).unsqueeze(0).tile((N_sample_axis**2, 1)))
v3 = LabelTensor(torch.tensor(v_plot3).unsqueeze(0).tile((N_sample_axis**2, 1)))

predictions1 = scheme.pde.space.evaluate(x, v1, mu).w
predictions2 = scheme.pde.space.evaluate(x, v2, mu).w
predictions3 = scheme.pde.space.evaluate(x, v3, mu).w

# expected = f_ini(x, v, mu)


T_ = LabelTensor(torch.zeros_like(x.x))
expected1 = exact(x, v1, T_)
expected2 = exact(x, v2, T_)
expected3 = exact(x, v3, T_)

# vmin = expected2[:,0].min()
# vmax = expected2[:,0].max()

predicted1 = predictions1.reshape(N_sample_axis, N_sample_axis)
predicted2 = predictions2.reshape(N_sample_axis, N_sample_axis)
predicted3 = predictions3.reshape(N_sample_axis, N_sample_axis)

expected1 = expected1.reshape(N_sample_axis, N_sample_axis)
expected2 = expected2.reshape(N_sample_axis, N_sample_axis)
expected3 = expected3.reshape(N_sample_axis, N_sample_axis)

print("Initialization time: ", end_init - start_time)

fig = plt.figure(figsize=(15, 10))

gs = gridspec.GridSpec(2, 3, height_ratios=[2, 2])

ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[0, 1])
ax3 = fig.add_subplot(gs[0, 2])
ax4 = fig.add_subplot(gs[1, 0])
ax5 = fig.add_subplot(gs[1, 1])
ax6 = fig.add_subplot(gs[1, 2])

# im1 = ax1.imshow(predicted1.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
# ax1.set_title("Initialization for an angle of %s \n(predicted)" % v_plot1)
# cbar1 = fig.colorbar(im1, ax=ax1, orientation='vertical')
# im2 = ax2.imshow(predicted2.detach().numpy().T, origin = 'lower', cmap = 'jet', extent = (-1, 1, -1, 1))
# ax2.set_title("Initialization for an angle of %s \n(predicted in privileged direction)" % v_plot2)
# cbar2 = fig.colorbar(im2, ax=ax2, orientation='vertical')
# im3 = ax3.imshow(predicted3.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
# ax3.set_title("Initialization for an angle of %s \n(predicted)" % v_plot3)
# cbar3 = fig.colorbar(im3, ax=ax3, orientation='vertical')
# im4 = ax4.imshow(expected1.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
# ax4.set_title("Initialization for an angle of %s \n(expected)" % v_plot1)
# cbar4 = fig.colorbar(im4, ax=ax4, orientation='vertical')
# im5 = ax5.imshow(expected2.detach().numpy().T, origin = 'lower', cmap = 'jet', extent = (-1, 1, -1, 1))
# ax5.set_title("Initialization for an angle of %s \n(expected in privileged direction)" % v_plot2)
# cbar5 = fig.colorbar(im5, ax=ax5, orientation='vertical')
# im6 = ax6.imshow(expected3.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
# ax6.set_title("Initialization for an angle of %s \n(expected)" % v_plot3)
# cbar6 = fig.colorbar(im6, ax=ax6, orientation='vertical')

im1 = ax1.imshow(
    predicted1.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax1.set_title(
    "Initialization for an angle of %s \n(predicted in privileged direction \nfor the bottom gaussian)"
    % v_plot1
)
cbar1 = fig.colorbar(im1, ax=ax1, orientation="vertical")
im2 = ax2.imshow(
    predicted2.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax2.set_title(
    "Initialization for an angle of %s \n(predicted in privileged direction \nfor left gaussian)"
    % v_plot2
)
cbar2 = fig.colorbar(im2, ax=ax2, orientation="vertical")
im3 = ax3.imshow(
    predicted3.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax3.set_title("Initialization for an angle of %s \n(predicted)" % [0.707, 0.707])
cbar3 = fig.colorbar(im3, ax=ax3, orientation="vertical")
im4 = ax4.imshow(
    expected1.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax4.set_title(
    "Initialization for an angle of %s \n(expectedin privileged direction \nfor the bottom gaussian)"
    % v_plot1
)
cbar4 = fig.colorbar(im4, ax=ax4, orientation="vertical")
im5 = ax5.imshow(
    expected2.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax5.set_title(
    "Initialization for an angle of %s \n(expected in privileged direction \nfor left gaussian)"
    % v_plot2
)
cbar5 = fig.colorbar(im5, ax=ax5, orientation="vertical")
im6 = ax6.imshow(
    expected3.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax6.set_title("Initialization for an angle of %s \n(expected)" % [0.707, 0.707])
cbar6 = fig.colorbar(im6, ax=ax6, orientation="vertical")


plt.show()

# %%

T = 0.7
scheme.projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
scheme.projector.load("ini_radiative")
scheme.projector.space.load_from_best_approx()
scheme.solve(dt=0.04, T=T, epochs=20, n_collocation=20000, verbose=True)
scheme.projector.save("evo_radiative")
# scheme.projector.load("evo_radiative")

end_evo = time.time()

# %%

alpha = [0.0, 0.0]
# v_alpha = LabelTensor(torch.tensor([v_plot[0] + alpha[0], v_plot[1] + alpha[1]]).unsqueeze(0).tile((N_sample_axis**2,1)))
# predictions = scheme.pde.space.evaluate(x, v_alpha, mu).w

predictions1 = scheme.pde.space.evaluate(x, v1, mu).w
predictions2 = scheme.pde.space.evaluate(x, v2, mu).w
predictions3 = scheme.pde.space.evaluate(x, v3, mu).w
# expected = f_ini(x, v_alpha, mu)

T_ = LabelTensor(torch.ones_like(x.x) * T)
exact1 = exact(x, v1, T_)
exact2 = exact(x, v2, T_)
exact3 = exact(x, v3, T_)

vmin = exact2[:, 0].min()
vmax = exact2[:, 0].max()

predicted1 = predictions1.reshape(N_sample_axis, N_sample_axis)
predicted2 = predictions2.reshape(N_sample_axis, N_sample_axis)
predicted3 = predictions3.reshape(N_sample_axis, N_sample_axis)

expected1 = exact1.reshape(N_sample_axis, N_sample_axis)
expected2 = exact2.reshape(N_sample_axis, N_sample_axis)
expected3 = exact3.reshape(N_sample_axis, N_sample_axis)

erreur1 = 0
erreur2 = 0
erreur3 = 0

erreur1 = torch.sqrt(torch.sum((exact1 - predictions1) ** 2) / torch.sum(exact1**2))
erreur2 = torch.sqrt(torch.sum((exact2 - predictions2) ** 2) / torch.sum(exact2**2))
erreur3 = torch.sqrt(torch.sum((exact3 - predictions3) ** 2) / torch.sum(exact3**2))

print("Evolution time: ", end_evo - end_init)

print("Total time: ", end_evo - start_time)

fig = plt.figure(figsize=(15, 10))

gs = gridspec.GridSpec(2, 3, height_ratios=[2, 2])

ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[0, 1])
ax3 = fig.add_subplot(gs[0, 2])
ax4 = fig.add_subplot(gs[1, 0])
ax5 = fig.add_subplot(gs[1, 1])
ax6 = fig.add_subplot(gs[1, 2])

"""im1 = ax1.imshow(predicted1.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
ax1.set_title("Evolution for an angle of %s \n(predicted)" % v_plot1)
cbar1 = fig.colorbar(im1, ax=ax1, orientation='vertical')
im2 = ax2.imshow(predicted2.detach().numpy().T, origin = 'lower', cmap = 'jet', extent = (-1, 1, -1, 1))
ax2.set_title("Evolution for an angle of %s \n(predicted in privileged direction)" % v_plot2)
cbar2 = fig.colorbar(im2, ax=ax2, orientation='vertical')
im3 = ax3.imshow(predicted3.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
ax3.set_title("Evolution for an angle of %s \n(predicted)" % v_plot3)
cbar3 = fig.colorbar(im3, ax=ax3, orientation='vertical')
im4 = ax4.imshow(expected1.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
ax4.set_title("Evolution for an angle of %s \n(expected)" % v_plot1)
cbar4 = fig.colorbar(im4, ax=ax4, orientation='vertical')
im5 = ax5.imshow(expected2.detach().numpy().T, origin = 'lower', cmap = 'jet', extent = (-1, 1, -1, 1))
ax5.set_title("Evolution for an angle of %s \n(expected in privileged direction)" % v_plot2)
cbar5 = fig.colorbar(im5, ax=ax5, orientation='vertical')
im6 = ax6.imshow(expected3.detach().numpy().T, origin = 'lower',  cmap = 'jet', extent = (-1, 1, -1, 1))
ax6.set_title("Evolution for an angle of %s \n(expected)" % v_plot3)
cbar6 = fig.colorbar(im6, ax=ax6, orientation='vertical')"""

im1 = ax1.imshow(
    predicted1.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax1.set_title(
    "Evolution for an angle of %s \n(predicted in privileged direction \nfor the bottom gaussian) \nRelative L2 error = %e"
    % (v_plot1, float(erreur1))
)
cbar1 = fig.colorbar(im1, ax=ax1, orientation="vertical")
im2 = ax2.imshow(
    predicted2.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax2.set_title(
    "Evolution for an angle of %s \n(predicted in privileged direction \nfor left gaussian) \nRelative L2 error = %e"
    % (v_plot2, float(erreur2))
)
cbar2 = fig.colorbar(im2, ax=ax2, orientation="vertical")
im3 = ax3.imshow(
    predicted3.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax3.set_title(
    "Evolution for an angle of %s \n(predicted) \nRelative L2 error = %e"
    % ([0.707, 0.707], float(erreur3))
)
cbar3 = fig.colorbar(im3, ax=ax3, orientation="vertical")
im4 = ax4.imshow(
    expected1.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax4.set_title(
    "Evolution for an angle of %s \n(expected in privileged direction \nfor the bottom gaussian)"
    % v_plot1
)
cbar4 = fig.colorbar(im4, ax=ax4, orientation="vertical")
im5 = ax5.imshow(
    expected2.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax5.set_title(
    "Evolution for an angle of %s \n(expected in privileged direction \nfor left gaussian)"
    % v_plot2
)
cbar5 = fig.colorbar(im5, ax=ax5, orientation="vertical")
im6 = ax6.imshow(
    expected3.detach().numpy().T, origin="lower", cmap="jet", extent=(-1, 1, -1, 1)
)
ax6.set_title("Evolution for an angle of %s \n(expected)" % [0.707, 0.707])
cbar6 = fig.colorbar(im6, ax=ax6, orientation="vertical")

plt.show()

####################### Pinn discret #################
"""print(">>>> Begin natural gradient >>>>>")
space2 = NNxvSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    domain_v,
    sampler,
    layer_sizes=[16, 8],
)
pde2 = RadiativeTransfer(space2, init=f_ini, f=f_rhs, g=f_bc)
projector2 = NaturalGradientProjector(space2, rhs=f_ini)
scheme2 = DiscretePINN(pde2, projector2, scheme="rk4")

scheme2.initialization(epochs=50, verbose=True, n_collocation=3000)
T = 0.2
scheme2.solve(dt=0.02, T=T, epochs=15, n_collocation=3000, verbose=False)

# %%

predictions = scheme2.pde.space.evaluate(x, mu).w
expected = f_ini(x, mu)

T_ = LabelTensor(torch.ones_like(x.x) * T)
expected = exact(x, T_)

plt.scatter(x.x.detach().cpu(), predictions.detach().cpu(), label="pred")
plt.scatter(x.x.detach().cpu(), expected.detach().cpu(), label="true")
plt.legend()
plt.show()"""