#!/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
- Supports ANNULAR APERTURES via --inner-radius option

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

  # With annular aperture (central hole)
  python zernike_html_generator.py "path/to/solution.op2" --inner-radius 135.75

  # 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
Updated: 2025-01-06 - Added annular aperture support via --inner-radius
"""

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) - MFG_90 objective (relative 90-20, J1-J3 filtered)
        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=[
                ["MFG_90 Objective (90-20, J1-J3 filtered)",
                 "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_rel,  # Use RELATIVE (90-20) for mfg_90 objective match
        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_rel,  # Use RELATIVE (90-20) for mfg_90 objective match
        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)