#!/usr/bin/env python3
"""
Generate Zernike Validation Report Figures
==========================================

Creates publication-quality figures for the Zernike pipeline validation report.
Outputs PNG figures to a specified directory.

Author: Mario (Atomizer V&V)
Created: 2026-03-09
"""

import sys
import json
import csv
from pathlib import Path
from math import factorial
from typing import Dict, Tuple
import numpy as np
from numpy.linalg import lstsq
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import TwoSlopeNorm

# ============================================================================
# Zernike Math
# ============================================================================

def noll_indices(j):
    if j < 1:
        raise ValueError
    count, n = 0, 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, 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, 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):
    n, m = noll_indices(j)
    names = {
        (0,0):"Piston",(1,-1):"Tilt X",(1,1):"Tilt Y",
        (2,0):"Defocus",(2,-2):"Astig 45°",(2,2):"Astig 0°",
        (3,-1):"Coma X",(3,1):"Coma Y",(3,-3):"Trefoil X",(3,3):"Trefoil Y",
        (4,0):"Spherical",(4,-2):"2nd Astig X",(4,2):"2nd Astig Y",
        (6,0):"2nd Sph",
    }
    return names.get((n, m), f"Z{j}")

def fit_zernike(x_m, y_m, dz_m, diameter_mm=None, n_modes=50):
    if diameter_mm is None:
        diameter_mm = 2.0 * np.max(np.sqrt(x_m**2 + y_m**2) * 1000.0)
    outer_r_m = diameter_mm / 2000.0
    r_norm = np.sqrt(x_m**2 + y_m**2) / outer_r_m
    theta = np.arctan2(y_m, x_m)
    dz_nm = dz_m * 1e9
    Z = np.zeros((len(x_m), n_modes))
    for j in range(1, n_modes + 1):
        Z[:, j-1] = zernike_noll(j, r_norm, theta)
    coeffs, _, _, _ = lstsq(Z, dz_nm, rcond=None)
    return coeffs

def read_fea_csv(path):
    x, y, z, dx, dy, dz = [], [], [], [], [], []
    with open(path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            x.append(float(row['X'])); y.append(float(row['Y'])); z.append(float(row['Z']))
            dx.append(float(row['DX'])); dy.append(float(row['DY'])); dz.append(float(row['DZ']))
    return np.array(x), np.array(y), np.array(z), np.array(dx), np.array(dy), np.array(dz)

# ============================================================================
# Figure Generation
# ============================================================================

def fig_single_mode_bar_chart(suite_dir, output_dir):
    """Bar chart showing input vs recovered for each single-mode test."""
    single_modes = ["Z05_astig_0deg", "Z06_astig_45deg", "Z07_coma_x", "Z08_coma_y",
                    "Z09_trefoil_x", "Z10_trefoil_y", "Z11_spherical", "Z22_2nd_spherical", "Z37_high_order"]

    fig, ax = plt.subplots(figsize=(14, 6))

    labels = []
    input_vals = []
    recovered_vals = []
    errors = []

    for name in single_modes:
        csv_path = suite_dir / f"{name}.csv"
        truth_path = suite_dir / f"{name}_truth.json"
        if not csv_path.exists():
            continue

        x, y, z, dx, dy, dz = read_fea_csv(csv_path)
        with open(truth_path) as f:
            truth = json.load(f)

        coeffs_in = {int(k): v for k, v in truth["input_coefficients"].items()}
        recovered = fit_zernike(x, y, dz, truth.get("diameter_mm"))

        for j, amp in coeffs_in.items():
            labels.append(f"Z{j}\n{zernike_name(j)}")
            input_vals.append(amp)
            recovered_vals.append(recovered[j-1])
            errors.append(abs(recovered[j-1] - amp))

    x_pos = np.arange(len(labels))
    width = 0.35

    bars1 = ax.bar(x_pos - width/2, input_vals, width, label='Input (Truth)', color='#2196F3', alpha=0.9)
    bars2 = ax.bar(x_pos + width/2, recovered_vals, width, label='Recovered', color='#FF9800', alpha=0.9)

    ax.set_xlabel('Zernike Mode', fontsize=12)
    ax.set_ylabel('Coefficient Amplitude (nm)', fontsize=12)
    ax.set_title('Phase 1: Single-Mode Validation — Input vs Recovered\n(Real M2 Mesh, 357 nodes)', fontsize=14, fontweight='bold')
    ax.set_xticks(x_pos)
    ax.set_xticklabels(labels, fontsize=9)
    ax.legend(fontsize=11)
    ax.grid(axis='y', alpha=0.3)

    # Add error annotations
    for i, err in enumerate(errors):
        ax.annotate(f'Δ={err:.1e}', (x_pos[i], max(input_vals[i], recovered_vals[i]) + 2),
                   ha='center', fontsize=7, color='green')

    plt.tight_layout()
    fig.savefig(output_dir / "fig1_single_mode_validation.png", dpi=150, bbox_inches='tight')
    plt.close()
    print("  ✅ fig1_single_mode_validation.png")


def fig_multi_mode_comparison(suite_dir, output_dir):
    """Multi-mode realistic test: coefficient comparison."""
    csv_path = suite_dir / "realistic_gravity.csv"
    truth_path = suite_dir / "realistic_gravity_truth.json"

    x, y, z, dx, dy, dz = read_fea_csv(csv_path)
    with open(truth_path) as f:
        truth = json.load(f)

    coeffs_in = {int(k): v for k, v in truth["input_coefficients"].items()}
    recovered = fit_zernike(x, y, dz, truth.get("diameter_mm"))

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

    # Left: Bar chart of all 50 modes
    modes = np.arange(1, 51)
    input_arr = np.array([coeffs_in.get(j, 0.0) for j in modes])
    recov_arr = np.array([recovered[j-1] for j in modes])

    ax1.bar(modes - 0.2, input_arr, 0.4, label='Input', color='#2196F3', alpha=0.8)
    ax1.bar(modes + 0.2, recov_arr, 0.4, label='Recovered', color='#FF9800', alpha=0.8)
    ax1.set_xlabel('Noll Index j', fontsize=11)
    ax1.set_ylabel('Amplitude (nm)', fontsize=11)
    ax1.set_title('Multi-Mode Realistic Test\n(9 modes, gravity-like pattern)', fontsize=12, fontweight='bold')
    ax1.legend()
    ax1.grid(axis='y', alpha=0.3)

    # Right: Error for each mode
    error = np.abs(recov_arr - input_arr)
    colors = ['#4CAF50' if e < 0.5 else '#F44336' for e in error]
    ax2.bar(modes, error, color=colors, alpha=0.8)
    ax2.axhline(y=0.5, color='red', linestyle='--', alpha=0.5, label='Tolerance (0.5 nm)')
    ax2.set_xlabel('Noll Index j', fontsize=11)
    ax2.set_ylabel('|Error| (nm)', fontsize=11)
    ax2.set_title('Coefficient Recovery Error\n(all modes ≤ machine epsilon)', fontsize=12, fontweight='bold')
    ax2.legend()
    ax2.grid(axis='y', alpha=0.3)
    ax2.set_yscale('symlog', linthresh=1e-15)
    ax2.set_ylim(-1e-16, 1)

    plt.tight_layout()
    fig.savefig(output_dir / "fig2_multimode_validation.png", dpi=150, bbox_inches='tight')
    plt.close()
    print("  ✅ fig2_multimode_validation.png")


def fig_noise_sensitivity(suite_dir, output_dir):
    """Noise sensitivity analysis: error vs noise level."""
    noise_levels = [0, 1, 5, 10]
    names = ["realistic_gravity", "realistic_noise_1nm", "realistic_noise_5nm", "realistic_noise_10nm"]

    active_modes = [5, 6, 7, 8, 9, 11, 13, 16, 22]
    mode_labels = [f"Z{j}" for j in active_modes]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

    # Collect errors per noise level per mode
    all_errors = {}
    max_errors = []
    rms_errors = []

    for noise, name in zip(noise_levels, names):
        csv_path = suite_dir / f"{name}.csv"
        truth_path = suite_dir / f"{name}_truth.json"

        x, y, z, dx, dy, dz = read_fea_csv(csv_path)
        with open(truth_path) as f:
            truth = json.load(f)

        coeffs_in = {int(k): v for k, v in truth["input_coefficients"].items()}
        recovered = fit_zernike(x, y, dz, truth.get("diameter_mm"))

        mode_errors = []
        for j in active_modes:
            err = abs(recovered[j-1] - coeffs_in.get(j, 0.0))
            mode_errors.append(err)

        all_errors[noise] = mode_errors
        # Max error across ALL 50 modes
        full_err = [abs(recovered[j-1] - coeffs_in.get(j+1, 0.0)) for j in range(50)]
        max_errors.append(max(full_err))
        rms_errors.append(np.sqrt(np.mean(np.array(full_err)**2)))

    # Left: Error per active mode at each noise level
    x_pos = np.arange(len(active_modes))
    width = 0.2
    for i, noise in enumerate(noise_levels):
        offset = (i - 1.5) * width
        label = f"Noise={noise}nm" if noise > 0 else "Clean"
        ax1.bar(x_pos + offset, all_errors[noise], width, label=label, alpha=0.85)

    ax1.set_xlabel('Zernike Mode', fontsize=11)
    ax1.set_ylabel('|Coefficient Error| (nm)', fontsize=11)
    ax1.set_title('Noise Sensitivity: Active Mode Errors\n(357-node M2 mesh)', fontsize=12, fontweight='bold')
    ax1.set_xticks(x_pos)
    ax1.set_xticklabels(mode_labels)
    ax1.legend()
    ax1.grid(axis='y', alpha=0.3)

    # Right: Max error vs noise level
    ax2.plot(noise_levels, max_errors, 'o-', color='#F44336', linewidth=2, markersize=8, label='Max error (any mode)')
    ax2.plot(noise_levels, rms_errors, 's-', color='#2196F3', linewidth=2, markersize=8, label='RMS error (all modes)')
    ax2.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5, label='Tolerance (0.5 nm)')
    ax2.set_xlabel('Input Noise RMS (nm)', fontsize=11)
    ax2.set_ylabel('Error (nm)', fontsize=11)
    ax2.set_title('Error Growth vs Noise Level\n(357 nodes — sparse mesh amplifies noise)', fontsize=12, fontweight='bold')
    ax2.legend()
    ax2.grid(alpha=0.3)

    plt.tight_layout()
    fig.savefig(output_dir / "fig3_noise_sensitivity.png", dpi=150, bbox_inches='tight')
    plt.close()
    print("  ✅ fig3_noise_sensitivity.png")


def fig_surface_maps(suite_dir, output_dir):
    """2D surface maps showing input OPD, recovered, and residual."""
    csv_path = suite_dir / "realistic_gravity.csv"
    truth_path = suite_dir / "realistic_gravity_truth.json"

    x, y, z, dx, dy, dz = read_fea_csv(csv_path)
    with open(truth_path) as f:
        truth = json.load(f)

    diameter_mm = truth.get("diameter_mm")
    coeffs_in = {int(k): v for k, v in truth["input_coefficients"].items()}
    recovered = fit_zernike(x, y, dz, diameter_mm)

    # Compute surfaces in nm
    dz_nm = dz * 1e9

    outer_r_m = diameter_mm / 2000.0
    r_norm = np.sqrt(x**2 + y**2) / outer_r_m
    theta = np.arctan2(y, x)

    # Reconstructed surface
    recon_nm = np.zeros_like(x)
    for j in range(1, 51):
        recon_nm += recovered[j-1] * zernike_noll(j, r_norm, theta)

    residual_nm = dz_nm - recon_nm

    fig, axes = plt.subplots(1, 3, figsize=(18, 5.5))

    x_mm = x * 1000
    y_mm = y * 1000

    for ax, data, title, cmap in zip(
        axes,
        [dz_nm, recon_nm, residual_nm],
        ['Input OPD (DZ)', 'Reconstructed (50 modes)', 'Residual'],
        ['RdBu_r', 'RdBu_r', 'RdBu_r']
    ):
        vmax = max(abs(data.max()), abs(data.min())) if np.any(data != 0) else 1
        if vmax < 1e-10:
            vmax = 1e-10
        norm = TwoSlopeNorm(vmin=-vmax, vcenter=0, vmax=vmax)
        sc = ax.scatter(x_mm, y_mm, c=data, cmap=cmap, norm=norm, s=15, edgecolors='none')
        ax.set_aspect('equal')
        ax.set_xlabel('X (mm)')
        ax.set_ylabel('Y (mm)')
        ax.set_title(f'{title}\nRMS: {np.sqrt(np.mean(data**2)):.4f} nm')
        plt.colorbar(sc, ax=ax, label='nm', shrink=0.8)

    fig.suptitle('Phase 1: Realistic Multi-Mode Surface Validation\n(Real M2 Mesh, 357 nodes, 308mm diameter)',
                 fontsize=13, fontweight='bold', y=1.02)
    plt.tight_layout()
    fig.savefig(output_dir / "fig4_surface_maps.png", dpi=150, bbox_inches='tight')
    plt.close()
    print("  ✅ fig4_surface_maps.png")


def fig_mesh_visualization(suite_dir, output_dir):
    """Show the real M2 mesh node distribution."""
    csv_path = suite_dir / "realistic_gravity.csv"
    x, y, z, dx, dy, dz = read_fea_csv(csv_path)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

    # Left: Node positions (XY plane)
    x_mm = x * 1000
    y_mm = y * 1000
    ax1.scatter(x_mm, y_mm, c='#2196F3', s=8, alpha=0.7)
    ax1.set_aspect('equal')
    ax1.set_xlabel('X (mm)', fontsize=11)
    ax1.set_ylabel('Y (mm)', fontsize=11)
    ax1.set_title(f'M2 Mirror FEA Mesh — {len(x)} Nodes\n(Normand Fullum / Optiques Fullum)', fontsize=12, fontweight='bold')
    ax1.grid(alpha=0.2)

    # Add diameter annotation
    r_mm = np.sqrt(x_mm**2 + y_mm**2)
    circle = plt.Circle((0, 0), r_mm.max(), fill=False, color='red', linestyle='--', alpha=0.5)
    ax1.add_patch(circle)
    ax1.annotate(f'Ø {r_mm.max()*2:.0f} mm', (r_mm.max()*0.7, r_mm.max()*0.9), color='red', fontsize=10)

    # Right: Surface profile (Z vs radius)
    z_mm = z * 1000
    ax2.scatter(r_mm, z_mm, c='#FF5722', s=8, alpha=0.7)
    ax2.set_xlabel('Radial Distance (mm)', fontsize=11)
    ax2.set_ylabel('Surface Height Z (mm)', fontsize=11)
    ax2.set_title(f'Surface Profile — Sag: {z_mm.max()-z_mm.min():.2f} mm', fontsize=12, fontweight='bold')
    ax2.grid(alpha=0.3)

    plt.tight_layout()
    fig.savefig(output_dir / "fig5_mesh_visualization.png", dpi=150, bbox_inches='tight')
    plt.close()
    print("  ✅ fig5_mesh_visualization.png")


def fig_suite_summary(suite_dir, output_dir):
    """Summary table-like figure showing all test results."""
    manifest_path = suite_dir / "suite_manifest.json"
    with open(manifest_path) as f:
        manifest = json.load(f)

    test_names = []
    statuses = []
    max_errors = []

    for name, meta in sorted(manifest["cases"].items()):
        csv_path = suite_dir / f"{name}.csv"
        truth_path = suite_dir / f"{name}_truth.json"
        if not csv_path.exists():
            continue

        x, y, z, dx, dy, dz = read_fea_csv(csv_path)
        with open(truth_path) as f:
            truth = json.load(f)

        coeffs_in = {int(k): v for k, v in truth["input_coefficients"].items()}
        recovered = fit_zernike(x, y, dz, truth.get("diameter_mm"))

        max_err = 0
        for j in range(1, 51):
            err = abs(recovered[j-1] - coeffs_in.get(j, 0.0))
            max_err = max(max_err, err)

        test_names.append(name.replace("_", "\n"))
        max_errors.append(max_err)
        statuses.append("PASS" if max_err < 0.5 else "NOISE")

    fig, ax = plt.subplots(figsize=(16, 6))

    colors = ['#4CAF50' if s == "PASS" else '#FF9800' for s in statuses]
    bars = ax.barh(range(len(test_names)), max_errors, color=colors, alpha=0.85)
    ax.axvline(x=0.5, color='red', linestyle='--', alpha=0.6, label='Tolerance (0.5 nm)')
    ax.set_yticks(range(len(test_names)))
    ax.set_yticklabels(test_names, fontsize=8)
    ax.set_xlabel('Max Coefficient Error (nm)', fontsize=11)
    ax.set_title('Validation Suite Summary — All Test Cases\n(Green = PASS, Orange = Expected noise degradation)',
                 fontsize=13, fontweight='bold')
    ax.legend(fontsize=10)
    ax.grid(axis='x', alpha=0.3)
    ax.set_xscale('symlog', linthresh=1e-14)

    # Add value labels
    for i, (bar, err, status) in enumerate(zip(bars, max_errors, statuses)):
        label = f"{err:.2e} nm — {status}"
        ax.text(max(err, 1e-15) * 1.5, i, label, va='center', fontsize=8)

    plt.tight_layout()
    fig.savefig(output_dir / "fig6_suite_summary.png", dpi=150, bbox_inches='tight')
    plt.close()
    print("  ✅ fig6_suite_summary.png")


# ============================================================================
# Main
# ============================================================================

def main():
    suite_dir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("tools/validation_suite")
    output_dir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("tools/validation_report_figures")
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Generating validation report figures...")
    print(f"  Suite: {suite_dir}")
    print(f"  Output: {output_dir}")
    print()

    fig_single_mode_bar_chart(suite_dir, output_dir)
    fig_multi_mode_comparison(suite_dir, output_dir)
    fig_noise_sensitivity(suite_dir, output_dir)
    fig_surface_maps(suite_dir, output_dir)
    fig_mesh_visualization(suite_dir, output_dir)
    fig_suite_summary(suite_dir, output_dir)

    print(f"\nAll figures generated in: {output_dir}")


if __name__ == "__main__":
    main()