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

Generates synthetic FEA-style CSV files with known Zernike content for
validating the WFE_from_CSV_OPD tool. Output matches the exact format
used by Atomizer's Zernike analysis pipeline.

Output CSV format (matching WFE_from_CSV_OPD):
    ,X,Y,Z,DX,DY,DZ
    0, x_meters, y_meters, z_meters, dx_meters, dy_meters, dz_meters

Where:
    X, Y, Z = undeformed node positions (meters)
    DX, DY, DZ = displacement vector (meters)
    The tool computes OPD from displacements (rigorous method accounts for
    lateral DX/DY via interpolation on the reference surface).

Two operating modes:
    1. PURE SYNTHETIC: Generate a flat/spherical mirror mesh + inject known
       Zernike displacements as DZ. Ground truth is exact.
    2. INJECT INTO REAL: Load a real FEA CSV, zero out DZ, then inject known
       Zernike content. Tests the full pipeline including real mesh geometry.

Usage:
    # Pure synthetic (flat mirror, M1 params)
    python generate_synthetic_wfe.py --zernike "5:100,7:50" -o test.csv

    # Pure synthetic with spherical surface
    python generate_synthetic_wfe.py --preset realistic --surface sphere --roc 5000 -o test.csv

    # Inject into real M2 data
    python generate_synthetic_wfe.py --inject m2_data.csv --zernike "5:100" -o test_injected.csv

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

    # Full suite with injection into real mesh
    python generate_synthetic_wfe.py --suite --inject m2_data.csv -d validation_suite/

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

import sys
import os
import argparse
import json
import csv
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})"


# ============================================================================
# Mesh Generation
# ============================================================================

def generate_mesh(
    diameter_mm: float = 1200.0,
    inner_radius_mm: float = 0.0,
    n_rings: int = 20,
    surface_type: str = "flat",
    roc_mm: float = 0.0,
    conic: float = 0.0,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Generate a mirror mesh with realistic node distribution.

    Args:
        diameter_mm: Mirror outer diameter in mm
        inner_radius_mm: Inner radius (0 for full disk)
        n_rings: Number of radial rings
        surface_type: "flat", "sphere", "parabola"
        roc_mm: Radius of curvature in mm (for sphere/parabola)
        conic: Conic constant (-1 = parabola, 0 = sphere)

    Returns:
        x_m, y_m, z_m: Node positions in meters
    """
    outer_r_mm = diameter_mm / 2.0

    # Generate rings with increasing node count
    all_x = []
    all_y = []

    if inner_radius_mm > 0:
        r_values = np.linspace(inner_radius_mm, outer_r_mm, n_rings)
    else:
        # Include center point
        r_values = np.linspace(0, outer_r_mm, n_rings + 1)

    for i, r in enumerate(r_values):
        if r < 1e-6:
            all_x.append(0.0)
            all_y.append(0.0)
        else:
            # More nodes at larger radii (like real FEA mesh)
            n_pts = max(6, int(6 + i * 3))
            angles = np.linspace(0, 2 * np.pi, n_pts, endpoint=False)
            for a in angles:
                all_x.append(r * np.cos(a))
                all_y.append(r * np.sin(a))

    x_mm = np.array(all_x)
    y_mm = np.array(all_y)

    # Compute Z (surface sag)
    if surface_type == "flat":
        z_mm = np.zeros_like(x_mm)
    elif surface_type in ("sphere", "parabola"):
        if roc_mm <= 0:
            raise ValueError("roc_mm must be > 0 for curved surfaces")
        r2 = x_mm**2 + y_mm**2
        if surface_type == "sphere":
            # z = R - sqrt(R^2 - r^2)
            z_mm = roc_mm - np.sqrt(roc_mm**2 - r2)
        else:
            # Parabola: z = r^2 / (2R)
            z_mm = r2 / (2 * roc_mm)
    else:
        raise ValueError(f"Unknown surface_type: {surface_type}")

    # Convert to meters
    return x_mm / 1000.0, y_mm / 1000.0, z_mm / 1000.0


def load_real_mesh(csv_path: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Load node positions from a real FEA CSV file.

    Returns:
        x_m, y_m, z_m: Node positions in meters
    """
    x_list, y_list, z_list = [], [], []
    with open(csv_path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            x_list.append(float(row['X']))
            y_list.append(float(row['Y']))
            z_list.append(float(row['Z']))

    return np.array(x_list), np.array(y_list), np.array(z_list)


# ============================================================================
# Surface Synthesis
# ============================================================================

def synthesize_displacements(
    x_m: np.ndarray,
    y_m: np.ndarray,
    coefficients: Dict[int, float],
    diameter_mm: float = None,
    noise_rms_nm: float = 0.0,
    include_lateral: bool = False,
    seed: int = 42,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict]:
    """
    Generate synthetic displacement vectors from Zernike coefficients.

    The Zernike content is injected into DZ (surface-normal displacement).
    DX and DY are set to zero unless include_lateral is True (adds small
    realistic lateral displacements).

    Args:
        x_m, y_m: Node positions in meters
        coefficients: {Noll_index: amplitude_nm}
        diameter_mm: Mirror diameter (auto-detected if None)
        noise_rms_nm: Gaussian noise RMS in nm
        include_lateral: Add small realistic DX/DY
        seed: Random seed

    Returns:
        dx_m, dy_m, dz_m: Displacement vectors in meters
        metadata: Ground truth info
    """
    # Auto-detect diameter from mesh
    if diameter_mm is None:
        r_mm = np.sqrt(x_m**2 + y_m**2) * 1000.0
        diameter_mm = 2.0 * np.max(r_mm)

    outer_r_m = diameter_mm / 2000.0  # meters

    # Normalize to unit disk
    r_m = np.sqrt(x_m**2 + y_m**2)
    r_norm = r_m / outer_r_m
    theta = np.arctan2(y_m, x_m)

    # Build DZ from Zernike modes (in nm, then convert to meters)
    dz_nm = np.zeros_like(x_m)
    for j, amp_nm in coefficients.items():
        Z_j = zernike_noll(j, r_norm, theta)
        dz_nm += amp_nm * Z_j

    rms_clean_nm = np.sqrt(np.mean(dz_nm**2))

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

    rms_noisy_nm = np.sqrt(np.mean(dz_nm**2))

    # Convert to meters
    dz_m = dz_nm * 1e-9

    # DX, DY
    if include_lateral:
        # Small lateral displacements (~1/10 of DZ magnitude)
        scale = np.max(np.abs(dz_m)) * 0.1
        dx_m = rng.normal(0, scale, size=x_m.shape)
        dy_m = rng.normal(0, scale, size=y_m.shape)
    else:
        dx_m = np.zeros_like(x_m)
        dy_m = np.zeros_like(y_m)

    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_m),
        "diameter_mm": float(diameter_mm),
        "rms_nm_clean": float(rms_clean_nm),
        "rms_nm_with_noise": float(rms_noisy_nm),
        "noise_rms_nm": noise_rms_nm,
        "include_lateral": include_lateral,
        "seed": seed,
        "units": {
            "positions": "meters",
            "displacements": "meters",
            "coefficients": "nanometers",
        }
    }

    return dx_m, dy_m, dz_m, metadata


# ============================================================================
# Output (matching WFE_from_CSV_OPD format exactly)
# ============================================================================

def write_csv_fea(
    filepath: str,
    x_m: np.ndarray,
    y_m: np.ndarray,
    z_m: np.ndarray,
    dx_m: np.ndarray,
    dy_m: np.ndarray,
    dz_m: np.ndarray,
):
    """
    Write FEA-style CSV matching the WFE_from_CSV_OPD input format.

    Format:
        ,X,Y,Z,DX,DY,DZ
        0,x,y,z,dx,dy,dz
        1,...
    """
    filepath = Path(filepath)
    filepath.parent.mkdir(parents=True, exist_ok=True)

    with open(filepath, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['', 'X', 'Y', 'Z', 'DX', 'DY', 'DZ'])
        for i in range(len(x_m)):
            writer.writerow([
                i,
                f"{x_m[i]:.16e}",
                f"{y_m[i]:.16e}",
                f"{z_m[i]:.16e}",
                f"{dx_m[i]:.5e}",
                f"{dy_m[i]:.16e}",
                f"{dz_m[i]:.16e}",
            ])

    print(f"  Written: {filepath} ({len(x_m)} nodes)")


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},
    "large_amplitude":  {5: 500.0, 7: 300.0},
    "many_modes":       {j: 100.0 / j for j in range(5, 51)},
}


def generate_single_case(
    name: str,
    coeffs: Dict[int, float],
    output_dir: Path,
    x_m: np.ndarray,
    y_m: np.ndarray,
    z_m: np.ndarray,
    diameter_mm: float,
    noise_rms_nm: float = 0.0,
    include_lateral: bool = False,
):
    """Generate a single test case."""
    print(f"\nGenerating: {name}")
    dx, dy, dz, meta = synthesize_displacements(
        x_m, y_m, coeffs, diameter_mm,
        noise_rms_nm=noise_rms_nm,
        include_lateral=include_lateral,
    )
    write_csv_fea(output_dir / f"{name}.csv", x_m, y_m, z_m, dx, dy, dz)
    write_metadata(output_dir / f"{name}_truth.json", meta)
    return meta


def generate_test_suite(
    output_dir: str,
    inject_csv: str = None,
    diameter_mm: float = None,
    inner_radius_mm: float = 0.0,
    surface_type: str = "flat",
    roc_mm: float = 0.0,
):
    """Generate the full validation test suite."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Get mesh
    if inject_csv:
        print(f"Loading real mesh from: {inject_csv}")
        x_m, y_m, z_m = load_real_mesh(inject_csv)
        if diameter_mm is None:
            r_mm = np.sqrt(x_m**2 + y_m**2) * 1000.0
            diameter_mm = 2.0 * np.max(r_mm)
        mesh_source = f"real FEA mesh ({inject_csv})"
    else:
        if diameter_mm is None:
            diameter_mm = 1200.0
        x_m, y_m, z_m = generate_mesh(
            diameter_mm, inner_radius_mm,
            n_rings=25, surface_type=surface_type, roc_mm=roc_mm
        )
        mesh_source = f"synthetic {surface_type} (D={diameter_mm}mm)"

    print(f"\n  Mesh: {mesh_source}")
    print(f"  Nodes: {len(x_m)}")
    print(f"  Diameter: {diameter_mm:.1f} mm")

    all_cases = {}

    # 1. Single-mode tests
    print("\n=== Single-Mode Tests ===")
    for name, coeffs in SINGLE_MODE_TESTS.items():
        meta = generate_single_case(name, coeffs, output_dir, x_m, y_m, z_m, diameter_mm)
        all_cases[name] = meta

    # 2. Realistic multi-mode
    print("\n=== Realistic Multi-Mode ===")
    meta = generate_single_case("realistic_gravity", PRESET_REALISTIC, output_dir,
                                 x_m, y_m, z_m, diameter_mm)
    all_cases["realistic_gravity"] = 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"
        meta = generate_single_case(name, PRESET_REALISTIC, output_dir,
                                     x_m, y_m, z_m, diameter_mm,
                                     noise_rms_nm=noise_level)
        all_cases[name] = meta

    # 4. Edge cases
    print("\n=== Edge Case Tests ===")
    for name, coeffs in EDGE_CASE_TESTS.items():
        meta = generate_single_case(name, coeffs, output_dir, x_m, y_m, z_m, diameter_mm)
        all_cases[name] = meta

    # 5. With lateral displacement (tests rigorous OPD method)
    print("\n=== Lateral Displacement Test ===")
    meta = generate_single_case("realistic_with_lateral", PRESET_REALISTIC, output_dir,
                                 x_m, y_m, z_m, diameter_mm, include_lateral=True)
    all_cases["realistic_with_lateral"] = meta

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

    # Write suite manifest
    manifest = {
        "suite": "Zernike Pipeline Validation",
        "generated": "2026-03-09",
        "mesh_source": mesh_source,
        "diameter_mm": diameter_mm,
        "n_modes": 50,
        "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) on flat synthetic mesh
  python generate_synthetic_wfe.py --zernike "5:100" -o test_astig.csv

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

  # Inject into real M2 mesh
  python generate_synthetic_wfe.py --inject m2_real.csv --zernike "5:100,7:50" -o test.csv

  # Full suite using real mesh geometry
  python generate_synthetic_wfe.py --suite --inject m2_real.csv -d validation_suite/

  # Full suite with synthetic spherical mirror
  python generate_synthetic_wfe.py --suite --surface sphere --roc 5000 -d validation_suite/
        """
    )

    parser.add_argument("--zernike", "-z", type=str,
                        help="Zernike coefficients as 'j:amp_nm,...'")
    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("--inject", type=str,
                        help="Real FEA CSV to use as mesh source (keeps geometry, replaces displacements)")
    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=None,
                        help="Mirror diameter in mm (auto-detected from mesh if not set)")
    parser.add_argument("--inner-radius", type=float, default=0.0,
                        help="Inner radius in mm for annular aperture")
    parser.add_argument("--surface", choices=["flat", "sphere", "parabola"],
                        default="flat", help="Surface type for synthetic mesh")
    parser.add_argument("--roc", type=float, default=5000.0,
                        help="Radius of curvature in mm (for sphere/parabola)")
    parser.add_argument("--noise", type=float, default=0.0,
                        help="Gaussian noise RMS in nm")
    parser.add_argument("--lateral", action="store_true",
                        help="Include small lateral DX/DY displacements")

    args = parser.parse_args()

    print("=" * 60)
    print("Synthetic WFE Surface Generator")
    print("Zernike Pipeline Validation — GigaBIT")
    print("=" * 60)

    if args.suite:
        generate_test_suite(
            args.output_dir,
            inject_csv=args.inject,
            diameter_mm=args.diameter,
            inner_radius_mm=args.inner_radius,
            surface_type=args.surface,
            roc_mm=args.roc,
        )
    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)

        # Get mesh
        if args.inject:
            x_m, y_m, z_m = load_real_mesh(args.inject)
            diameter_mm = args.diameter
        else:
            diameter_mm = args.diameter or 1200.0
            x_m, y_m, z_m = generate_mesh(
                diameter_mm, args.inner_radius,
                n_rings=25, surface_type=args.surface, roc_mm=args.roc
            )

        print(f"  Nodes: {len(x_m)}")
        if diameter_mm:
            print(f"  Diameter: {diameter_mm:.1f} mm")

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

        dx, dy, dz, meta = synthesize_displacements(
            x_m, y_m, coeffs, diameter_mm,
            noise_rms_nm=args.noise,
            include_lateral=args.lateral,
        )

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

        write_csv_fea(args.output, x_m, y_m, z_m, dx, dy, dz)
        meta_path = Path(args.output).with_suffix('.json')
        write_metadata(str(meta_path), meta)

    print("\nDone!")


if __name__ == "__main__":
    main()