import math
import numpy as np
from allytools.logger import get_logger

log = get_logger(__name__)
class ZernikeBasis:

    # Convert (n, m) → Noll/Zemax 1-based index j
    @staticmethod
    def _noll_j_from_nm(n: int, m: int) -> int:
        """
        Convert (n, m) → Noll / Zemax Standard index j = 1,2,...

        Validity rules:
        - |m| <= n
        - (n - m) must be even
        """

        log.trace(f" Computing Noll index for (n={n}, m={m}).")

        if abs(m) > n or ((n - m) % 2) != 0:
            raise ValueError(f"Invalid (n,m) for Zernike: n={n}, m={m}")

        a = n * (n + 1) // 2 + abs(m)
        n_mod = n % 4

        if (m > 0 and n_mod in (0, 1)) or (m < 0 and n_mod in (2, 3)):
            offset = 0
        else:
            offset = 1

        j = a + offset
        log.trace(f" Noll index j={j} for (n={n}, m={m}).")

        return j


    # Convert Zemax Standard index j → (n, m)
    @staticmethod
    def zemax_index_to_nm(j: int) -> tuple[int, int]:
        """
        Convert j = 1..N (Zemax Zernike Standard Coefficients)
        → (n, m) following the Noll/OpticStudio convention.
        """

        log.trace(f" Converting Zemax index j={j} → (n, m).")

        j = int(j)

        # n up to 20 easily covers all 231 Zemax terms
        for n in range(0, 20):
            for m in range(-n, n + 1, 2):
                if ZernikeBasis._noll_j_from_nm(n, m) == j:
                    log.trace(f" j={j} maps to (n={n}, m={m}).")
                    return n, m

        raise ValueError(f"Cannot map Zernike index j={j} to (n,m)")


    # OSA/ANSI mapping j(0-based) → (n,m)
    @staticmethod
    def osa_index_to_nm(j: int) -> tuple[int, int]:
        """
        OSA/ANSI mapping (0-based index):
            j = (n(n+2) + m) / 2
        """

        n = int(math.floor((math.sqrt(8*j + 1) - 1) / 2))
        m = int(2*j - n*(n + 2))
        log.trace(f"OSA index j={j} → (n={n}, m={m}).")
        return n, m

    # Radial polynomial R_n^{|m|}
    @staticmethod
    def radial_poly(n: int, m: int, rho: np.ndarray) -> np.ndarray:
        """
        Compute the radial Zernike polynomial R_n^{|m|}(rho)
        using the OSA/ANSI definition.
        """

        log.trace(f"Computing radial polynomial R_{n}^{{|{m}|}} for rho grid.")
        m = abs(m)
        R = np.zeros_like(rho)

        # Standard OSA radial sum
        for k in range((n - m)//2 + 1):
            num = math.factorial(n - k)
            den = (
                math.factorial(k)
                * math.factorial((n + m)//2 - k)
                * math.factorial((n - m)//2 - k)
            )
            c = ((-1)**k) * num / den
            R += c * rho**(n - 2*k)

        return R


    # Build first N Zernike modes on rho<=1
    def build_basis(self, N: int, rho: np.ndarray, phi: np.ndarray, mask: np.ndarray):
        """
        Construct the first N Zernike basis functions on the unit disk.

        Inputs:
            N   - number of Zernike modes
            rho - normalized radius grid (0..1)
            phi - azimuth grid
            mask - pupil mask (1 inside, 0 outside)

        Returns:
            basis_flat: (N, grid_size*grid_size)
            mask      : (grid_size, grid_size)
        """

        log.trace(f"Building Zernike basis for N={N} modes.")

        gs = rho.shape[0]
        basis = np.zeros((N, gs, gs), dtype=float)

        # Loop over all modes
        for k in range(1, N+1):
            j = k  # in Zemax / Noll indexing, mode index = j
            n, m = self.zemax_index_to_nm(j)

            log.trace(f"Mode k={k}: using (n={n}, m={m}).")

            # Radial component
            R = self.radial_poly(n, m, rho)

            # Angular part according to Noll normalization
            if m == 0:
                Z = math.sqrt(n + 1) * R
                log.trace(f"Mode k={k}: m=0 → cosine term omitted.")
            elif m > 0:
                Z = math.sqrt(2 * (n + 1)) * R * np.cos(m * phi)
                log.trace(f"Mode k={k}: positive m → cos(mφ).")
            else:
                Z = math.sqrt(2 * (n + 1)) * R * np.sin(-m * phi)
                log.trace(f"Mode k={k}: negative m → sin(-mφ).")

            # Apply pupil mask
            basis[k-1] = Z * mask
            log.trace(f"Mode k={k}: applied mask. Non-zero count = {np.sum(mask > 0)}")

        # Flatten basis into shape (N, grid_size²)
        basis_flat = basis.reshape(N, -1)
        log.debug(f"Basis assembled. Shape={basis_flat.shape}")
        return basis_flat, mask
