"""Postprocess the results of the 3D level-set transport problem."""

# %%

import matplotlib.pyplot as plt
import numpy as np
import torch
from matplotlib.colors import LightSource
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from skimage import measure

from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_3d import Cube3D
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.collocation_projector import (
    NaturalGradientProjector,
)
from scimba_torch.numerical_solvers.temporal_pde.neural_semilagrangian import (
    Characteristic,
    NeuralSemiLagrangian,
)
from scimba_torch.physical_models.temporal_pde.advection_diffusion_equation import (
    AdvectionReactionDiffusionDirichletStrongForm,
)
from scimba_torch.utils.scimba_tensors import LabelTensor, MultiLabelTensor

DEFAULT_BOUNDS = [(0, 1), (0, 1), (0, 1)]


def f_ini(x, mu):
    x1, x2, x3 = x.get_components()
    x0 = 0.35
    return (x1 - x0) ** 2 + (x2 - x0) ** 2 + (x3 - x0) ** 2 - 0.15**2


def f_ini_multi_label_tensor(x, mu):
    x1, x2, x3 = x.get_components()
    x0 = 0.35
    return MultiLabelTensor((x1 - x0) ** 2 + (x2 - x0) ** 2 + (x3 - x0) ** 2 - 0.15**2)


def plot_implicit(func, bbox=(0, 1), iter=None):
    """Create a plot of an implicit function.

    Args:
        func: Implicit function (plot where func==0).
        bbox: The x, y, and z limits of plotted interval.
        iter: Iteration number (optional).
    """
    import numpy as np

    def make_xyz(x, y, z):
        if not isinstance(x, np.ndarray):
            x = x * torch.ones_like(torch.tensor(y).flatten().unsqueeze(1))
            y = torch.tensor(y).flatten().unsqueeze(1)
            z = torch.tensor(z).flatten().unsqueeze(1)
        if not isinstance(y, np.ndarray):
            y = y * torch.ones_like(torch.tensor(z).flatten().unsqueeze(1))
            x = torch.tensor(x).flatten().unsqueeze(1)
            z = torch.tensor(z).flatten().unsqueeze(1)
        if not isinstance(z, np.ndarray):
            z = z * torch.ones_like(torch.tensor(x).flatten().unsqueeze(1))
            x = torch.tensor(x).flatten().unsqueeze(1)
            y = torch.tensor(y).flatten().unsqueeze(1)
        x_ = torch.cat((x, y, z), dim=1)
        mu = torch.zeros((x_.shape[0], 0))
        return LabelTensor(x_), LabelTensor(mu)

    xmin, xmax, ymin, ymax, zmin, zmax = bbox * 3
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    A = np.linspace(xmin, xmax, 100)  # resolution of the contour
    B = np.linspace(xmin, xmax, 15)  # number of slices
    A1, A2 = np.meshgrid(A, A)  # grid on which the contour is plotted

    for i, fn in enumerate([func, f_ini_multi_label_tensor]):
        alpha = 1 if i == 0 else 0.5
        cmap = "jet" if i == 0 else None

        for z in B:  # plot contours in the XY plane
            X, Y = A1, A2
            Z = fn(*make_xyz(X, Y, z)).w[:, 0].detach().cpu().reshape((100, 100))
            ax.contour(X, Y, Z + z, [z], zdir="z", alpha=alpha, cmap=cmap)
            # [z] defines the only level to plot for this contour for this value of z

        for y in B:  # plot contours in the XZ plane
            X, Z = A1, A2
            Y = fn(*make_xyz(X, y, Z)).w[:, 0].detach().cpu().reshape((100, 100))
            ax.contour(X, Y + y, Z, [y], zdir="y", alpha=alpha, cmap=cmap)

        for x in B:  # plot contours in the YZ plane
            Y, Z = A1, A2
            X = fn(*make_xyz(x, Y, Z)).w[:, 0].detach().cpu().reshape((100, 100))
            ax.contour(X + x, Y, Z, [x], zdir="x", alpha=alpha, cmap=cmap)

        # must set plot limits because the contour will likely extend
        # way beyond the displayed level.  Otherwise matplotlib extends the plot limits
        # to encompass all values in the contour.
        ax.set_zlim3d(zmin, zmax)
        ax.set_xlim3d(xmin, xmax)
        ax.set_ylim3d(ymin, ymax)

    ax.set_aspect("equal")

    plt.gca().set_rasterization_zorder(-1)

    if iter:
        plt.savefig(f"{iter}_transport_3D_level_set_large_adaptive.pdf")

    # plt.show()


def plot_results(f, sampler=None, T=0, n_visu=85, iter=None):
    plot_implicit(f)


def save_to_csv(f, file_name, sampler=None, n_visu=32):
    import pandas as pd

    if sampler is None:
        x = torch.linspace(0.2, 0.9, n_visu)
        y = torch.linspace(0.17, 0.77, n_visu)
        z = torch.linspace(0.17, 0.77, n_visu)

        X, Y, Z = torch.meshgrid(x, y, z)

        x_ = LabelTensor(torch.stack((X.flatten(), Y.flatten(), Z.flatten()), dim=1))
        mu = LabelTensor(torch.zeros((x_.shape[0], 0)))

    else:
        x_, mu = sampler(n_visu**3)

    y = f(x_, mu).w[:, 0].detach().cpu()

    df = pd.DataFrame(
        {
            "x": x_.x[:, 0].detach().cpu().numpy(),
            "y": x_.x[:, 1].detach().cpu().numpy(),
            "z": x_.x[:, 2].detach().cpu().numpy(),
            "value": y.numpy(),
        }
    )

    df.to_csv(f"{file_name}", index=False)


def plot_slice(f, direction, ax, n_visu=768, border=0.02, center=0.35, radius=0.15):
    shift = radius + border

    x = torch.linspace(center - shift, center + shift, n_visu)
    y = torch.linspace(center - shift, center + shift, n_visu)
    X, Y = torch.meshgrid(x, y)
    z = center * torch.ones_like(X)

    if direction == "x":
        x_ = LabelTensor(torch.stack((z.flatten(), X.flatten(), Y.flatten()), dim=1))
        x_axis = "y"
        y_axis = "z"
    elif direction == "y":
        x_ = LabelTensor(torch.stack((X.flatten(), z.flatten(), Y.flatten()), dim=1))
        x_axis = "x"
        y_axis = "z"
    elif direction == "z":
        x_ = LabelTensor(torch.stack((X.flatten(), Y.flatten(), z.flatten()), dim=1))
        x_axis = "x"
        y_axis = "y"
    else:
        raise ValueError("Invalid direction argument")

    mu = LabelTensor(torch.zeros((x_.shape[0], 0)))

    Z = f(x_, mu).w[:, 0].detach().cpu().reshape((n_visu, n_visu))

    X = X.detach().cpu()
    Y = Y.detach().cpu()
    Z = Z.detach().cpu()

    ax.contour(X, Y, Z, levels=[0], colors="black")

    v = ax.collections[0].get_paths()[0].vertices
    x_approx, y_approx = v[:, 0], v[:, 1]

    Z_true = (X - 0.35) ** 2 + (Y - 0.35) ** 2 - 0.15**2

    ax.contour(X, Y, Z_true, levels=[0], colors="red")

    v = ax.collections[1].get_paths()[0].vertices
    x_true, y_true = v[:, 0], v[:, 1]

    ax.set_title(f"Slice in the {direction}-direction")
    ax.set_xlabel(f"{x_axis}")
    ax.set_ylabel(f"{y_axis}")

    ax.set_aspect("equal")

    return x_approx, y_approx, x_true, y_true


def plot_slices(f, n_visu=768, save_to_csv=False):
    fig, ax = plt.subplots(1, 3, figsize=(15, 5))

    def save(direction):
        if save_to_csv:
            np.savetxt(
                f"approx_slice_in_{direction}_direction.csv",
                np.stack((x_approx, y_approx)).T,
                delimiter=",",
                header="x_approx,y_approx",
                comments="",
            )
            np.savetxt(
                f"true_slice_in_{direction}_direction.csv",
                np.stack((x_true, y_true)).T,
                delimiter=",",
                header="x_true,y_true",
                comments="",
            )

    x_approx, y_approx, x_true, y_true = plot_slice(f, "x", ax[0], n_visu)
    save("x")

    x_approx, y_approx, x_true, y_true = plot_slice(f, "y", ax[1], n_visu)
    save("y")

    x_approx, y_approx, x_true, y_true = plot_slice(f, "z", ax[2], n_visu)
    save("z")

    plt.show()


# %%


def plot_intermediate(file_number=10, n_visu=96, color="gray"):
    if file_number == 0:
        file_name = "ini_transport_3D_level_set_larger_truly_adaptive"
    else:
        file_name = f"{file_number}_transport_3D_level_set_larger_truly_adaptive_large_time_step"

    scheme, sampler = load_and_postprocess_projector(file_name)
    f = scheme.pde.space.evaluate

    # x_, mu = sampler.sample(n_visu**3)

    x = torch.linspace(0.2, 0.9, n_visu)
    y = torch.linspace(0.17, 0.77, n_visu)
    z = torch.linspace(0.17, 0.77, n_visu)

    X, Y, Z = torch.meshgrid(x, y, z, indexing="ij")

    x_ = LabelTensor(torch.stack((X.flatten(), Y.flatten(), Z.flatten()), dim=1))
    mu = LabelTensor(torch.zeros((x_.shape[0], 0)))

    y = f(x_, mu).w[:, 0].detach().cpu().numpy().reshape((n_visu, n_visu, n_visu))

    # Make the contour, marching cubes
    marchCubeSpace = 1.0 / n_visu
    verts, faces, normals, values = measure.marching_cubes(
        y, 0, spacing=(marchCubeSpace, marchCubeSpace, marchCubeSpace)
    )

    # Create Ploy3D and set up a light source
    mesh = Poly3DCollection(verts[faces], alpha=1.0)
    ls = LightSource(azdeg=315.0, altdeg=15.0)

    # First change - normals are per vertex, so I made it per face.
    normalsarray = np.array(
        [
            np.array(
                (
                    np.sum(normals[face[:], 0] / 3),
                    np.sum(normals[face[:], 1] / 3),
                    np.sum(normals[face[:], 2] / 3),
                )
                / np.sqrt(
                    np.sum(normals[face[:], 0] / 3) ** 2
                    + np.sum(normals[face[:], 1] / 3) ** 2
                    + np.sum(normals[face[:], 2] / 3) ** 2
                )
            )
            for face in faces
        ]
    )

    # Next this is more asthetic, but it prevents the shadows of the image being too dark. (linear interpolation to correct)
    min = np.min(ls.shade_normals(normalsarray, fraction=1.0))  # min shade value
    max = np.max(ls.shade_normals(normalsarray, fraction=1.0))  # max shade value
    diff = max - min
    newMin = 0.3
    newMax = 0.95
    newdiff = newMax - newMin

    # Using a constant color, put in desired RGB values here.
    if color == "gray":
        colourRGB = np.array((0.75, 0.75, 0.75, 1))
    else:
        colourRGB = np.array((228 / 255, 26 / 255, 28 / 255, 1))

    # The correct shading for shadows are now applied. Use the face normals and light orientation to generate a shading value and apply to the RGB colors for each face.
    rgbNew = np.array(
        [
            colourRGB * (newMin + newdiff * ((shade - min) / diff))
            for shade in ls.shade_normals(normalsarray, fraction=1.0)
        ]
    )

    # Apply color to face
    mesh.set_facecolor(rgbNew)
    mesh.set_edgecolor(rgbNew)

    # Plot
    fig = plt.figure(figsize=(7.5, 7.5))
    ax = fig.add_subplot(111, projection="3d")
    ax.view_init(elev=-45, azim=-90)
    ax.add_collection3d(mesh)

    ax.set_xlim3d(0.15, 0.8)
    ax.set_ylim3d(0, 0.94)
    ax.set_zlim3d(0.17, 0.77)

    plt.axis("off")
    plt.grid(b=None)

    plt.tight_layout()

    plt.savefig(
        f"3D_plots/three_dimensional_level_set_{file_number}.png", transparent=True
    )

    plt.show()
    plt.close()


# %%


class NonUniformSampler(DomainSampler):
    def __init__(self, f, domain, **kwargs):
        super().__init__(domain, **kwargs)
        self.f = f

    def sample_around_zeroes(self, n, sigma=0.05, keep=0.75):
        candidates = LabelTensor(self.vol_sampler.sample(10 * n))
        parameters = LabelTensor(torch.zeros((10 * n, 0)))
        f_values = self.f(candidates, parameters).w[:, 0]
        weights = torch.exp(-(f_values**2) / (2 * sigma**2))
        return candidates[weights > keep].x

    def sample_new_points(self, n: int):
        """Samples `n` new points from the volumetric domain and labels them by subdomains.

        Args:
            n (int): Number of points to sample.

        Returns:
            ScimbaTensor: A tensor of sampled points with their subdomain labels.
        """
        pts = torch.zeros((0, 3))

        iter = 0
        max_iter = 100

        while pts.shape[0] < 3 * n // 4 and iter < max_iter:
            pts_ = self.sample_around_zeroes(n, sigma=0.02)
            pts = torch.cat((pts, pts_), dim=0)
            iter += 1

        pts_ = self.sample_around_zeroes(n // 4, sigma=0.15)
        pts = torch.cat((pts, pts_), dim=0)

        pts = pts[:n]

        labels = torch.zeros(n, dtype=int)

        # Label points based on subdomains
        for i, subdomain in enumerate(self.domain.list_subdomains):
            condition = subdomain.is_inside(pts)
            labels[condition] = i + 1

        # Wrap the points and labels in a ScimbaTensor
        data = LabelTensor(pts, labels)
        data.x.requires_grad_()
        return data


def load_and_postprocess_projector(file_name="final_level_set") -> NeuralSemiLagrangian:
    torch.random.manual_seed(0)

    domain_x = Cube3D(DEFAULT_BOUNDS, is_main_domain=True)

    sampler_unif = TensorizedSampler(
        [
            DomainSampler(domain_x),
            UniformParametricSampler([]),
        ]
    )

    if "larger" in file_name:
        layer_sizes = [35, 50, 50, 35]
    else:
        layer_sizes = [35, 50, 35]

    space = NNxSpace(1, 0, GenericMLP, domain_x, sampler_unif, layer_sizes=layer_sizes)
    pde = AdvectionReactionDiffusionDirichletStrongForm(
        # space, a=a, u0=f_ini, constant_advection=False, zero_diffusion=True
        space,
        u0=f_ini,
        constant_advection=True,
        zero_diffusion=True,
    )
    characteristic = Characteristic(pde, periodic=True)

    sampler = TensorizedSampler(
        [
            NonUniformSampler(space.evaluate, domain_x),
            # DomainSampler(domain_x),
            UniformParametricSampler([]),
        ]
    )
    space.integrator = sampler
    pde.space = space
    characteristic = Characteristic(pde, periodic=True)

    projector = NaturalGradientProjector(space, rhs=f_ini)
    projector.load(file_name)

    scheme = NeuralSemiLagrangian(characteristic, projector)

    return scheme, sampler


# %%

if __name__ == "__main__":
    # print("#" * 80)
    # print("INITIAL TEST")
    # print("#" * 80)

    # scheme, sampler = load_and_postprocess_projector(
    #     "final_3D_level_set_large_adaptive"
    # )
    # plot_slices(scheme.pde.space.evaluate)

    # print("#" * 80)
    # print("LARGER NETWORK")
    # print("#" * 80)

    # scheme, sampler = load_and_postprocess_projector(
    #     "final_3D_level_set_larger_adaptive"
    # )
    # plot_slices(scheme.pde.space.evaluate)

    # print("#" * 80)
    # print("LARGER NETWORK, ADAPTIVE SAMPLING")
    # print("#" * 80)

    # scheme, sampler = load_and_postprocess_projector(
    #     "final_3D_level_set_larger_truly_adaptive"
    # )
    # plot_slices(scheme.pde.space.evaluate)

    print("#" * 80)
    print("LARGER NETWORK, ADAPTIVE SAMPLING, LARGE TIME STEP")
    print("#" * 80)

    scheme, sampler = load_and_postprocess_projector(
        "final_3D_level_set_larger_truly_adaptive_large_time_step"
    )
    plot_slices(scheme.pde.space.evaluate, save_to_csv=True)

# %%

if __name__ == "__main__":

    def plot_intermediate_wrapper(i):
        print(f"{i} ", flush=True)
        plot_intermediate(file_number=i, n_visu=256)

    for i in range(11):
        plot_intermediate_wrapper(i)

# %%