import math
import numpy as np
from typing import Optional
from allytools.logger import get_logger
from allytools.units import Length, LengthUnit, Angle, AngleUnit
from gosti.base.pupil_grid import PupilGrid
from gosti.zernike.zernike_base import ZernikeBasis
from gosti.scalar2d.scalar_field2d import ScalarField2D

log = get_logger(__name__)
def compute_fft_psf(
    *,
    coefficients: np.ndarray,
    n_coef: int,
    wavelength: Length,
    exit_pupil_diameter: Length,
    exit_pupil_position: Length,
    focal_length: Length,
    grid_size: int,
    pad_factor: int = 4,
    field_x_img: Optional[Length] = None,
    field_y_img: Optional[Length] = None) -> ScalarField2D:
    """
    FFT-based PSF computation from Zernike phase map.

    Returns
    -------
    ScalarField2D
        PSF as a normalized intensity map in the image plane (mm).
    """

    # 1️⃣ Build the normalized pupil grid (handles off-axis / elliptical pupil)
    grid = PupilGrid(
        grid_size=grid_size,
        field_x_img=field_x_img ,
        field_y_img=field_y_img,
        exit_pupil_position=exit_pupil_position)
    x_norm, y_norm, rho, phi, mask = grid.build()
    log.debug("PupilGrid.build() returned mask with %d inside pixels.", mask.sum())

    # 2️⃣ Build Zernike basis
    zb = ZernikeBasis()
    basis_flat, pupil_mask = zb.build_basis(n_coef, rho, phi, mask)
    log.debug("Zernike basis built: basis_flat = %s, active pixels = %d.",basis_flat.shape,pupil_mask.sum())

    # 3️⃣ Construct phase map (waves)
    coeffs_used = coefficients[:n_coef]
    phase = (coeffs_used @ basis_flat).reshape(grid_size, grid_size)

    # RMS on pupil
    phase_rms = np.sqrt(((phase ** 2) * pupil_mask).sum() / pupil_mask.sum())
    log.debug("Phase RMS = %.6e waves.", phase_rms)

    # 5️⃣ Complex pupil function
    pupil = pupil_mask * np.exp(1j * 2 * math.pi * phase)
    log.debug("Complex pupil constructed.")

    # 6️⃣ Zero-padding
    h0 = w0 = grid_size
    h = pad_factor * h0
    w = pad_factor * w0
    pad_h = (h - h0) // 2
    pad_w = (w - w0) // 2

    pupil_padded = np.pad(
        pupil,
        ((pad_h, h - h0 - pad_h), (pad_w, w - w0 - pad_w)))
    log.debug("Zero-padding applied: %d→%d, %d→%d (Pad_H=%d, Pad_W=%d).",h0, h, w0, w, pad_h, pad_w)

    # 7️⃣ FFT → PSF
    field = np.fft.fftshift( np.fft.fft2(np.fft.ifftshift(pupil_padded)))
    psf = np.abs(field) ** 2
    psf_sum = psf.sum()
    if psf_sum > 0:
        psf /= psf_sum

    log.debug("PSF computed and normalized (sum=%.6f).", psf.sum())

    # 8️⃣ Image plane coordinates (mm), wrapped directly into ScalarField2D
    #
    # Angular sampling after padding:


    delta_theta = Angle(wavelength / (exit_pupil_diameter * pad_factor), AngleUnit.RAD)
    delta_x = exit_pupil_position * delta_theta.value_rad  # same for x and y
    log.debug("λ = %.1f [nm], f = %.2f [mm] EPD = %.2f [mm], pad = %d,  Δθ = %.3e [rad], Δx = %.3f [um]",
              wavelength.to(LengthUnit.NM),
              focal_length.to(LengthUnit.MM),
              exit_pupil_diameter.to(LengthUnit.MM),
              pad_factor,
              delta_theta.to(AngleUnit.RAD),
              delta_x.to(LengthUnit.UM))

    # Index grids centered at 0
    x_idx = np.arange(w) - w // 2
    y_idx = np.arange(h) - h // 2

    # Pixel pitch and min-coordinates
    dx = delta_x
    dy = delta_x
    min_x = x_idx[0] * dx
    min_y = y_idx[0] * dy

    log.debug("Image-plane sampling: delta_x=%.3e [um] , grid %dx%d, min_x=%.1e [um], min_y=%.1e [um]",
        dx.to(LengthUnit.UM), h, w, min_x.to(LengthUnit.UM), min_y.to(LengthUnit.UM))


    psf_field = ScalarField2D(
        values=psf.astype(np.float64, copy=False),
        min_x=min_x,
        min_y=min_y,
        dx=dx,
        dy=dy)

    return psf_field
