#!/usr/bin/env python3
"""
Generate PSD figures for CDR — using Atomizer's full OPD pipeline.
Ensures exact consistency with CDR WFE numbers.
"""
import sys
sys.path.insert(0, '/home/papa/repos/Atomizer')

import numpy as np
from scipy.interpolate import griddata
from numpy.fft import fft2, fftshift
import warnings
warnings.filterwarnings('ignore')

import os
os.environ['MPLBACKEND'] = 'Agg'
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from optimization_engine.extractors.extract_zernike import (
    zernike_noll, compute_zernike_coefficients, DEFAULT_FILTER_ORDERS
)

OP2_PATH = '/home/papa/obsidian-vault/2-Projects/P04-GigaBIT-M1/FEA/CDR-Optimization-V7/iter0_BEST_OF_ALL/assy_m1_assyfem1_sim1-solution_1.op2'
OUTPUT_DIR = '/home/papa/obsidian-vault/2-Projects/P04-GigaBIT-M1/CDR-Report/Images'
ANGLES = ['20', '40', '60', '90']
APERTURE_RADIUS_MM = 575.0
GRID_SIZE = 512

BANDS = {
    'LSF': (0.05, 2.0,  '#3b82f6'),
    'MSF': (2.0,  20.0, '#f59e0b'),
    'HSF': (20.0, 100.0,'#ef4444'),
}


def extract_filtered_wfe(ext, angle_str):
    """
    Use Atomizer's full OPD pipeline + Zernike filtering.
    Gets the WFE array (2× surface, figure-corrected) then removes J1-J4.
    """
    # Get the full OPD data (this is what extract_subcase uses internally)
    opd_data = ext._build_figure_opd_data(angle_str)
    
    X = opd_data['x_original']
    Y = opd_data['y_original']
    WFE = opd_data['wfe_nm']  # Already 2× surface error, figure-corrected
    
    # Now do the exact same Zernike filtering as extract_subcase
    coeffs, R_max = compute_zernike_coefficients(X, Y, WFE, n_modes=50)
    
    x_c = X - np.mean(X)
    y_c = Y - np.mean(Y)
    r = np.hypot(x_c / R_max, y_c / R_max)
    theta = np.arctan2(y_c, x_c)
    
    filter_orders = DEFAULT_FILTER_ORDERS  # 4
    Z_low = np.column_stack([
        zernike_noll(j, r, theta) for j in range(1, filter_orders + 1)
    ])
    
    wfe_filtered = WFE - Z_low @ coeffs[:filter_orders]
    
    X_c = X - np.mean(X)
    Y_c = Y - np.mean(Y)
    
    rms_filtered = np.sqrt(np.mean(wfe_filtered**2))
    
    return X_c, Y_c, wfe_filtered, rms_filtered


def compute_psd(X, Y, Z_nm, n_grid=GRID_SIZE):
    extent = APERTURE_RADIUS_MM
    x_grid = np.linspace(-extent, extent, n_grid)
    y_grid = np.linspace(-extent, extent, n_grid)
    Xg, Yg = np.meshgrid(x_grid, y_grid)
    Rg = np.sqrt(Xg**2 + Yg**2)
    mask = Rg <= APERTURE_RADIUS_MM
    
    Zg = griddata((X, Y), Z_nm, (Xg, Yg), method='linear', fill_value=0.0)
    Zg[~mask] = 0.0
    Zg = np.nan_to_num(Zg, nan=0.0)
    
    hx = np.hanning(n_grid)
    hy = np.hanning(n_grid)
    Zw = Zg * np.outer(hy, hx) * mask
    
    F = fftshift(fft2(Zw))
    psd_2d = np.abs(F)**2
    
    dx = 2 * extent / n_grid
    freqs = fftshift(np.fft.fftfreq(n_grid, d=dx))
    FX, FY = np.meshgrid(freqs, freqs)
    FR_cpa = np.sqrt(FX**2 + FY**2) * (2 * APERTURE_RADIUS_MM)
    
    bin_edges = np.logspace(np.log10(0.03), np.log10(n_grid/4), 80)
    freqs_out, psd_out = [], []
    for i in range(len(bin_edges) - 1):
        sel = (FR_cpa >= bin_edges[i]) & (FR_cpa < bin_edges[i+1]) & mask
        if np.sum(sel) > 2:
            psd_out.append(np.mean(psd_2d[sel]) * dx**2)
            freqs_out.append(np.sqrt(bin_edges[i] * bin_edges[i+1]))
    
    return np.array(freqs_out), np.array(psd_out), Zg, mask


def compute_band_rms(X, Y, Z_nm, f_lo, f_hi, n_grid=GRID_SIZE):
    extent = APERTURE_RADIUS_MM
    x_grid = np.linspace(-extent, extent, n_grid)
    y_grid = np.linspace(-extent, extent, n_grid)
    Xg, Yg = np.meshgrid(x_grid, y_grid)
    Rg = np.sqrt(Xg**2 + Yg**2)
    mask = Rg <= APERTURE_RADIUS_MM
    
    Zg = griddata((X, Y), Z_nm, (Xg, Yg), method='linear', fill_value=0.0)
    Zg[~mask] = 0.0
    Zg = np.nan_to_num(Zg, nan=0.0)
    
    F = np.fft.fft2(Zg)
    dx = 2 * extent / n_grid
    FX, FY = np.meshgrid(np.fft.fftfreq(n_grid, d=dx), np.fft.fftfreq(n_grid, d=dx))
    FR_cpa = np.sqrt(FX**2 + FY**2) * (2 * APERTURE_RADIUS_MM)
    
    band = (FR_cpa >= f_lo) & (FR_cpa < f_hi)
    F_band = np.zeros_like(F)
    F_band[band] = F[band]
    Z_band = np.real(np.fft.ifft2(F_band))
    
    return np.sqrt(np.mean(Z_band[mask]**2))


def main():
    from optimization_engine.extractors.extract_zernike_figure import ZernikeOPDExtractor
    
    print("Loading FEA data (Atomizer full OPD pipeline)...")
    ext = ZernikeOPDExtractor(OP2_PATH)
    
    # Reference values from Atomizer
    all_subcases = ext.extract_all_subcases()
    print("\n  Atomizer reference WFE (filtered_rms_nm):")
    for a in ANGLES:
        print(f"    {a}°: {all_subcases[int(a)]['filtered_rms_nm']:.2f} nm")
    
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    all_results = {}
    
    print("\n  Computing PSD from Atomizer OPD pipeline:")
    for angle in ANGLES:
        print(f"\n  {angle}°:")
        X, Y, Z_filt, rms_filt = extract_filtered_wfe(ext, angle)
        ref_rms = all_subcases[int(angle)]['filtered_rms_nm']
        print(f"    Our filtered RMS: {rms_filt:.2f} nm | Atomizer ref: {ref_rms:.2f} nm | Δ: {abs(rms_filt - ref_rms):.2f} nm")
        
        freqs, psd, Zg, mask = compute_psd(X, Y, Z_filt)
        band_rms = {bname: compute_band_rms(X, Y, Z_filt, flo, fhi) for bname, (flo, fhi, _) in BANDS.items()}
        
        all_results[angle] = {
            'freqs': freqs, 'psd': psd, 'Zg': Zg, 'mask': mask,
            'band_rms': band_rms, 'total_rms': rms_filt, 'ref_rms': ref_rms,
        }
        print(f"    LSF: {band_rms['LSF']:.2f}  |  MSF: {band_rms['MSF']:.2f}  |  HSF: {band_rms['HSF']:.2f} nm")
    
    # ── FIGURES ──
    print("\n--- Generating figures ---")
    ca = {'20': '#22c55e', '40': '#3b82f6', '60': '#f59e0b', '90': '#ef4444'}
    
    # FIG 1: PSD overlay
    fig, ax = plt.subplots(figsize=(12, 7))
    for angle in ANGLES:
        r = all_results[angle]
        v = r['psd'] > 0
        ax.loglog(r['freqs'][v], r['psd'][v], '-', color=ca[angle], linewidth=2, alpha=0.85,
                  label=f"{angle}° (WFE = {r['total_rms']:.1f} nm RMS)")
    for bname, (flo, fhi, color) in BANDS.items():
        ax.axvspan(flo, fhi, alpha=0.08, color=color, zorder=0)
    ax.text(0.3, 0.02, 'LSF', transform=ax.transAxes, fontsize=12, fontweight='bold', color='#3b82f6', alpha=0.7, ha='center')
    ax.text(0.6, 0.02, 'MSF', transform=ax.transAxes, fontsize=12, fontweight='bold', color='#f59e0b', alpha=0.7, ha='center')
    ax.text(0.88, 0.02, 'HSF', transform=ax.transAxes, fontsize=12, fontweight='bold', color='#ef4444', alpha=0.7, ha='center')
    ax.set_xlabel('Spatial Frequency [cycles / aperture]', fontsize=13)
    ax.set_ylabel('Power Spectral Density [nm²·mm²]', fontsize=13)
    ax.set_title('Residual WFE Power Spectral Density — Tony Hull Methodology\nM1 Mirror, CDR V7 (Zernike J1–J4 corrected)', fontsize=14, fontweight='bold')
    ax.legend(fontsize=11, loc='upper right', framealpha=0.9)
    ax.set_xlim(0.05, 100); ax.grid(True, which='both', alpha=0.3); ax.tick_params(labelsize=11)
    plt.tight_layout()
    fig.savefig(os.path.join(OUTPUT_DIR, 'psd_multi_angle.png'), dpi=200, bbox_inches='tight')
    plt.close(fig); print("  ✓ psd_multi_angle.png")
    
    # FIG 2: Band decomposition
    fig, ax = plt.subplots(figsize=(10, 6))
    x_pos = np.arange(len(ANGLES)); width = 0.22
    for i, (bname, (flo, fhi, color)) in enumerate(BANDS.items()):
        vals = [all_results[a]['band_rms'][bname] for a in ANGLES]
        bars = ax.bar(x_pos + i*width - width, vals, width, label=f'{bname} ({flo}–{fhi} c/a)',
                      color=color, alpha=0.8, edgecolor='white', linewidth=0.5)
        for bar, val in zip(bars, vals):
            if val > 0.3:
                ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
                        f'{val:.1f}', ha='center', va='bottom', fontsize=9, fontweight='bold')
    totals = [all_results[a]['total_rms'] for a in ANGLES]
    ax.plot(x_pos, totals, 'ks--', markersize=8, linewidth=1.5, label='Total WFE (J5+)',
            markeredgecolor='white', markeredgewidth=1, zorder=5)
    for xp, t in zip(x_pos, totals):
        ax.text(xp, t + 0.8, f'{t:.1f}', ha='center', va='bottom', fontsize=9, fontweight='bold', color='#333')
    ax.set_xticks(x_pos); ax.set_xticklabels([f'{a}°' for a in ANGLES], fontsize=12)
    ax.set_xlabel('Elevation Angle', fontsize=13); ax.set_ylabel('RMS WFE [nm]', fontsize=13)
    ax.set_title('Residual WFE Band Decomposition (LSF / MSF / HSF)\nZernike J1–J4 Corrected (Piston, Tilt, Defocus)', fontsize=14, fontweight='bold')
    ax.legend(fontsize=10, loc='upper left'); ax.grid(True, axis='y', alpha=0.3); ax.tick_params(labelsize=11)
    plt.tight_layout()
    fig.savefig(os.path.join(OUTPUT_DIR, 'psd_band_decomposition.png'), dpi=200, bbox_inches='tight')
    plt.close(fig); print("  ✓ psd_band_decomposition.png")
    
    # FIG 3: 40° composite
    r40 = all_results['40']
    fig, axes = plt.subplots(1, 2, figsize=(16, 7))
    ax1 = axes[0]
    Zp = np.where(r40['mask'], r40['Zg'], np.nan)
    vmax = np.nanpercentile(np.abs(Zp), 99)
    im = ax1.imshow(Zp, extent=[-APERTURE_RADIUS_MM, APERTURE_RADIUS_MM, -APERTURE_RADIUS_MM, APERTURE_RADIUS_MM],
                    cmap='RdBu_r', vmin=-vmax, vmax=vmax, origin='lower')
    plt.colorbar(im, ax=ax1, shrink=0.8, label='Residual WFE [nm]')
    ax1.set_xlabel('X [mm]', fontsize=12); ax1.set_ylabel('Y [mm]', fontsize=12)
    ax1.set_title(f'Residual WFE at 40° Elevation\n(J1–J4 removed, RMS = {r40["total_rms"]:.1f} nm)', fontsize=13, fontweight='bold')
    ax1.set_aspect('equal')
    
    ax2 = axes[1]
    v = r40['psd'] > 0
    ax2.loglog(r40['freqs'][v], r40['psd'][v], '-', color='#3b82f6', linewidth=2.5)
    for bname, (flo, fhi, color) in BANDS.items():
        ax2.axvspan(flo, fhi, alpha=0.1, color=color)
        ax2.annotate(f'{bname}\n{r40["band_rms"][bname]:.1f} nm',
                     xy=(np.sqrt(flo*fhi), 1), xycoords=('data', 'axes fraction'),
                     fontsize=10, fontweight='bold', color=color,
                     ha='center', va='top', xytext=(0, -10), textcoords='offset points')
    ax2.set_xlabel('Spatial Frequency [cycles/aperture]', fontsize=12); ax2.set_ylabel('PSD [nm²·mm²]', fontsize=12)
    ax2.set_title('Power Spectral Density at 40°\nTony Hull Method', fontsize=13, fontweight='bold')
    ax2.set_xlim(0.05, 100); ax2.grid(True, which='both', alpha=0.3)
    plt.tight_layout()
    fig.savefig(os.path.join(OUTPUT_DIR, 'psd_40deg_composite.png'), dpi=200, bbox_inches='tight')
    plt.close(fig); print("  ✓ psd_40deg_composite.png")
    
    # FIG 4: Trajectory
    fig, ax = plt.subplots(figsize=(10, 6))
    an = [int(a) for a in ANGLES]
    for bname, (flo, fhi, color) in BANDS.items():
        ax.plot(an, [all_results[a]['band_rms'][bname] for a in ANGLES], 'o-', color=color, linewidth=2.5, markersize=8,
                label=f'{bname} ({flo}–{fhi} c/a)', markeredgecolor='white', markeredgewidth=1)
    ax.plot(an, [all_results[a]['total_rms'] for a in ANGLES], 's--', color='#1e293b', linewidth=2, markersize=8,
            label='Total WFE (J5+)', markeredgecolor='white', markeredgewidth=1)
    ax.set_xlabel('Elevation Angle [°]', fontsize=13); ax.set_ylabel('RMS WFE [nm]', fontsize=13)
    ax.set_title('Residual WFE PSD Band Trajectory vs Elevation\nZernike J1–J4 Corrected', fontsize=14, fontweight='bold')
    ax.legend(fontsize=11); ax.grid(True, alpha=0.3); ax.set_xticks(an); ax.tick_params(labelsize=11)
    plt.tight_layout()
    fig.savefig(os.path.join(OUTPUT_DIR, 'psd_trajectory.png'), dpi=200, bbox_inches='tight')
    plt.close(fig); print("  ✓ psd_trajectory.png")
    
    # Summary
    print("\n" + "="*75)
    print("RESIDUAL WFE PSD — Tony Hull Method (Atomizer OPD Pipeline)")
    print("="*75)
    print(f"{'Angle':>6}  {'WFE RMS':>8}  {'Ref':>8}  {'Δ':>6}  {'LSF':>8}  {'MSF':>8}  {'HSF':>8}")
    print("-"*75)
    for a in ANGLES:
        r = all_results[a]
        d = abs(r['total_rms'] - r['ref_rms'])
        print(f"{a:>5}°  {r['total_rms']:>8.2f}  {r['ref_rms']:>8.2f}  {d:>6.2f}  {r['band_rms']['LSF']:>8.2f}  {r['band_rms']['MSF']:>8.2f}  {r['band_rms']['HSF']:>8.2f}")
    print("="*75)


if __name__ == '__main__':
    main()