#!/usr/bin/env python3
"""
Atomizer Zernike WFE Analyzer
=============================

Analyze Zernike wavefront error from NX Nastran OP2 results.

IMPORTANT: This script requires numpy/scipy. Run from command line with
the atomizer conda environment, NOT from within NX.

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

  # Or without argument - searches current directory for OP2 files:
  python analyze_wfe_zernike.py

Output:
  - Zernike coefficients for each subcase
  - Relative WFE metrics (filtered RMS)
  - Manufacturing workload (J1-J3 filtered)
  - Weighted sum calculation

Author: Atomizer
Created: 2025-12-18
"""

import sys
import os
from pathlib import Path


def log(msg):
    """Print to console."""
    print(msg)


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)

    # Look for OP2 files
    op2_files = list(working_dir.glob("*solution*.op2")) + list(working_dir.glob("*.op2"))

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

    if not op2_files:
        return None

    # Return most recently modified
    return max(op2_files, key=lambda p: p.stat().st_mtime)

def analyze_zernike(op2_path):
    """Run Zernike analysis on OP2 file."""

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

    try:
        from optimization_engine.extractors import ZernikeExtractor
    except ImportError as e:
        log(f"ERROR: Could not import ZernikeExtractor: {e}")
        log(f"Make sure Atomizer is properly installed.")
        log(f"Atomizer root: {atomizer_root}")
        return None

    log("=" * 70)
    log("ZERNIKE WAVEFRONT ERROR ANALYSIS")
    log("=" * 70)
    log(f"OP2 File: {op2_path.name}")
    log(f"Directory: {op2_path.parent}")
    log("")

    # Create extractor
    try:
        extractor = ZernikeExtractor(
            op2_path,
            bdf_path=None,
            displacement_unit='mm',
            n_modes=50,
            filter_orders=4
        )
    except Exception as e:
        log(f"ERROR creating extractor: {e}")
        return None

    # Get available subcases from the extractor's displacement data
    subcases = list(extractor.displacements.keys())
    log(f"Available subcases: {subcases}")
    log("")

    # Standard subcase mapping for M1 mirror
    subcase_labels = {
        '1': '90 deg (Manufacturing/Polishing)',
        '2': '20 deg (Reference)',
        '3': '40 deg (Operational)',
        '4': '60 deg (Operational)'
    }

    # Extract absolute Zernike for each subcase
    log("-" * 70)
    log("ABSOLUTE ZERNIKE ANALYSIS (per subcase)")
    log("-" * 70)

    results = {}
    for sc in subcases:
        try:
            result = extractor.extract_subcase(sc)
            results[sc] = result
            label = subcase_labels.get(sc, f'Subcase {sc}')
            log(f"\n{label}:")
            log(f"  Global RMS:     {result['global_rms_nm']:.2f} nm")
            log(f"  Filtered RMS:   {result['filtered_rms_nm']:.2f} nm  (J4+ only)")
        except Exception as e:
            log(f"  ERROR extracting subcase {sc}: {e}")

    # Relative analysis (using subcase 2 as reference)
    ref_subcase = '2'
    if ref_subcase in subcases:
        log("")
        log("-" * 70)
        log(f"RELATIVE ANALYSIS (vs {subcase_labels.get(ref_subcase, ref_subcase)})")
        log("-" * 70)

        relative_results = {}
        for sc in subcases:
            if sc == ref_subcase:
                continue
            try:
                rel = extractor.extract_relative(sc, ref_subcase)
                relative_results[sc] = rel
                label = subcase_labels.get(sc, f'Subcase {sc}')
                log(f"\n{label} vs Reference:")
                log(f"  Relative Filtered RMS: {rel['relative_filtered_rms_nm']:.2f} nm")
                if 'relative_rms_filter_j1to3' in rel:
                    log(f"  J1-J3 Filtered RMS:    {rel['relative_rms_filter_j1to3']:.2f} nm")
            except Exception as e:
                log(f"  ERROR: {e}")

        # Calculate weighted sum (M1 mirror optimization objectives)
        log("")
        log("-" * 70)
        log("OPTIMIZATION OBJECTIVES")
        log("-" * 70)

        obj_40_20 = relative_results.get('3', {}).get('relative_filtered_rms_nm', 0)
        obj_60_20 = relative_results.get('4', {}).get('relative_filtered_rms_nm', 0)
        obj_mfg = relative_results.get('1', {}).get('relative_rms_filter_j1to3', 0)

        log(f"\n  40-20 Filtered RMS:  {obj_40_20:.2f} nm")
        log(f"  60-20 Filtered RMS:  {obj_60_20:.2f} nm")
        log(f"  MFG 90 (J1-J3):      {obj_mfg:.2f} nm")

        # Weighted sums for different weight configurations
        log("")
        log("Weighted Sum Calculations:")

        # V4 weights: 5*40 + 5*60 + 2*mfg + mass
        ws_v4 = 5*obj_40_20 + 5*obj_60_20 + 2*obj_mfg
        log(f"  V4 weights (5/5/2): {ws_v4:.2f}  (+ mass)")

        # V5 weights: 5*40 + 5*60 + 3*mfg + mass
        ws_v5 = 5*obj_40_20 + 5*obj_60_20 + 3*obj_mfg
        log(f"  V5 weights (5/5/3): {ws_v5:.2f}  (+ mass)")

        return {
            'absolute': results,
            'relative': relative_results,
            'objectives': {
                '40_20': obj_40_20,
                '60_20': obj_60_20,
                'mfg_90': obj_mfg,
                'ws_v4': ws_v4,
                'ws_v5': ws_v5
            }
        }

    return {'absolute': results}

def main(args):
    """Main entry point."""
    log("")
    log("=" * 70)
    log("  ATOMIZER ZERNIKE WFE ANALYZER")
    log("=" * 70)
    log("")

    # Determine OP2 file
    op2_path = None

    if args and len(args) > 0 and args[0]:
        # OP2 path provided as argument
        op2_path = Path(args[0])
        if not op2_path.exists():
            log(f"ERROR: OP2 file not found: {op2_path}")
            return
    else:
        # Try to find OP2 in current directory
        log("No OP2 file specified, searching...")
        op2_path = find_op2_file()

        if op2_path is None:
            log("ERROR: No OP2 file found in current directory.")
            log("Usage: Run after solving, or provide OP2 path as argument.")
            return

        log(f"Found: {op2_path}")

    # Run analysis
    results = analyze_zernike(op2_path)

    if results:
        log("")
        log("=" * 70)
        log("ANALYSIS COMPLETE")
        log("=" * 70)
    else:
        log("")
        log("Analysis failed. Check errors above.")

if __name__ == '__main__':
    # Get arguments (works both in NX and command line)
    if len(sys.argv) > 1:
        main(sys.argv[1:])
    else:
        main([])