"""
pycutwm.power_xy.py
------------

Main routines to compute optical power from the x-y field distribution.

Functions:
    - compute_power_from_h5: Compute optical power for a given field from its HDF5 file.
    - compute_power_vs_time_from_h5: Compute optical power for each time (x-y slice) in 
      (3+1)D problems from HDF5 files.
    - plot_power_vs_time: From the previously computed power vs time, plot it.

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



import os
import glob
import h5py
import json
import numpy as np
import matplotlib.pyplot as plt
from typing import Optional, Dict, Any, Tuple

# 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 _auto_find_config_json(folder: str) -> str:
    candidates = sorted(glob.glob(os.path.join(folder, "config_*.json")))
    if len(candidates) == 0:
        raise FileNotFoundError(f"No config_*.json found in {folder}")
    if len(candidates) > 1:
        raise RuntimeError(f"Multiple config_*.json files found in {folder}: {candidates}")
    return candidates[0]


def _read_json_key(config_path: str, key: str) -> Any:
    with open(config_path, "r") as f:
        data = json.load(f)
    for part in key.split('.'):
        data = data[part]
    return data


def _grid_from_json(config_path: str) -> Tuple[int, int, float, float]:
    NX = int(_read_json_key(config_path, "grid.grid_points.NX"))
    NY = int(_read_json_key(config_path, "grid.grid_points.NY"))
    LX_um = float(_read_json_key(config_path, "crystal.dimensions.LX"))
    LY_um = float(_read_json_key(config_path, "crystal.dimensions.LY"))
    dx = LX_um / (NX-1)
    dy = LY_um / (NY-1)
    return NX, NY, dx, dy


def _load_complex_field_h5(field_path: str, shape: Tuple[int, int]) -> np.ndarray:
    with h5py.File(field_path, 'r') as f:
        # Names must be adjusted if your datasets are different!
        Er = f["real"][:]
        Ei = f["imag"][:]
    # Reshape as (NY, NX)
    NY, NX = shape
    Er = Er.reshape(NY, NX)
    Ei = Ei.reshape(NY, NX)
    return Er + 1j * Ei


def _intensity(E: np.ndarray, n: float, field_scale: float = 1.0) -> np.ndarray:
    E_vpm = E * field_scale
    return 0.5 * n * EPS0 * C0 * np.abs(E_vpm)**2


def _power(I: np.ndarray, dx: float, dy: float) -> float:
    return float(I.sum() * dx * dy)


def compute_power_from_h5(
    folder: str,
    field: str,
    n: float,
    config_path: Optional[str] = None,
    field_scale: float = 1.0
) -> Dict[str, float]:
    """
    Compute optical power for a given field (p/s/i) from its .h5 file.
    Args:
        folder: Folder containing output files and config.
        field:  'p', 's', or 'i' (pump, signal, idler).
        n:      Refractive index for the field.
        config_path: Path to config JSON (optional).
        field_scale: Multiply E by this factor to convert to V/μm (default: 1.0).
    Returns:
        Dictionary: { "power": value_in_W }
    """
    if config_path is None:
        config_path = _auto_find_config_json(folder)
    NX, NY, dx, dy = _grid_from_json(config_path)
    h5_file = os.path.join(folder, f"{'pump' if field=='p' else 'signal' if field=='s' else 'idler'}_output_XY.h5")
    if not os.path.isfile(h5_file):
        raise FileNotFoundError(f"{h5_file} not found.")
    E = _load_complex_field_h5(h5_file, shape=(NY, NX))
    I = _intensity(E, n=n, field_scale=field_scale)
    P = _power(I, dx, dy)
    return {"power": P}


def compute_power_vs_time_from_h5(
    h5_path: str,
    time_h5_path: str,
    field: str,
    n_dict: Dict[str, float],
    config_path: str,
    field_scale: float = 1.0
):
    # --- Load config for grid info ---
    with open(config_path, "r") as f:
        config = json.load(f)
    NX_cfg = int(config["grid"]["grid_points"]["NX"])
    NY_cfg = int(config["grid"]["grid_points"]["NY"])
    LX_um = float(config["crystal"]["dimensions"]["LX"])
    LY_um = float(config["crystal"]["dimensions"]["LY"])
    dx = LX_um / (NX_cfg - 1)
    dy = LY_um / (NY_cfg - 1)

    # --- Read field from HDF5 (shape: NT × NY × NX) ---
    with h5py.File(h5_path, "r") as f:
        Er = f["real"][:]  # (NT, NY, NX)
        Ei = f["imag"][:]  # (NT, NY, NX)

    E = (Er + 1j * Ei) * field_scale
    NT, NY, NX = E.shape

    # Pequeña comprobación de coherencia
    assert NX == NX_cfg and NY == NY_cfg, "Dimensiones de HDF5 y config.json no coinciden"

    n = n_dict[field]

    # --- Compute power vs time ---
    powers = np.zeros(NT)
    for it in range(NT):
        E_slice = E[it, :, :]
        I = 0.5 * n * EPS0 * C0 * np.abs(E_slice)**2
        powers[it] = I.sum() * dx * dy

    # --- Load times from time.h5 ---
    with h5py.File(time_h5_path, "r") as f:
        times = f["data"][:]   # asumimos longitud NT y mismo orden

    return times, powers


def plot_power_vs_time(times: np.ndarray, powers: np.ndarray, field: str, figname: str):
    """Plot optical power vs time."""
    fig = plt.figure(figsize=(6, 6))
    plt.plot(times, powers, label=f"{field} power")
    plt.xlabel("Time (ps)")
    plt.ylabel("Power (W)")
    plt.title(f"Power vs Time for {field}")
    plt.legend()
    plt.grid(True, ls="--", alpha=0.6)
    plt.tight_layout()
    plt.show()
    fig.savefig(f"{figname}.png", dpi=150, bbox_inches='tight')
    plt.close(fig)