#!/usr/bin/env python3
"""
Synthetic WFE Surface Generator for Zernike Pipeline Validation
================================================================

Generates synthetic Optical Path Difference (OPD) maps from user-defined
Zernike coefficients. Used to validate the Atomizer Zernike fitting pipeline
by creating "known truth" surfaces that can be round-tripped through the
WFE_from_CSV_OPD tool.

Features:
- Noll-indexed Zernike polynomials (standard optical convention)
- Full-disk or annular aperture support
- Configurable grid density, mirror diameter, noise level
- Multiple output formats: CSV (for WFE_from_CSV_OPD), NumPy, plot
- Preset test cases for common validation scenarios

Usage:
    # Single mode test (pure astigmatism)
    python generate_synthetic_wfe.py --mode single --zernike "5:100" --output test_astig.csv

    # Multi-mode realistic mirror
    python generate_synthetic_wfe.py --mode multi --output test_multi.csv

    # Custom coefficients
    python generate_synthetic_wfe.py --zernike "5:80,7:45,9:25,11:15" --output test_custom.csv

    # With noise
    python generate_synthetic_wfe.py --mode multi --noise 2.0 --output test_noisy.csv

    # Full test suite (generates all validation cases)
    python generate_synthetic_wfe.py --suite --output-dir validation_suite/

Output CSV format (matching WFE_from_CSV_OPD):
    x(mm), y(mm), dz(mm)
    where dz is surface displacement (OPD = 2 * dz for reflective surfaces)

Author: Mario (Atomizer V&V)
Created: 2026-03-09
Project: P-Zernike-Validation (GigaBIT M1)
"""

import sys
import os
import argparse
import json
from pathlib import Path
from math import factorial
from typing import Dict, List, Tuple, Optional
import numpy as np

# ============================================================================
# Zernike Polynomial Mathematics (matching extract_zernike.py conventions)
# ============================================================================

def noll_indices(j: int) -> Tuple[int, int]:
    """Convert Noll index j to radial order n and azimuthal frequency m."""
    if j < 1:
        raise ValueError("Noll index j must be >= 1")
    count = 0
    n = 0
    while True:
        if n == 0:
            ms = [0]
        elif n % 2 == 0:
            ms = [0] + [m for k in range(1, n // 2 + 1) for m in (-2 * k, 2 * k)]
        else:
            ms = [m for k in range(0, (n + 1) // 2) for m in (-(2 * k + 1), (2 * k + 1))]
        for m in ms:
            count += 1
            if count == j:
                return n, m
        n += 1


def zernike_radial(n: int, m: int, r: np.ndarray) -> np.ndarray:
    """Compute radial component R_n^m(r)."""
    R = np.zeros_like(r)
    m_abs = abs(m)
    for s in range((n - m_abs) // 2 + 1):
        coef = ((-1) ** s * factorial(n - s) /
                (factorial(s) *
                 factorial((n + m_abs) // 2 - s) *
                 factorial((n - m_abs) // 2 - s)))
        R += coef * r ** (n - 2 * s)
    return R


def zernike_noll(j: int, r: np.ndarray, theta: np.ndarray) -> np.ndarray:
    """Evaluate Noll-indexed Zernike polynomial Z_j(r, theta)."""
    n, m = noll_indices(j)
    R = zernike_radial(n, m, r)
    if m == 0:
        return R
    elif m > 0:
        return R * np.cos(m * theta)
    else:
        return R * np.sin(-m * theta)


def zernike_name(j: int) -> str:
    """Get common optical name for Zernike mode."""
    n, m = noll_indices(j)
    names = {
        (0, 0): "Piston",
        (1, -1): "Tilt X", (1, 1): "Tilt Y",
        (2, 0): "Defocus",
        (2, -2): "Astigmatism 45°", (2, 2): "Astigmatism 0°",
        (3, -1): "Coma X", (3, 1): "Coma Y",
        (3, -3): "Trefoil X", (3, 3): "Trefoil Y",
        (4, 0): "Primary Spherical",
        (4, -2): "2nd Astig X", (4, 2): "2nd Astig Y",
        (4, -4): "Quadrafoil X", (4, 4): "Quadrafoil Y",
        (5, -1): "2nd Coma X", (5, 1): "2nd Coma Y",
        (5, -3): "2nd Trefoil X", (5, 3): "2nd Trefoil Y",
        (5, -5): "Pentafoil X", (5, 5): "Pentafoil Y",
        (6, 0): "2nd Spherical",
    }
    if (n, m) in names:
        return names[(n, m)]
    return f"Z({n},{m:+d})"


# ============================================================================
# Surface Generation
# ============================================================================

def generate_grid(
    diameter_mm: float = 1200.0,
    inner_radius_mm: float = 135.75,
    n_points_radial: int = 200,
    grid_type: str = "cartesian",
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Generate a 2D grid of points within the mirror aperture.

    Returns:
        x_mm, y_mm: Physical coordinates in mm
        mask: Boolean array (True = inside aperture)
    """
    outer_radius = diameter_mm / 2.0

    if grid_type == "cartesian":
        # Regular Cartesian grid
        n = n_points_radial * 2
        x_1d = np.linspace(-outer_radius, outer_radius, n)
        y_1d = np.linspace(-outer_radius, outer_radius, n)
        x_mm, y_mm = np.meshgrid(x_1d, y_1d)
        x_mm = x_mm.ravel()
        y_mm = y_mm.ravel()
    elif grid_type == "scattered":
        # Random scattered points (simulates real mesh)
        n_total = (n_points_radial * 2) ** 2
        rng = np.random.default_rng(42)
        x_mm = rng.uniform(-outer_radius, outer_radius, n_total)
        y_mm = rng.uniform(-outer_radius, outer_radius, n_total)
    else:
        raise ValueError(f"Unknown grid_type: {grid_type}")

    # Apply aperture mask
    r = np.sqrt(x_mm**2 + y_mm**2)
    if inner_radius_mm > 0:
        mask = (r <= outer_radius) & (r >= inner_radius_mm)
    else:
        mask = r <= outer_radius

    return x_mm[mask], y_mm[mask], mask


def synthesize_surface(
    x_mm: np.ndarray,
    y_mm: np.ndarray,
    coefficients: Dict[int, float],
    diameter_mm: float = 1200.0,
    noise_rms_nm: float = 0.0,
    seed: int = 42,
) -> Tuple[np.ndarray, Dict]:
    """
    Generate a synthetic OPD surface from Zernike coefficients.

    Args:
        x_mm, y_mm: Physical coordinates in mm
        coefficients: Dict of {Noll_index: amplitude_nm}
        diameter_mm: Mirror diameter in mm
        noise_rms_nm: RMS of Gaussian noise to add (nm)
        seed: Random seed for reproducibility

    Returns:
        opd_nm: OPD values in nanometers at each point
        metadata: Dict with ground truth info
    """
    outer_radius = diameter_mm / 2.0

    # Normalize to unit disk
    r_norm = np.sqrt(x_mm**2 + y_mm**2) / outer_radius
    theta = np.arctan2(y_mm, x_mm)

    # Build surface from Zernike modes
    opd_nm = np.zeros_like(x_mm)
    for j, amp_nm in coefficients.items():
        Z_j = zernike_noll(j, r_norm, theta)
        opd_nm += amp_nm * Z_j

    # Compute ground truth RMS (before noise)
    rms_total = np.sqrt(np.mean(opd_nm**2))

    # Add noise if requested
    if noise_rms_nm > 0:
        rng = np.random.default_rng(seed)
        noise = rng.normal(0, noise_rms_nm, size=opd_nm.shape)
        opd_nm += noise

    rms_with_noise = np.sqrt(np.mean(opd_nm**2))

    metadata = {
        "input_coefficients": {str(j): amp for j, amp in coefficients.items()},
        "coefficient_names": {str(j): zernike_name(j) for j in coefficients},
        "n_points": len(x_mm),
        "diameter_mm": diameter_mm,
        "rms_nm_clean": float(rms_total),
        "rms_nm_with_noise": float(rms_with_noise),
        "noise_rms_nm": noise_rms_nm,
        "seed": seed,
    }

    return opd_nm, metadata


# ============================================================================
# Output Formats
# ============================================================================

def write_csv_opd(
    filepath: str,
    x_mm: np.ndarray,
    y_mm: np.ndarray,
    opd_nm: np.ndarray,
    opd_unit: str = "nm",
):
    """
    Write OPD surface to CSV format compatible with WFE_from_CSV_OPD.

    Default format: x(mm), y(mm), opd(nm)
    Can also output in mm for displacement: x(mm), y(mm), dz(mm)

    NOTE: Update this function once the exact WFE_from_CSV_OPD format is confirmed.
    """
    filepath = Path(filepath)
    filepath.parent.mkdir(parents=True, exist_ok=True)

    if opd_unit == "nm":
        header = "x(mm),y(mm),opd(nm)"
        data = np.column_stack([x_mm, y_mm, opd_nm])
    elif opd_unit == "mm":
        # Convert nm to mm for displacement
        dz_mm = opd_nm / 1e6
        header = "x(mm),y(mm),dz(mm)"
        data = np.column_stack([x_mm, y_mm, dz_mm])
    else:
        raise ValueError(f"Unknown opd_unit: {opd_unit}")

    np.savetxt(filepath, data, delimiter=",", header=header, comments="",
               fmt="%.10e")
    print(f"  Written: {filepath} ({len(x_mm)} points)")


def write_metadata(filepath: str, metadata: Dict):
    """Write ground truth metadata as JSON."""
    filepath = Path(filepath)
    filepath.parent.mkdir(parents=True, exist_ok=True)
    with open(filepath, 'w') as f:
        json.dump(metadata, f, indent=2)
    print(f"  Metadata: {filepath}")


# ============================================================================
# Preset Test Cases
# ============================================================================

# Realistic M1 mirror: typical gravity-induced deformation pattern
PRESET_REALISTIC = {
    5: 80.0,    # Astigmatism 0° — dominant gravity mode
    6: 45.0,    # Astigmatism 45°
    7: 30.0,    # Coma X
    8: 20.0,    # Coma Y
    9: 15.0,    # Trefoil X
    11: 10.0,   # Primary Spherical
    13: 5.0,    # Secondary Astigmatism
    16: 3.0,    # Secondary Coma
    22: 2.0,    # Secondary Spherical
}

# Single-mode test cases
SINGLE_MODE_TESTS = {
    "Z05_astig_0deg":    {5: 100.0},
    "Z06_astig_45deg":   {6: 100.0},
    "Z07_coma_x":        {7: 100.0},
    "Z08_coma_y":        {8: 100.0},
    "Z09_trefoil_x":     {9: 100.0},
    "Z10_trefoil_y":     {10: 100.0},
    "Z11_spherical":     {11: 100.0},
    "Z22_2nd_spherical": {22: 50.0},
    "Z37_high_order":    {37: 30.0},
}

# Edge case tests
EDGE_CASE_TESTS = {
    "near_zero":        {5: 0.1, 7: 0.05, 11: 0.01},  # Sub-nm coefficients
    "large_amplitude":  {5: 500.0, 7: 300.0},          # Large deformation
    "many_modes":       {j: 100.0 / j for j in range(5, 51)},  # All modes, decreasing
}


def generate_test_suite(
    output_dir: str,
    diameter_mm: float = 1200.0,
    inner_radius_mm: float = 135.75,
    n_points_radial: int = 200,
):
    """Generate the full validation test suite."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    all_cases = {}

    # 1. Single-mode tests
    print("\n=== Single-Mode Tests ===")
    for name, coeffs in SINGLE_MODE_TESTS.items():
        print(f"\nGenerating: {name}")
        x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
        opd, meta = synthesize_surface(x, y, coeffs, diameter_mm)
        write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
        write_metadata(output_dir / f"{name}_truth.json", meta)
        all_cases[name] = meta

    # 2. Realistic multi-mode
    print("\n=== Realistic Multi-Mode ===")
    name = "realistic_gravity"
    print(f"\nGenerating: {name}")
    x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
    opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm)
    write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
    write_metadata(output_dir / f"{name}_truth.json", meta)
    all_cases[name] = meta

    # 3. Noisy versions
    print("\n=== Noisy Tests ===")
    for noise_level in [1.0, 5.0, 10.0]:
        name = f"realistic_noise_{noise_level:.0f}nm"
        print(f"\nGenerating: {name}")
        x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
        opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm, noise_rms_nm=noise_level)
        write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
        write_metadata(output_dir / f"{name}_truth.json", meta)
        all_cases[name] = meta

    # 4. Edge cases
    print("\n=== Edge Case Tests ===")
    for name, coeffs in EDGE_CASE_TESTS.items():
        print(f"\nGenerating: {name}")
        x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial)
        opd, meta = synthesize_surface(x, y, coeffs, diameter_mm)
        write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
        write_metadata(output_dir / f"{name}_truth.json", meta)
        all_cases[name] = meta

    # 5. Full-disk (no central hole) for comparison
    print("\n=== Full-Disk (No Hole) ===")
    name = "realistic_full_disk"
    print(f"\nGenerating: {name}")
    x, y, _ = generate_grid(diameter_mm, 0.0, n_points_radial)
    opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm)
    meta["inner_radius_mm"] = 0.0
    write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
    write_metadata(output_dir / f"{name}_truth.json", meta)
    all_cases[name] = meta

    # 6. Scattered points (simulates real FEA mesh)
    print("\n=== Scattered Grid (FEA-like) ===")
    name = "realistic_scattered"
    print(f"\nGenerating: {name}")
    x, y, _ = generate_grid(diameter_mm, inner_radius_mm, n_points_radial, grid_type="scattered")
    opd, meta = synthesize_surface(x, y, PRESET_REALISTIC, diameter_mm)
    meta["grid_type"] = "scattered"
    write_csv_opd(output_dir / f"{name}.csv", x, y, opd)
    write_metadata(output_dir / f"{name}_truth.json", meta)
    all_cases[name] = meta

    # Summary
    print(f"\n{'='*60}")
    print(f"Generated {len(all_cases)} test cases in: {output_dir}")
    print(f"{'='*60}")

    # Write suite manifest
    manifest = {
        "suite": "Zernike Pipeline Validation",
        "generated": "2026-03-09",
        "mirror_diameter_mm": diameter_mm,
        "inner_radius_mm": inner_radius_mm,
        "n_zernike_modes": 50,
        "n_points_radial": n_points_radial,
        "cases": all_cases,
    }
    manifest_path = output_dir / "suite_manifest.json"
    with open(manifest_path, 'w') as f:
        json.dump(manifest, f, indent=2)
    print(f"Manifest: {manifest_path}")

    return all_cases


# ============================================================================
# CLI
# ============================================================================

def parse_zernike_string(s: str) -> Dict[int, float]:
    """Parse 'j1:amp1,j2:amp2,...' format."""
    coeffs = {}
    for pair in s.split(","):
        parts = pair.strip().split(":")
        if len(parts) != 2:
            raise ValueError(f"Invalid format: '{pair}'. Use 'j:amplitude'")
        j = int(parts[0])
        amp = float(parts[1])
        coeffs[j] = amp
    return coeffs


def main():
    parser = argparse.ArgumentParser(
        description="Generate synthetic WFE surfaces for Zernike validation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Pure astigmatism (100nm)
  python generate_synthetic_wfe.py --zernike "5:100" -o test_astig.csv

  # Realistic multi-mode mirror
  python generate_synthetic_wfe.py --preset realistic -o test_realistic.csv

  # Full validation suite
  python generate_synthetic_wfe.py --suite -d validation_suite/

  # Custom with noise
  python generate_synthetic_wfe.py --zernike "5:80,7:45,11:15" --noise 2.0 -o test.csv
        """
    )

    parser.add_argument("--zernike", "-z", type=str,
                        help="Zernike coefficients as 'j:amp_nm,...' (e.g. '5:100,7:50')")
    parser.add_argument("--preset", choices=["realistic"],
                        help="Use preset coefficient set")
    parser.add_argument("--suite", action="store_true",
                        help="Generate full validation test suite")
    parser.add_argument("--output", "-o", type=str, default="synthetic_wfe.csv",
                        help="Output CSV file path")
    parser.add_argument("--output-dir", "-d", type=str, default="validation_suite",
                        help="Output directory for --suite mode")
    parser.add_argument("--diameter", type=float, default=1200.0,
                        help="Mirror diameter in mm (default: 1200)")
    parser.add_argument("--inner-radius", type=float, default=135.75,
                        help="Inner radius in mm for annular aperture (default: 135.75, 0=full disk)")
    parser.add_argument("--grid-points", type=int, default=200,
                        help="Number of radial grid points (default: 200, total ~ 4*N^2)")
    parser.add_argument("--noise", type=float, default=0.0,
                        help="Gaussian noise RMS in nm (default: 0)")
    parser.add_argument("--grid-type", choices=["cartesian", "scattered"],
                        default="cartesian",
                        help="Grid type (default: cartesian)")
    parser.add_argument("--unit", choices=["nm", "mm"], default="nm",
                        help="OPD output unit (default: nm)")

    args = parser.parse_args()

    print("=" * 60)
    print("Synthetic WFE Surface Generator")
    print("Zernike Pipeline Validation — GigaBIT M1")
    print("=" * 60)
    print(f"  Mirror diameter: {args.diameter} mm")
    print(f"  Inner radius:    {args.inner_radius} mm")
    print(f"  Grid points:     ~{(args.grid_points*2)**2}")

    if args.suite:
        generate_test_suite(
            args.output_dir,
            args.diameter,
            args.inner_radius,
            args.grid_points,
        )
    else:
        # Determine coefficients
        if args.preset == "realistic":
            coeffs = PRESET_REALISTIC
        elif args.zernike:
            coeffs = parse_zernike_string(args.zernike)
        else:
            print("\nERROR: Specify --zernike, --preset, or --suite")
            sys.exit(1)

        print(f"\n  Coefficients:")
        for j, amp in sorted(coeffs.items()):
            print(f"    Z{j:2d} ({zernike_name(j):20s}): {amp:8.2f} nm")

        # Generate
        x, y, _ = generate_grid(args.diameter, args.inner_radius,
                                 args.grid_points, args.grid_type)
        opd, meta = synthesize_surface(x, y, coeffs, args.diameter,
                                        noise_rms_nm=args.noise)

        print(f"\n  Points in aperture: {len(x)}")
        print(f"  RMS (clean):  {meta['rms_nm_clean']:.3f} nm")
        if args.noise > 0:
            print(f"  RMS (noisy):  {meta['rms_nm_with_noise']:.3f} nm")

        # Write outputs
        print(f"\n  Output unit: {args.unit}")
        write_csv_opd(args.output, x, y, opd, opd_unit=args.unit)

        meta_path = Path(args.output).with_suffix('.json')
        write_metadata(str(meta_path), meta)

    print("\nDone!")


if __name__ == "__main__":
    main()