"""
pycutwm.plot_tools.py
------------

Main routines for plot the package results.

Functions:
    - plot_signal_intensity: Plot a single electric field.
    - plot_XY_profile: Plot the XY profile of an electric field
    - plot_fields_from_slices: View the propagation of enabled fields in config.json
    - make_field_propagation_gif: Generates a GIF showing the progression of 
        field intensity (x-y plane) at each slice along the crystal
        for each field enabled in config.json. Each frame is a 
        horizontal row of 2D intensity maps (one per field).

Dependencies:
    - Python >= 3.7
    - numpy, imageio, json, matplotlib
    - nvidia-smi in PATH (for compile_cutwm)
    - Custom utils (pycutwm.utils.load_config, save_config)

Author: Alfredo Daniel Sánchez Rossi
Date: 2025-25-12
License: MIT
"""


import numpy as np
import matplotlib.pyplot as plt
import os
import json
import imageio
import h5py

from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas



def plot_intensity_cut_YZ_from_h5(
    files_folder,
    config_path,
    n_dict,
    cmap='inferno',
    image_dir="intensity_cut_YZ/",
    image_name="intensity_profile",
    image_extension=".png",
    y_limits=None,     # e.g. (-500, 500)
    z_limits=None      # e.g. (0, 20000)
):
    # Read config.json
    with open(config_path, "r") as f:
        cfg = json.load(f)
    save_mode = cfg.get("save_mode", {})
    grid_points = cfg.get("grid", {}).get("grid_points", {})
    NX = int(grid_points.get("NX", 32))
    NY = int(grid_points.get("NY", 32))
    NZ = int(grid_points.get("NZ", 1))
    npasses = cfg["mul_pass_scheme"]["npasses"]
    NZ = NZ*npasses

    LY  = cfg.get("crystal", {}).get("dimensions", {}).get("LY", 1000.0)
    LZ  = cfg.get("crystal", {}).get("dimensions", {}).get("Lcr", 20000.0)

    # Qué campos graficar
    field_list = [field for field in ('pump', 'signal', 'idler') if save_mode.get(f"save_{field}", False)]
    if not field_list:
        print("No fields are selected for graphing according to config.json.")
        return

    # Leer archivos H5 de cada campo
    E_fields = {}
    for field in field_list:
        h5_path = os.path.join(files_folder, f"{field}_inside_crystal.h5")
        if not os.path.isfile(h5_path):
            raise FileNotFoundError(f"{h5_path} not found.")
        with h5py.File(h5_path, "r") as f:
            real = f["real"][:]
            imag = f["imag"][:]
            # real, imag tienen shape (NY, NX, NZ)
            E_cmplx = real + 1j * imag
            # Reordenar a (NZ, NY, NX)
            E_cmplx = np.transpose(E_cmplx, (2, 0, 1)) # (NZ, NY, NX)
        E_fields[field] = E_cmplx

    # Physical axis in microns
    dy, dz = LY/(NY-1), LZ/(NZ-1)
    y = np.linspace(-1, 1, NY)*LY/2-0.5*dy
    z_pos = np.linspace(0, 1, NZ)*LZ

    # Define limit
    if y_limits is None:
        y_limits = (y[0], y[-1])
    if z_limits is None:
        z_limits = (z_pos[0], z_pos[-1])

    # Plot
    x_idx = NX // 2
    fig, axes = plt.subplots(len(field_list), 1, figsize=(5, 5), constrained_layout=True)
    if len(field_list) == 1:
        axes = [axes]

    for i, field in enumerate(field_list):
        n = n_dict[field]
        E = E_fields[field]
        # Physical constans
        C0 = 299792458 * 1e6 / 1e12        # [μm/ps]
        EPS0 = 8.8541878128e-12 * 1e12 / 1e6 # [W·ps/V²·μm]
        intensity_Wum2 = C0*EPS0*n*np.abs(E[:, :, x_idx])**2 / 2.0
        intensity_MWcm2 = intensity_Wum2 * 1e8 * 1e-6

        ax = axes[i]
        im = ax.imshow(
            intensity_MWcm2.T,
            origin='lower',
            aspect='auto',
            extent=[z_pos[0], z_pos[-1], y[0], y[-1]],  # X: Z (μm), Y: Y (μm)
            cmap=cmap
        )
        ax.set_xlim(z_limits)
        ax.set_ylim(y_limits)
        ax.set_title(f"{field.capitalize()} field")
        ax.set_xlabel('z (μm)')
        ax.set_ylabel('y (μm)')
        cb = fig.colorbar(im, ax=ax, fraction=0.045)
        cb.set_label("Intensity (MW/cm²)")

    plt.suptitle("Fields propagation (cut at x=NX/2)")

    # Ensure output directory exists
    if image_dir and not os.path.exists(image_dir):
        os.makedirs(image_dir)
    fig.savefig(os.path.join(image_dir, image_name+image_extension), dpi=150, bbox_inches='tight')
    print(f"Figure saved in {os.path.join(image_dir, image_name+image_extension)}")
    return


def field_prop_in_time_make_gif(
    files_folder,
    config_path,
    n_dict,           # {'pump': n_p, 'signal': n_s, 'idler': n_i}
    gif_path,
    cmap='inferno',
    fps=10,
    vmin=None,
    vmax=None
):
    """
    Generates a GIF showing the progression of field intensity (x-y plane) at each slice along the crystal
    for each field enabled in config.json. Each frame is a horizontal row of 2D intensity maps (one per field).

    Args:
        files_folder (str): Path to the folder containing the .dat files.
        config_path (str): Path to config.json.
        n_dict (dict): Refractive indices for each field. Example: {'pump': 2.23, 'signal': 2.21, 'idler': 2.19}
        gif_path (str): Output path for the GIF file (should end with .gif).
        cmap (str): Colormap to use for intensity images.
        fps (int): Frames per second for the GIF animation.
        vmin, vmax (float or None): Set fixed min/max for color scale. If None, scale automatically for each field.
    """
    print(f"Generating GIF in folder: {gif_path}. Please wait a few seconds...")

    # --- Load config.json ---
    with open(config_path, "r") as f:
        cfg = json.load(f)

    save_mode = cfg.get("save_mode", {})
    grid_points = cfg.get("grid", {}).get("grid_points", {})
    NX = int(grid_points.get("NX", 32))
    NY = int(grid_points.get("NY", 32))
    NT = int(grid_points.get("NT", 1))
    crystal = cfg.get("crystal", {})
    Lcr = float(crystal.get("dimensions", {}).get("Lcr", 10000.0))  # in micrometers
    twin = cfg.get("fields",{}).get("time_freq_vect",{}).get("time_window",0.2)

    # --- Which fields to plot? ---
    field_list = [field for field in ('pump', 'signal', 'idler') if save_mode.get(f"save_{field}", False)]
    if not field_list:
        print("No fields selected for plotting in config.json.")
        return

    dt = twin / (NT - 1) if NT > 1 else 0

    # --- Precompute color limits for all fields over all slices (for consistent colorbar) ---
    if vmin is None or vmax is None:
        all_intensities = {field: [] for field in field_list}
        for tt in range(NT):
            for field in field_list:
                fname_real = os.path.join(files_folder, f"{field}_output_XY_timeslice_{tt}_r.dat")
                fname_imag = os.path.join(files_folder, f"{field}_output_XY_timeslice_{tt}_i.dat")

                real = np.loadtxt(fname_real)
                imag = np.loadtxt(fname_imag)
                E = real + 1j * imag
                # Physical constants in user units:
                C0 = 299792458 * 1e6 / 1e12        # Speed of light in vacuum [μm/ps]
                EPS0 = 8.8541878128e-12 * 1e12 / 1e6 # Vacuum permittivity [W·ps/V²·μm]
                intensity_Wum2 = C0*EPS0*n_dict[field]*np.abs(E)**2 / 2.0  # (NY, NX)
                intensity_Wcm2 = intensity_Wum2 * 1e8
                all_intensities[field].append(intensity_Wcm2)
        # Flatten and find global min/max for all fields
        flat_all = np.concatenate([np.ravel(arr) for field in field_list for arr in all_intensities[field]])
        vmin = float(np.nanmin(flat_all))
        vmax = float(np.nanmax(flat_all))

    # --- Ensure output directory exists ---
    output_dir = os.path.dirname(gif_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # --- Create and save frames ---
    images = []
    for tt in range(NT):
        # Create figure and attach Agg canvas
        fig = Figure(figsize=(6*len(field_list), 5), constrained_layout=True)
        FigureCanvas(fig)
        axes = fig.subplots(1, len(field_list))
        if len(field_list) == 1:
            axes = [axes]

        for i, field in enumerate(field_list):
            fname_real = os.path.join(files_folder, f"{field}_output_XY_timeslice_{tt}_r.dat")
            fname_imag = os.path.join(files_folder, f"{field}_output_XY_timeslice_{tt}_i.dat")
            real = np.loadtxt(fname_real)
            imag = np.loadtxt(fname_imag)
            E = real + 1j * imag
            # Physical constants in user units:
            C0 = 299792458 * 1e6 / 1e12        # Speed of light in vacuum [μm/ps]
            EPS0 = 8.8541878128e-12 * 1e12 / 1e6 # Vacuum permittivity [W·ps/V²·μm]
            intensity_Wum2 = C0*EPS0*n_dict[field]*np.abs(E)**2 / 2.0  # shape (NY, NX)
            intensity_Wcm2 = intensity_Wum2 * 1e8

            ax = axes[i]
            im = ax.imshow(
                intensity_Wcm2,
                origin='lower',
                aspect='auto',
                extent=[0, NX-1, 0, NY-1],
                cmap=cmap,
                vmin=vmin,
                vmax=vmax
            )
            ax.set_title(f"{field.capitalize()} field")
            ax.set_xlabel('x')
            ax.set_ylabel('y')
            cb = fig.colorbar(im, ax=ax, fraction=0.045)
            cb.set_label("Intensity (W/cm²)")

        t_ps = tt * dt
        fig.suptitle(f"time = {t_ps:.1f} ps")

        # Draw and convert figure to numpy RGB array
        fig.canvas.draw()
        image = np.asarray(fig.canvas.buffer_rgba())[:, :, :3]
        images.append(image)
        plt.close('all')
    
    # --- Save GIF ---
    # Ensure output directory exists
    output_dir = os.path.dirname(gif_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    imageio.mimsave(gif_path, images, fps=fps)
    print(f"GIF saved to {gif_path}")