"""
Debug script to investigate the lateral displacement discrepancy between
Zernike OPD output and Simcenter post-processing.

User observation:
- Zernike OPD shows max lateral XY displacement: ~0.2238 µm
- Simcenter shows XX displacement at 20deg: 7.457e-05 mm = 0.0746 µm

Hypothesis: The Zernike "max_lateral_disp_um" is sqrt(dx² + dy²), not just dx or dy.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np
from pathlib import Path

# Find a mirror study OP2 to analyze
STUDIES_PATH = Path(r"c:\Users\antoi\Atomizer\studies\M1_Mirror")

def find_latest_op2():
    """Find a recent OP2 file to analyze."""
    op2_files = list(STUDIES_PATH.rglob("*.op2"))
    if not op2_files:
        print("No OP2 files found!")
        return None
    # Get most recent
    return max(op2_files, key=lambda p: p.stat().st_mtime)

def analyze_displacements(op2_path: Path):
    """Analyze displacement components for all subcases."""
    from pyNastran.op2.op2 import OP2
    from pyNastran.bdf.bdf import BDF

    print(f"\n{'='*70}")
    print(f"Analyzing: {op2_path.name}")
    print(f"Path: {op2_path}")
    print(f"{'='*70}")

    # Find BDF
    bdf_path = None
    for ext in ['.dat', '.bdf']:
        candidate = op2_path.with_suffix(ext)
        if candidate.exists():
            bdf_path = candidate
            break

    if not bdf_path:
        print("No BDF found, searching parent...")
        for f in op2_path.parent.iterdir():
            if f.suffix.lower() in ['.dat', '.bdf']:
                bdf_path = f
                break

    if not bdf_path:
        print("ERROR: No geometry file found!")
        return

    print(f"BDF: {bdf_path.name}")

    # Read data
    print("\nLoading OP2...")
    op2 = OP2()
    op2.read_op2(str(op2_path))

    print("Loading BDF...")
    bdf = BDF()
    bdf.read_bdf(str(bdf_path))

    node_geo = {int(nid): node.get_position() for nid, node in bdf.nodes.items()}

    SUBCASE_MAP = {'1': 90, '2': 20, '3': 40, '4': 60}

    print(f"\nNode count in BDF: {len(node_geo)}")
    print(f"\n{'='*70}")
    print("DISPLACEMENT ANALYSIS BY SUBCASE")
    print(f"{'='*70}")

    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]

        isubcase = getattr(darr, 'isubcase', None)
        label = str(isubcase) if isubcase else str(key)
        angle = SUBCASE_MAP.get(label, label)

        print(f"\n--- Subcase {label} ({angle} deg) ---")
        print(f"Nodes with displacement: {len(node_ids)}")

        # Extract displacement components
        dx = dmat[:, 0]  # X displacement (mm)
        dy = dmat[:, 1]  # Y displacement (mm)
        dz = dmat[:, 2]  # Z displacement (mm)

        # Compute statistics
        dx_um = dx * 1000.0  # Convert mm to µm
        dy_um = dy * 1000.0
        dz_um = dz * 1000.0

        # Lateral magnitude (XY combined)
        lateral_um = np.sqrt(dx_um**2 + dy_um**2)

        print(f"\nComponent Statistics (in um):")
        print(f"  X displacement (dx):")
        print(f"    Min:    {np.min(dx_um):+.4f} um")
        print(f"    Max:    {np.max(dx_um):+.4f} um")
        print(f"    RMS:    {np.sqrt(np.mean(dx_um**2)):.4f} um")

        print(f"  Y displacement (dy):")
        print(f"    Min:    {np.min(dy_um):+.4f} um")
        print(f"    Max:    {np.max(dy_um):+.4f} um")
        print(f"    RMS:    {np.sqrt(np.mean(dy_um**2)):.4f} um")

        print(f"  Z displacement (dz):")
        print(f"    Min:    {np.min(dz_um):+.4f} um")
        print(f"    Max:    {np.max(dz_um):+.4f} um")
        print(f"    RMS:    {np.sqrt(np.mean(dz_um**2)):.4f} um")

        print(f"\n  Lateral Magnitude (sqrt(dx^2 + dy^2)):")
        print(f"    Max:    {np.max(lateral_um):.4f} um  <-- This is what Zernike OPD reports!")
        print(f"    RMS:    {np.sqrt(np.mean(lateral_um**2)):.4f} um")

        # Find the node with max lateral displacement
        max_lat_idx = np.argmax(lateral_um)
        max_lat_nid = int(node_ids[max_lat_idx])
        max_dx = dx_um[max_lat_idx]
        max_dy = dy_um[max_lat_idx]

        print(f"\n  Node with max lateral displacement: Node {max_lat_nid}")
        print(f"    dx = {max_dx:+.4f} um")
        print(f"    dy = {max_dy:+.4f} um")
        print(f"    sqrt(dx^2+dy^2) = {lateral_um[max_lat_idx]:.4f} um")

        # Compare with just max(|dx|)
        max_abs_dx = np.max(np.abs(dx_um))
        max_abs_dy = np.max(np.abs(dy_um))
        print(f"\n  For comparison (what you see in Simcenter):")
        print(f"    max(|dx|) = {max_abs_dx:.4f} um")
        print(f"    max(|dy|) = {max_abs_dy:.4f} um")

    print(f"\n{'='*70}")
    print("EXPLANATION OF DISCREPANCY")
    print(f"{'='*70}")
    print("""
The Zernike OPD insight reports "max_lateral_disp_um" as:
    lateral = sqrt(dx^2 + dy^2)  -- combined XY magnitude at each node

Simcenter's "Displacement - Nodal, X" shows:
    Just the X component (dx) at each node

These are different metrics:
- If dx_max = 0.0746 um and dy is significant, then:
  lateral = sqrt(0.0746^2 + dy^2) > 0.0746 um

To match Simcenter exactly, look at the individual dx/dy/dz stats above.
The "max_lateral_um" in Zernike OPD is the MAGNITUDE of the XY vector,
not the individual X or Y components.

For a node where both dx and dy are non-zero:
  dx = 0.0746 um, dy = 0.21 um
  lateral = sqrt(0.0746^2 + 0.21^2) = sqrt(0.0056 + 0.044) = sqrt(0.0496) = 0.22 um
""")

if __name__ == '__main__':
    # Find and analyze an OP2 file
    if len(sys.argv) > 1:
        op2_path = Path(sys.argv[1])
    else:
        op2_path = find_latest_op2()

    if op2_path and op2_path.exists():
        analyze_displacements(op2_path)
    else:
        print("Please provide an OP2 file path as argument, or place OP2 files in the studies directory.")
        print(f"\nUsage: python {sys.argv[0]} <path_to_op2_file>")
