"""
pycutwm.gif_from_H5.py
------------

Main routines for plotting gif in (3+1)D problems.

Functions:
    - create_gif: Create a density plot GIF for a single electric field profile
    from a HDF5 file. Each frame corresponds to a time slice.
    - create_gif_3D: Create a 3D plot GIF for a single electric field profile
    from a HDF5 file. Each frame corresponds to a time slice.


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


import numpy as np
import h5py
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D  # needed for 3D
from typing import Optional, Dict, Any, Tuple

# === Matplotlib setup (optional) ===
plt.rcParams.update({
    "font.family": "serif",
    "mathtext.fontset": "cm",
    "font.size": 14
})

# 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]


def create_gif(h5_path: str, gif_name: str, field: str, fps=10, n_dict: Dict[str, float]=None):
    """
    Create a animated density plot GIF from data stored in an HDF5 file.
    Parameters:
    - h5_path: Path to the HDF5 file containing the electric field data.
    - gif_name: Name of the output GIF file.
    - field: Field to display ("pump", "signal", or "idler").
    - fps: Frames per second for the animation.
    - n_dict: Dictionary with refractive indices for each field.
    """

    # === 1. Read the entire field from the HDF5 and calculate the global maximum ===
    with h5py.File(h5_path, "r") as f:
        Er = f["real"][:]  # (NT, NY, NX)
        Ei = f["imag"][:]  # (NT, NY, NX)

    E_abs_all = np.abs(Er + 1j * Ei)  # (NT, NY, NX)
    
    n = n_dict[field]
    NT, ny, nx = E_abs_all.shape

    global_max = E_abs_all.max()
    iglobal_max = 0.5 * n * EPS0 * C0 * np.abs(global_max)**2
    if iglobal_max == 0:
        iglobal_max = 1.0  # avoid vmax = 0

    print(f"{NT} frames (times) were found. Global maximum I(x,y) = {iglobal_max:.3e}")


    # === 2. Prepare figure for animation ===
    fig, ax = plt.subplots(figsize=(6, 5))

    E_abs0 = E_abs_all[0, :, :]
    I = 0.5 * n * EPS0 * C0 * np.abs(E_abs0)**2
    
    im = ax.imshow(
        I,
        cmap='inferno',
        origin='lower',
        vmin=0.0,
        vmax=iglobal_max
    )

    cbar = fig.colorbar(im, ax=ax, label=r'I(x, y) (W/$\mu$m²)')

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title('Intensity profile (frame = 0)')

    plt.tight_layout()

    # === 3. Update function per frame ===
    def update(frame_index):
        E_abs = E_abs_all[frame_index, :, :]
        I = 0.5 * n * EPS0 * C0 * np.abs(E_abs)**2
        im.set_data(I)
        ax.set_title(f'Intensity profile (frame = {frame_index})')
        return [im]

    # === 4. Create and save GIF ===
    anim = FuncAnimation(
        fig,
        update,
        frames=NT,
        interval=1000 / fps,
        blit=False
    )

    anim.save(gif_name, writer='pillow', fps=fps)
    plt.close(fig)

    print(f"GIF guardado en: {gif_name}")


def create_gif_3D(h5_path: str, gif_name: str, field: str, 
                  fps=10, n_dict: Dict[str, float]=None,
                  elev=0, azim=-60):
    """
    Create a 3D animated GIF from data stored in an HDF5 file.
    Parameters:

    - h5_path: Path to the HDF5 file containing the electric field data.
    - gif_name: Name of the output GIF file.
    - field: Field to display ("pump", "signal", or "idler").
    - fps: Frames per second for the animation.
    - n_dict: Dictionary with refractive indices for each field.
    - elev: Elevation angle for the 3D view.
    - azim: Azimuth angle for the 3D view.
    """

    # === 1. Read the entire field from the HDF5 and calculate the global maximum ===
    with h5py.File(h5_path, "r") as f:
        Er = f["real"][:]  # (NT, NY, NX)
        Ei = f["imag"][:]  # (NT, NY, NX)
    
    n = n_dict[field]
    
    E_abs_all = np.abs(Er + 1j * Ei)  # (NT, NY, NX)
    NT, ny, nx = E_abs_all.shape

    # NT=2
    global_max = E_abs_all.max()
    iglobal_max = 0.5 * n * EPS0 * C0 * np.abs(global_max)**2
    if iglobal_max == 0:
        iglobal_max = 1.0  # avoid vmax = 0
        

    print(f"{NT} frames (times) were found. Global maximum I(x,y) = {iglobal_max:.3e}")

    # === 2. Prepare 3D figure for animation ===
    fig = plt.figure(figsize=(7, 6))
    ax = fig.add_subplot(111, projection='3d')

    # Spatial grid
    X, Y = np.meshgrid(np.arange(nx), np.arange(ny))

    # First frame
    E_abs0 = E_abs_all[0, :, :]
    I = 0.5 * n * EPS0 * C0 * np.abs(E_abs0)**2

    surface = [ax.plot_surface(
        X, Y, I,
        cmap='inferno',
        vmin=0.0,
        vmax=iglobal_max,
        linewidth=0,
        antialiased=False
    )]

    ax.set_xlim(0, nx - 1)
    ax.set_ylim(0, ny - 1)
    ax.set_zlim(0, iglobal_max)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel(r'I(x, y) (W/$\mu$m²)')
    ax.set_title('Intensity profile (frame = 0)')

    ax.view_init(elev, azim)

    cbar = fig.colorbar(surface[0], shrink=0.6, aspect=12, pad=0.1)
    cbar.set_label(r'I(x, y) (W/$\mu$m²)')
    
    plt.tight_layout()

    # === 3. Update function per frame ===
    def update(frame_index):
        E_abs = E_abs_all[frame_index, :, :]
        I = 0.5 * n * EPS0 * C0 * np.abs(E_abs)**2
        surface[0].remove()
        surface[0] = ax.plot_surface(
            X, Y, I,
            cmap='inferno',
            vmin=0.0,
            vmax=iglobal_max,
            linewidth=0,
            antialiased=False
        )

        ax.set_title(f'Intensity profile (frame = {frame_index})')
        return surface

    # === 4. Create and save GIF ===
    anim = FuncAnimation(
        fig,
        update,
        frames=NT,
        interval=1000 / fps,
        blit=False
    )

    anim.save(gif_name, writer='pillow', fps=fps)
    plt.close(fig)

    print(f"GIF saved in: {gif_name}")