#!/usr/bin/env python3
"""
Quick PSD implementation for Tony Hull's WFE approach
This adds PSD analysis to replace the LSF/MSF metrics
"""

import numpy as np
from scipy.interpolate import griddata
from scipy.fft import fft2, fftfreq
from typing import Tuple
import numpy as np

def compute_surface_psd(X: np.ndarray, Y: np.ndarray, Z: np.ndarray, 
                       aperture_radius: float) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute surface PSD using Tony Hull's methodology.
    
    Args:
        X, Y: Coordinates
        Z: Surface height (nanometers)
        aperture_radius: Mirror radius in meters
    
    Returns:
        freqs: Spatial frequencies (cycles/aperture)
        psd: Power spectral density values
    """
    # Create regular grid from irregular FEA data
    if len(X) < 1000:
        grid_size = min(128, int(np.sqrt(len(X))))
    else:
        grid_size = 256
    
    # Build regular grid
    x_range = np.linspace(X.min(), X.max(), grid_size)
    y_range = np.linspace(Y.min(), Y.max(), grid_size)
    X_grid, Y_grid = np.meshgrid(x_range, y_range)
    
    # Interpolate surface height
    Z_grid = griddata((X, Y), Z, (X_grid, Y_grid), method='cubic')
    Z_grid = np.nan_to_num(Z_grid, nan=0.0)
    
    # Apply Hann window to reduce edge effects
    hann_window = np.outer(np.hanning(grid_size), np.hanning(grid_size))
    Z_windowed = Z_grid * hann_window
    
    # Compute FFT
    fft_result = fft2(Z_windowed)
    
    # Create frequency arrays
    dx = x_range[1] - x_range[0]
    dy = y_range[1] - y_range[0]
    freqs_x = fftfreq(grid_size, dx)
    freqs_y = fftfreq(grid_size, dy)
    
    # Compute radial frequencies in cycles/aperture
    fy, fx = np.meshgrid(freqs_y, freqs_x)
    freqs_radial = np.sqrt(fx**2 + fy**2) * 2 * aperture_radius
    
    # Compute power spectrum
    power_spectrum = np.abs(fft_result)**2 
    
    # Create radial bins for averaging
    bin_edges = np.logspace(-2, 2, 50)  # Logarithmic spacing
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
    
    psd_radial = []
    for i in range(len(bin_edges)-1):
        mask = (freqs_radial >= bin_edges[i]) & (freqs_radial < bin_edges[i+1])
        if np.any(mask):
            psd_radial.append(np.mean(power_spectrum[mask]) * (dx * dy))
        else:
            psd_radial.append(0.0)
    
    return bin_centers, np.array(psd_radial)

def compute_psd_from_file(op2_path: str, angle_name: str, inner_radius: float = 135.75) -> Tuple[np.ndarray, np.ndarray, float]:
    """
    Compute PSD from FEA data file.
    
    Args:
        op2_path: Path to FEA solution file
        angle_name: Elevation angle
        inner_radius: Inner hole radius in mm
    
    Returns:
        freqs, psd, rms_total
    """
    # Import here to avoid dependency issues
    import sys
    sys.path.insert(0, '/home/papa/repos/Atomizer')
    from optimization_engine.extractors.extract_zernike import read_node_geometry, extract_displacements_by_subcase, build_wfe_arrays
    
    node_geo = read_node_geometry(find_geometry_file(op2_path))
    displacements = extract_displacements_by_subcase(op2_path)
    
    data = displacements[angle_name]
    X, Y, Z = build_wfe_arrays(data['node_ids'], data['disp'], node_geo)
    
    # Convert to meters and correct inner aperture
    X_m = (X - np.mean(X)) / 1000
    Y_m = (Y - np.mean(Y)) / 1000
    Z_nm = Z * 1000  # Convert to nanometers
    
    # Mask for inner hole
    r = np.sqrt(X_m**2 + Y_m**2)
    mask = r * 1000 <= 0.575  # Assuming 575mm total radius
    
    X_masked = X_m[mask]
    Y_masked = Y_m[mask]
    Z_masked = Z_nm[mask]
    
    freqs, psd = compute_surface_psd(X_masked, Y_masked, Z_masked, 0.575)
    rms_total = np.sqrt(np.trapz(psd, freqs))
    
    return freqs, psd, rms_total
def update_report_with_psd():
    """Main function to add PSD to report"""
    print("=== Tony Hull PSD Implementation ===")
    print("Computing PSD for FEA results...")
    
    op2_path = "/home/papa/obsidian-vault/2-Projects/P04-GigaBIT-M1/Technical-Analysis/Tensor_Project/assy_m1_assyfem1_sim1-solution_1.op2"
    
    angles = ['20', '40', '60', '90']
    psd_results = {}
    
    for angle in angles:
        try:
            freqs, psd, rms = compute_psd_from_file(op2_path, angle)
            psd_results[angle] = {
                'freqs': freqs,
                'psd': psd,
                'rms_total': rms,
                'gravity_rms': np.sqrt(np.trapz(psd[(freqs >= 0.1) & (freqs <= 2)], freqs[(freqs >= 0.1) & (freqs <= 2)])) if np.any((freqs >= 0.1) & (freqs <= 2)) else 0,
                'support_rms': np.sqrt(np.trapz(psd[(freqs >= 2) & (freqs <= 20)], freqs[(freqs >= 2) & (freqs <= 20)])) if np.any((freqs >= 2) & (freqs <= 20)) else 0
            }
            print(f"  {angle}°: RMS={rms:.2f}nm Gravity={psd_results[angle]['gravity_rms']:.2f}nm Support={psd_results[angle]['support_rms']:.2f}nm")
        except Exception as e:
            print(f"  Error processing {angle}°: {e}")
    
    print(f"\nTony Hull PSD complete for {len(psd_results)} angles")
    return psd_results

if __name__ == "__main__":
    result = update_report_with_psd()