#!/usr/bin/env python3
# ----------------------------------------------------------------------
# Copyright (C) 2025 Tristan Hoellinger
# Distributed under the GNU General Public License v3.0 (GPLv3).
# See the LICENSE file in the root directory for details.
# SPDX-License-Identifier: GPL-3.0-or-later
# ----------------------------------------------------------------------

__author__ = "Tristan Hoellinger"
__version__ = "0.1.0"
__date__ = "2025"
__license__ = "GPLv3"

"""Plotting utilities for the WIP-P3M package."""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.colors import TwoSlopeNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cmocean.cm as cm
import plotly.graph_objects as go

from wip3m.plot_params import *  # type: ignore

# Configure global plotting settings
setup_plotting()

fs = GLOBAL_FS_SMALL
fs_titles = GLOBAL_FS_LARGE
cols = COLOUR_LIST
cmap = cm.thermal  # type: ignore

def matplotlib_to_plotly(cmap, n=255):
    """Convert a matplotlib colormap to a Plotly colorscale."""
    colorscale = []
    for i in range(n):
        norm = i / (n - 1)
        r, g, b, _ = cmap(norm)
        colorscale.append([norm, f'rgb({int(r*255)}, {int(g*255)}, {int(b*255)})'])
    return colorscale

thermal_plotly = matplotlib_to_plotly(cmap)

def plotly_3d(
    field: np.ndarray,
    size: int = 128,
    L: float = None,  # type: ignore
    colormap: str | list = "RdYlBu",
    limits: str = "max",
) -> go.Figure:
    """
    Create an interactive 3D plot of volume slices using Plotly.

    Parameters
    ----------
    field : array-like
        3D data field to visualise.
    size : int, optional
        Size of the field along one dimension. Default is 128.
    L : float, optional
        Physical size of the field in Mpc/h. Used for axis labels only.
    colormap : str, optional
        Colour map for visualisation. Default is 'RdYlBu'.
    limits : str, optional
        Colour scale limits ('max', 'truncate', or 'default'). Default
        is 'max'.

    Returns
    -------
    go.Figure
        Plotly figure object.
    """
    volume = field.T
    rows, cols = volume[0].shape

    # Define colour scale limits
    if limits == "max":
        maxcol = np.max(np.abs(volume))
        mincol = -maxcol
    elif limits == "truncate":
        maxcol = min(np.max(-volume), np.max(volume))
        mincol = -maxcol
    else:
        maxcol = np.max(volume)
        mincol = np.min(volume)
    midcol = np.mean(volume)

    # Generate frames for the animation
    nb_frames = size
    frames = [
        go.Frame(
            data=go.Surface(
                z=(size - k) * np.ones((rows, cols)),
                surfacecolor=np.flipud(volume[cols - 1 - k]),
                cmin=mincol,
                cmid=midcol,
                cmax=maxcol,
            ),
            name=str(k),  # Frames must be named for proper animation
        )
        for k in range(nb_frames)
    ]

    # Initial plot configuration
    fig = go.Figure(
        frames=frames,
        data=go.Surface(
            z=size * np.ones((rows, cols)),
            surfacecolor=np.flipud(volume[cols // 2]),
            colorscale=colormap,
            cmin=mincol,
            cmid=midcol,
            cmax=maxcol,
            colorbar=dict(thickness=20, ticklen=4),
        ),
    )

    def frame_args(duration):
        """Helper function to set animation frame arguments."""
        return {
            "frame": {"duration": duration},
            "mode": "immediate",
            "fromcurrent": True,
            "transition": {"duration": duration, "easing": "linear"},
        }

    # Add animation slider
    sliders = [
        {
            "pad": {"b": 10, "t": 60},
            "len": 0.9,
            "x": 0.1,
            "y": 0,
            "steps": [
                {
                    "args": [[f.name], frame_args(0)],
                    "label": str(k),
                    "method": "animate",
                }
                for k, f in enumerate(fig.frames)
            ],
        }
    ]

    # Configure layout with or without physical size
    layout_config = dict(
        title="Slices in density field",
        width=600,
        height=600,
        scene=dict(
            zaxis=dict(range=[0, size - 1], autorange=False),
            xaxis_title="x [Mpc/h]",
            yaxis_title="y [Mpc/h]",
            zaxis_title="z [Mpc/h]",
            aspectratio=dict(x=1, y=1, z=1),
        ),
        updatemenus=[
            {
                "buttons": [
                    {"args": [None, frame_args(50)], "label": "&#9654;", "method": "animate"},
                    {"args": [[None], frame_args(0)], "label": "&#9724;", "method": "animate"},
                ],
                "direction": "left",
                "pad": {"r": 10, "t": 70},
                "type": "buttons",
                "x": 0.1,
                "y": 0,
            }
        ],
        sliders=sliders,
    )
    if L is not None:
        layout_config["scene"]["xaxis"] = dict(
            ticktext=[0, L / 2, L],
            tickvals=[0, size / 2, size],
            title="x [Mpc/h]",
        )
        layout_config["scene"]["yaxis"] = dict(
            ticktext=[0, L / 2, L],
            tickvals=[0, size / 2, size],
            title="y [Mpc/h]",
        )
        layout_config["scene"]["zaxis"]["ticktext"] = [0, L / 2, L]
        layout_config["scene"]["zaxis"]["tickvals"] = [0, size / 2, size]

    fig.update_layout(**layout_config)
    return fig

def clear_large_plot(fig):
    """
    Clear a figure to free up memory.

    Parameters
    ----------
    fig : matplotlib.figure.Figure
        The figure to clear.
    """
    from IPython.display import clear_output

    del fig
    clear_output()

def load_force_diagnostic(filename):
    """
    Load force diagnostic data.

    Parameters
    ----------
    filename : str
        Path to the CSV file written by `run_force_subtraction_test`
        from p3m.c.

    Returns
    -------
    r : ndarray
        Bin-centred distances.
    fmag : ndarray
        Total force magnitudes |f_after - f_before|.
    data : recarray
        Full structured array with all columns.
    """
    data = np.genfromtxt(filename, delimiter=",", names=True)
    # Append the magnitude of force difference to the data
    fmag = np.sqrt(data["fx"]**2 + data["fy"]**2 + data["fz"]**2)
    data = np.lib.recfunctions.append_fields(data, "fmag", fmag, usemask=False, asrecarray=True)

    return data["distance"], data["fmag"], data

def plot_force_distance(r, fmag, f_max=1e-1, a=None, title=None, save_path=None):
    """
    Plot total force magnitude vs distance and a theoretical 1/r² profile.

    Parameters
    ----------
    r : array_like
        Distance values (bin centers or actual distances).
    fmag : array_like
        Corresponding force magnitudes.
    f_max : float, optional
        Maximum force magnitude for including points in the plot (default: 1e-1).
    a : float or None, optional
        If provided, the theoretical prefactor for the 1/r² profile.
    title : str, optional
        Plot title.
    save_path : str or None, optional
        If provided, path to save the figure as a PDF.
    """
    from scipy.optimize import curve_fit

    # Select points to include in the plot
    mask = fmag < f_max
    rs = r[mask]
    fmag_select = fmag[mask]

    fig, ax = plt.subplots()
    ax.scatter(rs, fmag_select, s=10, alpha=0.8, label="Particle/hole pairs", color="tab:blue")

    # If a is provided, plot the theoretical curve
    if a is not None:
        def inverse_square(r, a):
            return a / r**2
        r = np.linspace(np.min(rs), np.max(rs), 300)
        f_th = inverse_square(r, a)
        ax.plot(r, f_th, "r-", label=r"$\propto 1/r^2$")
    
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel(r"Distance $r$ [Mpc/$h$]")
    ax.set_ylabel(r"Force magnitude [code units]")
    if title is not None:
        ax.set_title(title)
    ax.legend()
    ax.grid(True, which="both", ls=":", lw=0.5)

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path)
        print(f"Figure saved to: {save_path}")
    plt.show()

def plot_force_distance_comparison(rr, ff, ll, L, Np, Npm, ss=None, a=None, title=None, save_path=None):
    """
    Plot force magnitude vs distance.

    Parameters
    ----------
    rr : list of array_like
        List of distance arrays.
    ff : list of array_like
        List of force magnitude arrays corresponding to each distance array.
    ll : list of str
        List of labels for each dataset.
    ss : list of str or None, optional
        List of symbols for each dataset (default: None).
    L : float
        Physical size of the field in Mpc/h.
    Np : int
        Number of particles per dimension.
    Npm : int
        Number of PM cells per dimension.
    a : float or None, optional
        If provided, the theoretical prefactor for the 1/r² profile.
    title : str, optional
        Plot title.
    save_path : str or None, optional
        If provided, path to save the figure as a PDF.
    """
    fig, ax = plt.subplots()

    colours = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    handles1 = []

    if ss is None:
        ss = ["o"] * len(rr)
    for i, (r, f, label, symbol) in enumerate(zip(rr, ff, ll, ss)):
        scatter = ax.scatter(r, f, alpha=0.8, s=20, label=label,
                             color=colours[i % len(colours)], marker=symbol)
        handles1.append(scatter)

    loc1="lower right"

    # Theoretical curve
    theory_line = None
    if a is not None and len(rr) > 0:
        loc1="upper right"
        def inverse_square(r, a):
            return a / r**2
        r_min = min(np.min(r) for r in rr)
        r_max = max(np.max(r) for r in rr)
        r_plot = np.linspace(r_min, r_max, 300)
        f_th = inverse_square(r_plot, a)
        theory_line, = ax.plot(r_plot, f_th, "k-", label=r"theory $\propto 1/r^2$")
        handles1.append(theory_line)

    # Characteristic vertical reference scales
    nyquist = 2 * L / Npm
    epsilon = 0.03 * L / Np
    xs = 1.25 * L / Npm
    xr = 4.5 * xs
    line1 = ax.axvline(x=nyquist, color="black", linestyle="-", lw=1, label="Nyquist (PM grid)")
    line2 = ax.axvline(x=2*epsilon, color="gray", linestyle="--", lw=2, label=r"Particle length $2\epsilon$")
    line3 = ax.axvline(x=xs, color="gray", linestyle="-.", lw=2, label=r"Split scale $x_s$")
    line4 = ax.axvline(x=xr, color="gray", linestyle=":", lw=2, label=r"Short-range reach $x_r$")
    print(f"Nyquist: {nyquist:.2f} Mpc/h")
    print(f"Particle length: {2*epsilon:.2f} Mpc/h")
    print(f"Split scale: {xs:.2f} Mpc/h")
    print(f"Short-range reach: {xr:.2f} Mpc/h")

    # Set log-log axes and labels
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel(r"Distance $r$ [Mpc/$h$]")
    ax.set_ylabel(r"Force magnitude [code units]")
    if title is not None:
        ax.set_title(title)

    # Legend for data
    legend1 = ax.legend(handles=handles1, loc=loc1, frameon=True, fontsize=GLOBAL_FS_TINY)
    ax.add_artist(legend1)

    # Legend for vertical lines
    handles2 = [line1, line2, line3, line4]
    fig.subplots_adjust(bottom=0.3)
    legend2 = fig.legend(handles=handles2, loc='lower center', ncol=2, frameon=False)

    ax.grid(True, which="both", ls=":", lw=0.5)

    if save_path:
        plt.savefig(save_path)
        print(f"Figure saved to: {save_path}")
    plt.show()