import numpy as np
from pathlib import Path
from allytools.logger import get_logger
from allytools.units import Length,average_length
from gosti.fft_psf.fft_psf import FftPsf
from gosti.fft_psf.parse_zemax import read_zemax_fft_psf
from gosti.scalar2d.scalar_field2d import ScalarField2D

log = get_logger(__name__)


def fft_psf_from_zemax_report(
    *,
    path: Path,
) -> FftPsf:
    """
    Create FftPsf from a Zemax "Listing of FFT PSF Data" text export.

    Parameters
    ----------
    path : Path
        Path to the Zemax FFT PSF text report (UTF-16 "Listing of FFT PSF Data").

    Returns
    -------
    FftPsf
        PSF wrapper containing a ScalarField2D with metric coordinates.
    """
    psf, x_mm, y_mm, meta = read_zemax_fft_psf(path)
    if psf.ndim != 2:
        raise ValueError(f"Zemax FFT PSF must be 2D, got shape={psf.shape}")
    ny, nx = psf.shape
    if x_mm.shape[0] != nx or y_mm.shape[0] != ny:
        raise ValueError(
            "Mismatch between PSF shape and coordinate vectors: "
            f"psf.shape={psf.shape}, x_mm.shape={x_mm.shape}, y_mm.shape={y_mm.shape}")
    if nx > 1:
        dx_arr = np.diff(x_mm)
        dx_mean = float(dx_arr.mean())
        if not np.allclose(dx_arr, dx_mean, rtol=1e-6, atol=1e-12):
            log.warning(
                " Non-uniform X spacing detected; "
                "using mean dx=%.6g mm for ScalarField2D.",
                dx_mean)
        dx_mm = dx_mean
    else:
        dx_mm = 0.0
    if ny > 1:
        dy_arr = np.diff(y_mm)
        dy_mean = float(dy_arr.mean())
        if not np.allclose(dy_arr, dy_mean, rtol=1e-6, atol=1e-12):
            log.warning(
                " Non-uniform Y spacing detected; "
                "using mean dy=%.6g mm for ScalarField2D.",
                dy_mean)
        dy_mm = dy_mean
    else:
        dy_mm = 0.0
    min_x_mm = float(x_mm[0])
    min_y_mm = float(y_mm[0])
    log.debug(
        " Axes: min_x=%.6g mm, min_y=%.6g mm, "
        "dx=%.6g mm, dy=%.6g mm",
        min_x_mm,
        min_y_mm,
        dx_mm,
        dy_mm)
    scalar2d = ScalarField2D(
        values=psf,
        min_x=Length(min_x_mm),
        min_y=Length(min_y_mm),
        dx=Length(dx_mm),
        dy=Length(dy_mm))
    log.debug(" ScalarField2D created: shape=%s, extent_mm=%s",scalar2d.shape,scalar2d.extent_um)
    fft_psf = FftPsf(
        scalar2d=scalar2d,
        wavelength= average_length(meta.wavelength_min, meta.wavelength_max),
        field_x=meta.field_x,
        field_y=meta.field_y)
    log.info(" FftPsf created from Zemax FFT report: shape=%s",fft_psf.scalar2d.shape,)
    return fft_psf
