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.
    """

    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}"
        )

    # ----- Compute spacings -----
    if nx > 1:
        dx_arr = np.diff(x_mm)
        dx_mm = float(dx_arr.mean())
    else:
        dx_mm = 0.0

    if ny > 1:
        dy_arr = np.diff(y_mm)
        dy_mm = float(dy_arr.mean())
    else:
        dy_mm = 0.0

    # ============================================================
    # 🔧 CONSISTENT FIX: Zemax prints rows top→bottom → flip Y axis
    # ============================================================

    # Flip PSF vertically
    psf = psf[::-1, :]

    # Flip y_mm coordinates accordingly
    y_mm = y_mm[::-1].copy()

    # dy_mm should stay POSITIVE after flipping
    dy_mm = abs(dy_mm)

    # Recompute min_y after flip
    min_y_mm = float(y_mm.min())


    # x-axis untouched (Zemax left→right already consistent)
    min_x_mm = float(x_mm[0])

    log.debug(
        " Axes (after Y-flip): min_x=%.6g mm, min_y=%.6g mm, dx=%.6g mm, dy=%.6g mm",
        min_x_mm, min_y_mm, dx_mm, dy_mm
    )

    # ----- Build scalar field -----
    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
    )

    # ----- Wrap into FftPsf -----
    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
