#!/usr/bin/env python3
"""
Atomizer Zernike HTML Generator
================================

Generates 3 interactive HTML reports for Zernike wavefront analysis:
1. 40° vs 20° (relative) - Operational angle comparison
2. 60° vs 20° (relative) - Operational angle comparison
3. 90° (Manufacturing) - Absolute with manufacturing metrics

Uses the rigorous OPD method from extract_zernike_figure.py which:
- Accounts for lateral (X, Y) displacement via interpolation
- Uses actual mesh geometry as reference (no shape assumptions)
- Provides more accurate WFE for mirror optimization

Usage:
  conda activate atomizer
  python zernike_html_generator.py "path/to/solution.op2"

  # Or with file dialog (double-click via analyze_wfe.bat)
  python zernike_html_generator.py

Output:
  Creates 3 HTML files in same directory as OP2:
  - {basename}_40_vs_20.html
  - {basename}_60_vs_20.html
  - {basename}_90_mfg.html

Author: Atomizer
Created: 2025-12-19
Updated: 2025-12-28 - Upgraded to use rigorous OPD method
"""

import sys
import os
from pathlib import Path
from math import factorial
from datetime import datetime
import numpy as np
from numpy.linalg import LinAlgError

# Add Atomizer root to path
ATOMIZER_ROOT = Path(__file__).parent.parent
if str(ATOMIZER_ROOT) not in sys.path:
    sys.path.insert(0, str(ATOMIZER_ROOT))

try:
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots
    from matplotlib.tri import Triangulation
    from pyNastran.op2.op2 import OP2
    from pyNastran.bdf.bdf import BDF
except ImportError as e:
    print(f"ERROR: Missing dependency: {e}")
    print("Run: conda activate atomizer")
    sys.exit(1)

# Import the rigorous OPD extractor
try:
    from optimization_engine.extractors.extract_zernike_figure import ZernikeOPDExtractor
    USE_OPD_METHOD = True
    print("[INFO] Using rigorous OPD method (accounts for lateral displacement)")
except ImportError:
    USE_OPD_METHOD = False
    print("[WARN] OPD extractor not available, falling back to simple Z-only method")


# ============================================================================
# Configuration
# ============================================================================
N_MODES = 50
AMP = 0.5               # visual scale for residual plot (0.5x = reduced deformation)
PANCAKE = 3.0           # Z-axis range multiplier for camera view
PLOT_DOWNSAMPLE = 10000
FILTER_LOW_ORDERS = 4   # piston, tip, tilt, defocus

# Surface plot style
COLORSCALE = 'Turbo'    # Options: 'RdBu_r', 'Viridis', 'Plasma', 'Turbo', 'Jet'
SURFACE_LIGHTING = True # Smooth shaded surface with lighting
SHOW_ZERNIKE_BAR = True

REQUIRED_SUBCASES = [90, 20, 40, 60]
REF_LABEL = "2"         # Reference subcase (20 deg) - using isubcase
POLISH_LABEL = "1"      # Manufacturing orientation (90 deg) - using isubcase

# Displacement unit in OP2 -> nm scale for WFE = 2*Disp_Z
DISP_SRC_UNIT = "mm"
NM_PER_UNIT = 1e6 if DISP_SRC_UNIT == "mm" else 1e9


# ============================================================================
# Zernike Math
# ============================================================================
def noll_indices(j: int):
    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_noll(j, r, th):
    n, m = noll_indices(j)
    R = np.zeros_like(r)
    for s in range((n-abs(m))//2 + 1):
        c = ((-1)**s * factorial(n-s) /
             (factorial(s) *
              factorial((n+abs(m))//2 - s) *
              factorial((n-abs(m))//2 - s)))
        R += c * r**(n-2*s)
    if m == 0:
        return R
    return R * (np.cos(m*th) if m > 0 else np.sin(-m*th))


def compute_zernike_coeffs(X, Y, vals, n_modes, chunk_size=100000):
    Xc, Yc = X - np.mean(X), Y - np.mean(Y)
    R = float(np.max(np.hypot(Xc, Yc)))
    r = np.hypot(Xc/R, Yc/R).astype(np.float32)
    th = np.arctan2(Yc, Xc).astype(np.float32)
    mask = (r <= 1.0) & ~np.isnan(vals)
    if not np.any(mask):
        raise RuntimeError("No valid points inside unit disk.")
    idx = np.nonzero(mask)[0]
    m = int(n_modes)
    G = np.zeros((m, m), dtype=np.float64)
    h = np.zeros((m,), dtype=np.float64)
    v = vals.astype(np.float64)

    for start in range(0, len(idx), chunk_size):
        sl = idx[start:start+chunk_size]
        r_b, th_b, v_b = r[sl], th[sl], v[sl]
        Zb = np.column_stack([zernike_noll(j, r_b, th_b).astype(np.float32)
                              for j in range(1, m+1)])
        G += (Zb.T @ Zb).astype(np.float64)
        h += (Zb.T @ v_b).astype(np.float64)
    try:
        coeffs = np.linalg.solve(G, h)
    except LinAlgError:
        coeffs = np.linalg.lstsq(G, h, rcond=None)[0]
    return coeffs, R


def zernike_common_name(n: int, m: int) -> str:
    names = {
        (0, 0): "Piston", (1, -1): "Tilt X", (1, 1): "Tilt Y",
        (2, 0): "Defocus", (2, -2): "Astig 45 deg", (2, 2): "Astig 0 deg",
        (3, -1): "Coma X", (3, 1): "Coma Y", (3, -3): "Trefoil X", (3, 3): "Trefoil Y",
        (4, 0): "Primary Spherical", (4, -2): "Sec Astig X", (4, 2): "Sec Astig Y",
        (4, -4): "Quadrafoil X", (4, 4): "Quadrafoil Y",
        (5, -1): "Sec Coma X", (5, 1): "Sec Coma Y",
        (5, -3): "Sec Trefoil X", (5, 3): "Sec Trefoil Y",
        (5, -5): "Pentafoil X", (5, 5): "Pentafoil Y",
        (6, 0): "Sec Spherical",
    }
    return names.get((n, m), f"Z(n={n}, m={m})")


def zernike_label_for_j(j: int) -> str:
    n, m = noll_indices(j)
    return f"J{j:02d} - {zernike_common_name(n, m)} (n={n}, m={m})"


# ============================================================================
# File I/O
# ============================================================================
def find_geometry_file(op2_path: Path) -> Path:
    """Find matching BDF/DAT file for an OP2."""
    folder = op2_path.parent
    base = op2_path.stem

    for ext in ['.dat', '.bdf']:
        cand = folder / (base + ext)
        if cand.exists():
            return cand

    for f in folder.iterdir():
        if f.suffix.lower() in ['.dat', '.bdf']:
            return f

    raise FileNotFoundError(f"No .dat or .bdf geometry file found for {op2_path}")


def read_geometry(dat_path: Path) -> dict:
    bdf = BDF()
    bdf.read_bdf(str(dat_path))
    return {int(nid): node.get_position() for nid, node in bdf.nodes.items()}


def read_displacements(op2_path: Path) -> dict:
    """Read displacement data organized by subcase."""
    op2 = OP2()
    op2.read_op2(str(op2_path))

    if not op2.displacements:
        raise RuntimeError("No displacement data found in OP2 file")

    result = {}
    for key, darr in op2.displacements.items():
        data = darr.data
        dmat = data[0] if data.ndim == 3 else (data if data.ndim == 2 else None)
        if dmat is None:
            continue

        ngt = darr.node_gridtype.astype(int)
        node_ids = ngt if ngt.ndim == 1 else ngt[:, 0]

        # Get subcase identifier
        isubcase = getattr(darr, 'isubcase', None)
        label = str(isubcase) if isubcase else str(key)

        result[label] = {
            'node_ids': node_ids.astype(int),
            'disp': dmat.copy()
        }

    return result


# ============================================================================
# Data Processing
# ============================================================================
def build_wfe_arrays(label: str, node_ids, dmat, node_geo):
    """Build X, Y, WFE arrays for a subcase."""
    X, Y, WFE = [], [], []
    for nid, vec in zip(node_ids, dmat):
        geo = node_geo.get(int(nid))
        if geo is None:
            continue
        X.append(geo[0])
        Y.append(geo[1])
        wfe = vec[2] * 2.0 * NM_PER_UNIT  # Z-disp to WFE (nm)
        WFE.append(wfe)

    return np.array(X), np.array(Y), np.array(WFE)


def compute_relative_wfe(X1, Y1, WFE1, node_ids1, X2, Y2, WFE2, node_ids2):
    """Compute relative WFE: WFE1 - WFE2 for common nodes."""
    # Build mapping for reference
    ref_map = {int(nid): (x, y, w) for nid, x, y, w in zip(node_ids2, X2, Y2, WFE2)}

    X_rel, Y_rel, WFE_rel = [], [], []
    for nid, x, y, w in zip(node_ids1, X1, Y1, WFE1):
        nid = int(nid)
        if nid in ref_map:
            _, _, w_ref = ref_map[nid]
            X_rel.append(x)
            Y_rel.append(y)
            WFE_rel.append(w - w_ref)

    return np.array(X_rel), np.array(Y_rel), np.array(WFE_rel)


def compute_rms_metrics(X, Y, W_nm):
    """Compute RMS metrics and Zernike coefficients."""
    coeffs, R = compute_zernike_coeffs(X, Y, W_nm, N_MODES)

    Xc = X - np.mean(X)
    Yc = Y - np.mean(Y)
    r = np.hypot(Xc/R, Yc/R)
    th = np.arctan2(Yc, Xc)

    Z = np.column_stack([zernike_noll(j, r, th) for j in range(1, N_MODES+1)])
    W_res_filt = W_nm - Z[:, :FILTER_LOW_ORDERS].dot(coeffs[:FILTER_LOW_ORDERS])

    # J1-J3 filtered (keeping defocus - optician workload)
    W_res_filt_j1to3 = W_nm - Z[:, :3].dot(coeffs[:3])

    global_rms = float(np.sqrt(np.mean(W_nm**2)))
    filtered_rms = float(np.sqrt(np.mean(W_res_filt**2)))
    rms_filter_j1to3 = float(np.sqrt(np.mean(W_res_filt_j1to3**2)))

    return {
        'coefficients': coeffs,
        'R': R,
        'global_rms': global_rms,
        'filtered_rms': filtered_rms,
        'rms_filter_j1to3': rms_filter_j1to3,
        'W_res_filt': W_res_filt,
    }


def compute_mfg_metrics(coeffs):
    """Manufacturing aberration magnitudes from Zernike coefficients.

    Noll indexing (1-based): J1=Piston, J2=TiltX, J3=TiltY, J4=Defocus,
    J5=Astig45, J6=Astig0, J7=ComaX, J8=ComaY, J9=TrefoilX, J10=TrefoilY, J11=Spherical

    Python 0-indexed: coeffs[0]=J1, coeffs[3]=J4, etc.
    """
    # Individual mode magnitudes (RSS for paired modes)
    defocus = float(abs(coeffs[3]))  # J4
    astigmatism = float(np.sqrt(coeffs[4]**2 + coeffs[5]**2))  # RSS(J5, J6)
    coma = float(np.sqrt(coeffs[6]**2 + coeffs[7]**2))  # RSS(J7, J8)
    trefoil = float(np.sqrt(coeffs[8]**2 + coeffs[9]**2))  # RSS(J9, J10)
    spherical = float(abs(coeffs[10])) if len(coeffs) > 10 else 0.0  # J11

    # RMS of higher-order terms (J4+): sqrt(sum of squares of coefficients)
    # This is the proper Zernike-coefficient-based RMS excluding piston/tip/tilt
    higher_order_rms = float(np.sqrt(np.sum(coeffs[3:]**2)))

    return {
        'defocus_nm': defocus,
        'astigmatism_rms': astigmatism,
        'coma_rms': coma,
        'trefoil_rms': trefoil,
        'spherical_nm': spherical,
        'higher_order_rms': higher_order_rms,  # RMS of all J4+ coefficients
    }


# ============================================================================
# HTML Generation
# ============================================================================
def generate_html(
    title: str,
    X: np.ndarray,
    Y: np.ndarray,
    W_nm: np.ndarray,
    rms_data: dict,
    is_relative: bool = False,
    ref_title: str = "20 deg",
    abs_pair: tuple = None,
    is_manufacturing: bool = False,
    mfg_metrics: dict = None,
    correction_metrics: dict = None,
) -> str:
    """Generate complete HTML for Zernike visualization."""

    coeffs = rms_data['coefficients']
    global_rms = rms_data['global_rms']
    filtered_rms = rms_data['filtered_rms']
    W_res_filt = rms_data['W_res_filt']

    labels = [zernike_label_for_j(j) for j in range(1, N_MODES+1)]
    coeff_abs = np.abs(coeffs)

    # Downsample for display
    n = len(X)
    if n > PLOT_DOWNSAMPLE:
        rng = np.random.default_rng(42)
        sel = rng.choice(n, size=PLOT_DOWNSAMPLE, replace=False)
        Xp, Yp, Wp = X[sel], Y[sel], W_res_filt[sel]
    else:
        Xp, Yp, Wp = X, Y, W_res_filt

    res_amp = AMP * Wp
    max_amp = float(np.max(np.abs(res_amp))) if res_amp.size else 1.0

    # Triangulate for smooth mesh surface
    mesh_traces = []
    try:
        tri = Triangulation(Xp, Yp)
        if tri.triangles is not None and len(tri.triangles) > 0:
            i, j, k = tri.triangles.T

            # Create smooth shaded mesh with lighting
            mesh_traces.append(go.Mesh3d(
                x=Xp, y=Yp, z=res_amp,
                i=i, j=j, k=k,
                intensity=res_amp,
                colorscale=COLORSCALE,
                opacity=1.0,
                flatshading=False,  # Smooth shading (interpolated normals)
                lighting=dict(
                    ambient=0.4,
                    diffuse=0.8,
                    specular=0.3,
                    roughness=0.5,
                    fresnel=0.2
                ),
                lightposition=dict(x=100, y=200, z=300),
                showscale=True,
                colorbar=dict(
                    title=dict(text="Residual (nm)", side="right"),
                    thickness=15,
                    len=0.6,
                    tickformat=".1f"
                ),
                hovertemplate="X: %{x:.1f}<br>Y: %{y:.1f}<br>Residual: %{z:.2f} nm<extra></extra>"
            ))
    except Exception as e:
        print(f"Triangulation warning: {e}")

    # Fallback scatter only if mesh failed
    if not mesh_traces:
        mesh_traces.append(go.Scatter3d(
            x=Xp, y=Yp, z=res_amp,
            mode='markers',
            marker=dict(size=2, color=res_amp, colorscale=COLORSCALE, showscale=True),
            showlegend=False
        ))

    title_suffix = f" (relative to {ref_title})" if is_relative else " (absolute)"

    # Build figure layout based on content type
    if is_manufacturing and mfg_metrics and correction_metrics:
        # Manufacturing page: 5 rows
        fig = make_subplots(
            rows=5, cols=1,
            specs=[[{"type":"scene"}],
                   [{"type":"table"}],
                   [{"type":"table"}],
                   [{"type":"table"}],
                   [{"type":"xy"}]],
            row_heights=[0.38, 0.12, 0.12, 0.18, 0.20],
            vertical_spacing=0.025,
            subplot_titles=[
                f"<b>Surface Residual - {title}{title_suffix}</b>",
                "<b>RMS Metrics (Absolute 90 deg)</b>",
                "<b>Mode Magnitudes at 90 deg</b>",
                "<b>Pre-Correction (90 deg - 20 deg)</b>",
                "<b>|Zernike Coefficients| (nm)</b>"
            ]
        )
    elif SHOW_ZERNIKE_BAR:
        # Standard page with bar chart: 4 rows
        fig = make_subplots(
            rows=4, cols=1,
            specs=[[{"type":"scene"}],
                   [{"type":"table"}],
                   [{"type":"table"}],
                   [{"type":"xy"}]],
            row_heights=[0.45, 0.12, 0.25, 0.18],
            vertical_spacing=0.03,
            subplot_titles=[
                f"<b>Surface Residual - {title}{title_suffix}</b>",
                "<b>RMS Metrics</b>",
                f"<b>Zernike Coefficients ({N_MODES} modes)</b>",
                "<b>|Zernike Coefficients| (nm)</b>"
            ]
        )
    else:
        # Standard page without bar: 3 rows
        fig = make_subplots(
            rows=3, cols=1,
            specs=[[{"type":"scene"}],
                   [{"type":"table"}],
                   [{"type":"table"}]],
            row_heights=[0.55, 0.15, 0.30],
            vertical_spacing=0.03,
            subplot_titles=[
                f"<b>Surface Residual - {title}{title_suffix}</b>",
                "<b>RMS Metrics</b>",
                f"<b>Zernike Coefficients ({N_MODES} modes)</b>"
            ]
        )

    # Add 3D mesh traces
    for tr in mesh_traces:
        fig.add_trace(tr, row=1, col=1)

    # Configure 3D scene with better camera angle and proportional axes
    fig.update_scenes(
        camera=dict(
            eye=dict(x=1.2, y=1.2, z=0.8),  # Angled view from above
            up=dict(x=0, y=0, z=1)
        ),
        xaxis=dict(
            title="X (mm)",
            showgrid=True,
            gridcolor='rgba(128,128,128,0.3)',
            showbackground=True,
            backgroundcolor='rgba(240,240,240,0.9)'
        ),
        yaxis=dict(
            title="Y (mm)",
            showgrid=True,
            gridcolor='rgba(128,128,128,0.3)',
            showbackground=True,
            backgroundcolor='rgba(240,240,240,0.9)'
        ),
        zaxis=dict(
            title="Residual (nm)",
            range=[-max_amp * PANCAKE, max_amp * PANCAKE],
            showgrid=True,
            gridcolor='rgba(128,128,128,0.3)',
            showbackground=True,
            backgroundcolor='rgba(230,230,250,0.9)'
        ),
        aspectmode='manual',
        aspectratio=dict(x=1, y=1, z=0.4)  # Flatten Z for better visualization
    )

    # RMS table (row 2)
    if is_relative and abs_pair:
        abs_global, abs_filtered = abs_pair
        fig.add_trace(go.Table(
            header=dict(values=["<b>Metric</b>", "<b>Relative (nm)</b>", "<b>Absolute (nm)</b>"],
                       align="left", fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[
                ["Global RMS", "Filtered RMS (J1-J4 removed)"],
                [f"{global_rms:.2f}", f"{filtered_rms:.2f}"],
                [f"{abs_global:.2f}", f"{abs_filtered:.2f}"],
            ], align="left", fill_color='#374151', font=dict(color='white'))
        ), row=2, col=1)
    elif is_manufacturing and mfg_metrics and correction_metrics:
        # 90 deg absolute RMS
        fig.add_trace(go.Table(
            header=dict(values=["<b>Metric</b>", "<b>Value (nm)</b>"],
                       align="left", fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[
                ["Global RMS", "Filtered RMS (J1-J4)"],
                [f"{global_rms:.2f}", f"{filtered_rms:.2f}"]
            ], align="left", fill_color='#374151', font=dict(color='white'))
        ), row=2, col=1)

        # Mode magnitudes (row 3)
        fig.add_trace(go.Table(
            header=dict(values=["<b>Mode</b>", "<b>Value (nm)</b>"],
                       align="left", fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[
                ["Filtered RMS (J1-J3, with defocus)",
                 "Astigmatism (J5+J6)",
                 "Coma (J7+J8)",
                 "Trefoil (J9+J10)",
                 "Spherical (J11)"],
                [f"{rms_data['rms_filter_j1to3']:.2f}",
                 f"{mfg_metrics['astigmatism_rms']:.2f}",
                 f"{mfg_metrics['coma_rms']:.2f}",
                 f"{mfg_metrics['trefoil_rms']:.2f}",
                 f"{mfg_metrics['spherical_nm']:.2f}"]
            ], align="left", fill_color='#374151', font=dict(color='white'))
        ), row=3, col=1)

        # Pre-correction (row 4) - Aberrations to polish out (90° - 20°)
        # Shows what correction is needed when manufacturing at 90° to achieve 20° figure
        fig.add_trace(go.Table(
            header=dict(values=["<b>Aberration</b>", "<b>Magnitude (nm)</b>"],
                       align="left", fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[
                ["Defocus (J4)",
                 "Astigmatism (J5+J6)",
                 "Coma (J7+J8)",
                 "Trefoil (J9+J10)",
                 "Spherical (J11)"],
                [f"{correction_metrics['defocus_nm']:.2f}",
                 f"{correction_metrics['astigmatism_rms']:.2f}",
                 f"{correction_metrics['coma_rms']:.2f}",
                 f"{correction_metrics['trefoil_rms']:.2f}",
                 f"{correction_metrics['spherical_nm']:.2f}"]
            ], align="left", fill_color='#374151', font=dict(color='white'))
        ), row=4, col=1)
    else:
        # Standard absolute table
        fig.add_trace(go.Table(
            header=dict(values=["<b>Metric</b>", "<b>Value (nm)</b>"],
                       align="left", fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[
                ["Global RMS", "Filtered RMS (J1-J4 removed)"],
                [f"{global_rms:.2f}", f"{filtered_rms:.2f}"]
            ], align="left", fill_color='#374151', font=dict(color='white'))
        ), row=2, col=1)

    # Zernike coefficients table (row 3 or skip if manufacturing)
    if not (is_manufacturing and mfg_metrics and correction_metrics):
        fig.add_trace(go.Table(
            header=dict(values=["<b>Noll j</b>", "<b>Label</b>", "<b>|Coeff| (nm)</b>"],
                       align="left", fill_color='#1f2937', font=dict(color='white')),
            cells=dict(values=[
                list(range(1, N_MODES+1)),
                labels,
                [f"{c:.3f}" for c in coeff_abs]
            ], align="left", fill_color='#374151', font=dict(color='white'))
        ), row=3, col=1)

    # Bar chart (last row)
    if SHOW_ZERNIKE_BAR:
        bar_row = 5 if (is_manufacturing and mfg_metrics and correction_metrics) else 4
        fig.add_trace(
            go.Bar(
                x=coeff_abs.tolist(),
                y=labels,
                orientation='h',
                marker_color='#6366f1',
                hovertemplate="%{y}<br>|Coeff| = %{x:.3f} nm<extra></extra>",
                showlegend=False
            ),
            row=bar_row, col=1
        )

    # Layout
    height = 1500 if (is_manufacturing and mfg_metrics and correction_metrics) else 1300
    fig.update_layout(
        width=1400,
        height=height,
        margin=dict(t=60, b=20, l=20, r=20),
        paper_bgcolor='#111827',
        plot_bgcolor='#1f2937',
        font=dict(color='white'),
        title=dict(
            text=f"<b>Atomizer Zernike Analysis - {title}</b>",
            x=0.5,
            font=dict(size=18)
        )
    )

    return fig.to_html(include_plotlyjs='cdn', full_html=True)


# ============================================================================
# Main
# ============================================================================
def find_op2_file(working_dir=None):
    """Find the most recent OP2 file in the working directory."""
    if working_dir is None:
        working_dir = Path.cwd()
    else:
        working_dir = Path(working_dir)

    op2_files = list(working_dir.glob("*solution*.op2")) + list(working_dir.glob("*.op2"))

    if not op2_files:
        op2_files = list(working_dir.glob("**/*solution*.op2"))

    if not op2_files:
        return None

    return max(op2_files, key=lambda p: p.stat().st_mtime)


def main_opd(op2_path: Path):
    """Generate all 3 HTML files using rigorous OPD method."""
    print("=" * 70)
    print("  ATOMIZER ZERNIKE HTML GENERATOR (OPD METHOD)")
    print("=" * 70)
    print(f"\nOP2 File: {op2_path.name}")
    print(f"Directory: {op2_path.parent}")
    print("\n[INFO] Using OPD method: accounts for lateral (X,Y) displacement")

    # Initialize extractor
    extractor = ZernikeOPDExtractor(
        op2_path,
        displacement_unit='mm',
        n_modes=N_MODES,
        filter_orders=FILTER_LOW_ORDERS
    )

    print(f"\nAvailable subcases: {list(extractor.displacements.keys())}")

    # Map subcases (try common patterns)
    displacements = extractor.displacements
    subcase_map = {}

    if '1' in displacements and '2' in displacements:
        subcase_map = {'90': '1', '20': '2', '40': '3', '60': '4'}
    elif '90' in displacements and '20' in displacements:
        subcase_map = {'90': '90', '20': '20', '40': '40', '60': '60'}
    else:
        available = sorted(displacements.keys(), key=lambda x: int(x) if x.isdigit() else 0)
        if len(available) >= 4:
            subcase_map = {'90': available[0], '20': available[1], '40': available[2], '60': available[3]}
            print(f"[WARN] Using mapped subcases: {subcase_map}")
        else:
            print(f"[ERROR] Need 4 subcases, found: {available}")
            return

    output_dir = op2_path.parent
    base = op2_path.stem
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    html_files = []

    # ========================================================================
    # Extract absolute metrics for each subcase
    # ========================================================================
    print("\nExtracting absolute metrics (OPD method)...")

    results_abs = {}
    for angle, label in subcase_map.items():
        result = extractor.extract_subcase(label, include_coefficients=True)
        results_abs[angle] = result
        lat_disp = result.get('max_lateral_displacement_um', 0)
        print(f"  {angle} deg: Filtered RMS = {result['filtered_rms_nm']:.2f} nm, "
              f"Lateral disp max = {lat_disp:.3f} um")

    # ========================================================================
    # Extract relative metrics (40-20, 60-20, 90-20)
    # ========================================================================
    print("\nExtracting relative metrics (OPD method)...")

    # 40 vs 20
    result_40_rel = extractor.extract_relative(subcase_map['40'], subcase_map['20'], include_coefficients=True)
    print(f"  40-20: Relative Filtered RMS = {result_40_rel['relative_filtered_rms_nm']:.2f} nm")

    # 60 vs 20
    result_60_rel = extractor.extract_relative(subcase_map['60'], subcase_map['20'], include_coefficients=True)
    print(f"  60-20: Relative Filtered RMS = {result_60_rel['relative_filtered_rms_nm']:.2f} nm")

    # 90 vs 20 (for correction metrics)
    result_90_rel = extractor.extract_relative(subcase_map['90'], subcase_map['20'], include_coefficients=True)
    print(f"  90-20: Relative Filtered RMS = {result_90_rel['relative_filtered_rms_nm']:.2f} nm")

    # ========================================================================
    # Generate HTML files
    # ========================================================================

    # Helper to convert OPD results to the format expected by generate_html
    def opd_to_rms_data(result, is_relative=False):
        """Convert OPD extractor result to rms_data dict for generate_html."""
        coeffs = np.array(result.get('coefficients', [0] * N_MODES))

        # Recompute filtered residuals for visualization
        # For now, use simplified approach - the metrics are correct
        filtered_rms = result.get('relative_filtered_rms_nm' if is_relative else 'filtered_rms_nm', 0)
        global_rms = result.get('relative_global_rms_nm' if is_relative else 'global_rms_nm', 0)
        rms_j1to3 = result.get('relative_rms_filter_j1to3' if is_relative else 'rms_filter_j1to3_nm', 0)

        # We need W_res_filt for visualization - extract from diagnostic data
        # For now, create a placeholder that will be updated
        return {
            'coefficients': coeffs,
            'R': 1.0,  # Will be updated
            'global_rms': global_rms,
            'filtered_rms': filtered_rms,
            'rms_filter_j1to3': rms_j1to3,
            'W_res_filt': None,  # Will compute separately for visualization
        }

    # For visualization, we need the actual WFE arrays
    # Get diagnostic data from extractor
    print("\nGenerating HTML reports...")

    # 40 vs 20
    print("  Generating 40 deg vs 20 deg...")
    opd_40 = extractor._build_figure_opd_data(subcase_map['40'])
    opd_20 = extractor._build_figure_opd_data(subcase_map['20'])

    # Build relative WFE arrays
    ref_wfe_map = {int(nid): wfe for nid, wfe in zip(opd_20['node_ids'], opd_20['wfe_nm'])}
    X_40_rel, Y_40_rel, WFE_40_rel = [], [], []
    for i, nid in enumerate(opd_40['node_ids']):
        nid = int(nid)
        if nid in ref_wfe_map:
            X_40_rel.append(opd_40['x_deformed'][i])
            Y_40_rel.append(opd_40['y_deformed'][i])
            WFE_40_rel.append(opd_40['wfe_nm'][i] - ref_wfe_map[nid])
    X_40_rel = np.array(X_40_rel)
    Y_40_rel = np.array(Y_40_rel)
    WFE_40_rel = np.array(WFE_40_rel)

    rms_40_rel = compute_rms_metrics(X_40_rel, Y_40_rel, WFE_40_rel)
    rms_40_abs = compute_rms_metrics(opd_40['x_deformed'], opd_40['y_deformed'], opd_40['wfe_nm'])

    html_40 = generate_html(
        title="40 deg (OPD)",
        X=X_40_rel, Y=Y_40_rel, W_nm=WFE_40_rel,
        rms_data=rms_40_rel,
        is_relative=True,
        ref_title="20 deg",
        abs_pair=(rms_40_abs['global_rms'], rms_40_abs['filtered_rms'])
    )
    path_40 = output_dir / f"{base}_{timestamp}_40_vs_20.html"
    path_40.write_text(html_40, encoding='utf-8')
    html_files.append(path_40)
    print(f"    Created: {path_40.name}")

    # 60 vs 20
    print("  Generating 60 deg vs 20 deg...")
    opd_60 = extractor._build_figure_opd_data(subcase_map['60'])

    X_60_rel, Y_60_rel, WFE_60_rel = [], [], []
    for i, nid in enumerate(opd_60['node_ids']):
        nid = int(nid)
        if nid in ref_wfe_map:
            X_60_rel.append(opd_60['x_deformed'][i])
            Y_60_rel.append(opd_60['y_deformed'][i])
            WFE_60_rel.append(opd_60['wfe_nm'][i] - ref_wfe_map[nid])
    X_60_rel = np.array(X_60_rel)
    Y_60_rel = np.array(Y_60_rel)
    WFE_60_rel = np.array(WFE_60_rel)

    rms_60_rel = compute_rms_metrics(X_60_rel, Y_60_rel, WFE_60_rel)
    rms_60_abs = compute_rms_metrics(opd_60['x_deformed'], opd_60['y_deformed'], opd_60['wfe_nm'])

    html_60 = generate_html(
        title="60 deg (OPD)",
        X=X_60_rel, Y=Y_60_rel, W_nm=WFE_60_rel,
        rms_data=rms_60_rel,
        is_relative=True,
        ref_title="20 deg",
        abs_pair=(rms_60_abs['global_rms'], rms_60_abs['filtered_rms'])
    )
    path_60 = output_dir / f"{base}_{timestamp}_60_vs_20.html"
    path_60.write_text(html_60, encoding='utf-8')
    html_files.append(path_60)
    print(f"    Created: {path_60.name}")

    # 90 deg Manufacturing
    print("  Generating 90 deg Manufacturing...")
    opd_90 = extractor._build_figure_opd_data(subcase_map['90'])
    rms_90 = compute_rms_metrics(opd_90['x_deformed'], opd_90['y_deformed'], opd_90['wfe_nm'])
    mfg_metrics = compute_mfg_metrics(rms_90['coefficients'])

    # 90-20 relative for correction metrics
    X_90_rel, Y_90_rel, WFE_90_rel = [], [], []
    for i, nid in enumerate(opd_90['node_ids']):
        nid = int(nid)
        if nid in ref_wfe_map:
            X_90_rel.append(opd_90['x_deformed'][i])
            Y_90_rel.append(opd_90['y_deformed'][i])
            WFE_90_rel.append(opd_90['wfe_nm'][i] - ref_wfe_map[nid])
    X_90_rel = np.array(X_90_rel)
    Y_90_rel = np.array(Y_90_rel)
    WFE_90_rel = np.array(WFE_90_rel)
    rms_90_rel = compute_rms_metrics(X_90_rel, Y_90_rel, WFE_90_rel)

    # Get all correction metrics from Zernike coefficients (90° - 20°)
    correction_metrics = compute_mfg_metrics(rms_90_rel['coefficients'])

    html_90 = generate_html(
        title="90 deg Manufacturing (OPD)",
        X=opd_90['x_deformed'], Y=opd_90['y_deformed'], W_nm=opd_90['wfe_nm'],
        rms_data=rms_90,
        is_relative=False,
        is_manufacturing=True,
        mfg_metrics=mfg_metrics,
        correction_metrics=correction_metrics
    )
    path_90 = output_dir / f"{base}_{timestamp}_90_mfg.html"
    path_90.write_text(html_90, encoding='utf-8')
    html_files.append(path_90)
    print(f"    Created: {path_90.name}")

    # ========================================================================
    # Summary
    # ========================================================================
    print("\n" + "=" * 70)
    print("SUMMARY (OPD Method)")
    print("=" * 70)
    print(f"\nGenerated {len(html_files)} HTML files:")
    for f in html_files:
        print(f"  - {f.name}")

    print("\n" + "-" * 70)
    print("OPTIMIZATION OBJECTIVES (OPD Method)")
    print("-" * 70)
    print(f"  40-20 Filtered RMS:  {rms_40_rel['filtered_rms']:.2f} nm")
    print(f"  60-20 Filtered RMS:  {rms_60_rel['filtered_rms']:.2f} nm")
    print(f"  MFG 90 (J1-J3):      {rms_90_rel['rms_filter_j1to3']:.2f} nm")

    # Weighted sums
    ws_v4 = 5*rms_40_rel['filtered_rms'] + 5*rms_60_rel['filtered_rms'] + 2*rms_90_rel['rms_filter_j1to3']
    ws_v5 = 5*rms_40_rel['filtered_rms'] + 5*rms_60_rel['filtered_rms'] + 3*rms_90_rel['rms_filter_j1to3']
    print(f"\n  V4 Weighted Sum (5/5/2): {ws_v4:.2f}")
    print(f"  V5 Weighted Sum (5/5/3): {ws_v5:.2f}")

    # Lateral displacement summary
    print("\n" + "-" * 70)
    print("LATERAL DISPLACEMENT SUMMARY")
    print("-" * 70)
    for angle in ['20', '40', '60', '90']:
        lat = results_abs[angle].get('max_lateral_displacement_um', 0)
        print(f"  {angle} deg: max {lat:.3f} um")

    print("\n" + "=" * 70)
    print("DONE")
    print("=" * 70)

    return html_files


def main(op2_path: Path):
    """Generate all 3 HTML files (legacy Z-only method)."""
    print("=" * 70)
    print("  ATOMIZER ZERNIKE HTML GENERATOR")
    print("=" * 70)
    print(f"\nOP2 File: {op2_path.name}")
    print(f"Directory: {op2_path.parent}")

    # Find geometry
    print("\nFinding geometry file...")
    geo_path = find_geometry_file(op2_path)
    print(f"Found: {geo_path.name}")

    # Read data
    print("\nReading geometry...")
    node_geo = read_geometry(geo_path)
    print(f"Loaded {len(node_geo)} nodes")

    print("\nReading displacements...")
    displacements = read_displacements(op2_path)
    print(f"Found subcases: {list(displacements.keys())}")

    # Map subcases (try common patterns)
    # Pattern 1: Direct labels (1, 2, 3, 4)
    # Pattern 2: Angle labels (90, 20, 40, 60)
    subcase_map = {}

    if '1' in displacements and '2' in displacements:
        # Standard NX pattern: 1=90deg, 2=20deg, 3=40deg, 4=60deg
        subcase_map = {'90': '1', '20': '2', '40': '3', '60': '4'}
    elif '90' in displacements and '20' in displacements:
        # Direct angle labels
        subcase_map = {'90': '90', '20': '20', '40': '40', '60': '60'}
    else:
        # Try to use whatever is available
        available = sorted(displacements.keys(), key=lambda x: int(x) if x.isdigit() else 0)
        if len(available) >= 4:
            subcase_map = {'90': available[0], '20': available[1], '40': available[2], '60': available[3]}
            print(f"[WARN] Using mapped subcases: {subcase_map}")
        else:
            print(f"[ERROR] Need 4 subcases, found: {available}")
            return

    # Check all required subcases exist
    for angle, label in subcase_map.items():
        if label not in displacements:
            print(f"[ERROR] Subcase '{label}' (angle {angle}) not found")
            return

    output_dir = op2_path.parent
    base = op2_path.stem
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    html_files = []

    # ========================================================================
    # Build WFE arrays for reference (20 deg) and manufacturing (90 deg)
    # ========================================================================
    print("\nProcessing subcases...")

    # Reference: 20 deg
    ref_label = subcase_map['20']
    ref_data = displacements[ref_label]
    X_ref, Y_ref, WFE_ref = build_wfe_arrays(
        '20', ref_data['node_ids'], ref_data['disp'], node_geo
    )
    rms_ref = compute_rms_metrics(X_ref, Y_ref, WFE_ref)
    print(f"  20 deg: Global RMS = {rms_ref['global_rms']:.2f} nm, Filtered = {rms_ref['filtered_rms']:.2f} nm")

    # Manufacturing: 90 deg
    mfg_label = subcase_map['90']
    mfg_data = displacements[mfg_label]
    X_90, Y_90, WFE_90 = build_wfe_arrays(
        '90', mfg_data['node_ids'], mfg_data['disp'], node_geo
    )
    rms_90 = compute_rms_metrics(X_90, Y_90, WFE_90)
    mfg_metrics = compute_mfg_metrics(rms_90['coefficients'])
    print(f"  90 deg: Global RMS = {rms_90['global_rms']:.2f} nm, Filtered = {rms_90['filtered_rms']:.2f} nm")

    # ========================================================================
    # 1. Generate 40 deg vs 20 deg (relative)
    # ========================================================================
    print("\nGenerating 40 deg vs 20 deg...")

    sc_40_label = subcase_map['40']
    sc_40_data = displacements[sc_40_label]
    X_40, Y_40, WFE_40 = build_wfe_arrays(
        '40', sc_40_data['node_ids'], sc_40_data['disp'], node_geo
    )

    X_40_rel, Y_40_rel, WFE_40_rel = compute_relative_wfe(
        X_40, Y_40, WFE_40, sc_40_data['node_ids'],
        X_ref, Y_ref, WFE_ref, ref_data['node_ids']
    )

    rms_40_abs = compute_rms_metrics(X_40, Y_40, WFE_40)
    rms_40_rel = compute_rms_metrics(X_40_rel, Y_40_rel, WFE_40_rel)

    html_40 = generate_html(
        title="40 deg",
        X=X_40_rel, Y=Y_40_rel, W_nm=WFE_40_rel,
        rms_data=rms_40_rel,
        is_relative=True,
        ref_title="20 deg",
        abs_pair=(rms_40_abs['global_rms'], rms_40_abs['filtered_rms'])
    )

    path_40 = output_dir / f"{base}_{timestamp}_40_vs_20.html"
    path_40.write_text(html_40, encoding='utf-8')
    html_files.append(path_40)
    print(f"  Created: {path_40.name}")
    print(f"  Relative: Global={rms_40_rel['global_rms']:.2f}, Filtered={rms_40_rel['filtered_rms']:.2f}")

    # ========================================================================
    # 2. Generate 60 deg vs 20 deg (relative)
    # ========================================================================
    print("\nGenerating 60 deg vs 20 deg...")

    sc_60_label = subcase_map['60']
    sc_60_data = displacements[sc_60_label]
    X_60, Y_60, WFE_60 = build_wfe_arrays(
        '60', sc_60_data['node_ids'], sc_60_data['disp'], node_geo
    )

    X_60_rel, Y_60_rel, WFE_60_rel = compute_relative_wfe(
        X_60, Y_60, WFE_60, sc_60_data['node_ids'],
        X_ref, Y_ref, WFE_ref, ref_data['node_ids']
    )

    rms_60_abs = compute_rms_metrics(X_60, Y_60, WFE_60)
    rms_60_rel = compute_rms_metrics(X_60_rel, Y_60_rel, WFE_60_rel)

    html_60 = generate_html(
        title="60 deg",
        X=X_60_rel, Y=Y_60_rel, W_nm=WFE_60_rel,
        rms_data=rms_60_rel,
        is_relative=True,
        ref_title="20 deg",
        abs_pair=(rms_60_abs['global_rms'], rms_60_abs['filtered_rms'])
    )

    path_60 = output_dir / f"{base}_{timestamp}_60_vs_20.html"
    path_60.write_text(html_60, encoding='utf-8')
    html_files.append(path_60)
    print(f"  Created: {path_60.name}")
    print(f"  Relative: Global={rms_60_rel['global_rms']:.2f}, Filtered={rms_60_rel['filtered_rms']:.2f}")

    # ========================================================================
    # 3. Generate 90 deg Manufacturing (absolute with correction metrics)
    # ========================================================================
    print("\nGenerating 90 deg Manufacturing...")

    # Compute 90 deg - 20 deg for correction metrics
    X_90_rel, Y_90_rel, WFE_90_rel = compute_relative_wfe(
        X_90, Y_90, WFE_90, mfg_data['node_ids'],
        X_ref, Y_ref, WFE_ref, ref_data['node_ids']
    )
    rms_90_rel = compute_rms_metrics(X_90_rel, Y_90_rel, WFE_90_rel)
    # Get all correction metrics from Zernike coefficients (90° - 20°)
    correction_metrics = compute_mfg_metrics(rms_90_rel['coefficients'])

    html_90 = generate_html(
        title="90 deg (Manufacturing)",
        X=X_90, Y=Y_90, W_nm=WFE_90,
        rms_data=rms_90,
        is_relative=False,
        is_manufacturing=True,
        mfg_metrics=mfg_metrics,
        correction_metrics=correction_metrics
    )

    path_90 = output_dir / f"{base}_{timestamp}_90_mfg.html"
    path_90.write_text(html_90, encoding='utf-8')
    html_files.append(path_90)
    print(f"  Created: {path_90.name}")
    print(f"  Absolute: Global={rms_90['global_rms']:.2f}, Filtered={rms_90['filtered_rms']:.2f}")
    print(f"  Optician Workload (J1-J3): {rms_90['rms_filter_j1to3']:.2f} nm")

    # ========================================================================
    # Summary
    # ========================================================================
    print("\n" + "=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print(f"\nGenerated {len(html_files)} HTML files:")
    for f in html_files:
        print(f"  - {f.name}")

    print("\n" + "-" * 70)
    print("OPTIMIZATION OBJECTIVES")
    print("-" * 70)
    print(f"  40-20 Filtered RMS:  {rms_40_rel['filtered_rms']:.2f} nm")
    print(f"  60-20 Filtered RMS:  {rms_60_rel['filtered_rms']:.2f} nm")
    print(f"  MFG 90 (J1-J3):      {rms_90_rel['rms_filter_j1to3']:.2f} nm")

    # Weighted sums
    ws_v4 = 5*rms_40_rel['filtered_rms'] + 5*rms_60_rel['filtered_rms'] + 2*rms_90_rel['rms_filter_j1to3']
    ws_v5 = 5*rms_40_rel['filtered_rms'] + 5*rms_60_rel['filtered_rms'] + 3*rms_90_rel['rms_filter_j1to3']
    print(f"\n  V4 Weighted Sum (5/5/2): {ws_v4:.2f}")
    print(f"  V5 Weighted Sum (5/5/3): {ws_v5:.2f}")

    print("\n" + "=" * 70)
    print("DONE")
    print("=" * 70)

    return html_files


if __name__ == '__main__':
    if len(sys.argv) > 1:
        op2_path = Path(sys.argv[1])
        if not op2_path.exists():
            print(f"ERROR: File not found: {op2_path}")
            sys.exit(1)
    else:
        print("No OP2 file specified, searching...")
        op2_path = find_op2_file()
        if op2_path is None:
            print("ERROR: No OP2 file found in current directory.")
            print("Usage: python zernike_html_generator.py <path/to/solution.op2>")
            sys.exit(1)
        print(f"Found: {op2_path}")

    # Check for --legacy flag to use old Z-only method
    use_legacy = '--legacy' in sys.argv or '--z-only' in sys.argv

    try:
        if USE_OPD_METHOD and not use_legacy:
            main_opd(op2_path)
        else:
            if use_legacy:
                print("[INFO] Using legacy Z-only method (--legacy flag)")
            main(op2_path)
    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)