#!/usr/bin/env python
"""
Compare all Zernike extraction methods on M1 Mirror data.

Methods compared:
1. Standard - Z-displacement only at original (x,y)
2. Parabola OPD - Uses parabola approximation with prescription focal length
3. Figure OPD - Uses actual figure.dat geometry (most rigorous)

Run:
    python tests/test_zernike_methods_comparison.py studies/M1_Mirror/m1_mirror_cost_reduction_V9
"""

from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).parent.parent))

import numpy as np


def run_full_comparison(study_dir: Path):
    """Run comparison of all three Zernike methods."""
    from optimization_engine.extractors.extract_zernike import ZernikeExtractor
    from optimization_engine.extractors.extract_zernike_opd import ZernikeAnalyticExtractor
    from optimization_engine.extractors.extract_zernike_figure import ZernikeOPDExtractor

    # Find OP2 file
    op2_files = list(study_dir.glob('3_results/best_design_archive/**/*.op2'))
    if not op2_files:
        op2_files = list(study_dir.glob('2_iterations/iter1/*.op2'))
    if not op2_files:
        raise FileNotFoundError(f"No OP2 file found in {study_dir}")

    op2_file = op2_files[0]

    # Figure file is optional - extractor will use BDF geometry filtered to OP2 nodes
    figure_file = study_dir / '1_setup' / 'model' / 'figure.dat'
    use_figure_file = figure_file.exists()

    print("=" * 80)
    print("ZERNIKE METHODS COMPARISON - M1 MIRROR")
    print("=" * 80)
    print(f"\nOP2 file: {op2_file.name}")
    if use_figure_file:
        print(f"Figure file: {figure_file.name}")
    else:
        print("Figure file: Not found (using BDF geometry filtered to OP2 nodes)")

    # Initialize extractors
    print("\nInitializing extractors...")

    std_extractor = ZernikeExtractor(op2_file)
    print(f"  Standard: {len(std_extractor.node_geometry)} nodes in BDF")

    analytic_extractor = ZernikeAnalyticExtractor(op2_file, focal_length=1445.0, concave=True)
    print(f"  Analytic (Parabola f=1445mm): {len(analytic_extractor.node_geometry)} nodes")

    # OPD extractor - ALWAYS use BDF geometry filtered to OP2 nodes (RECOMMENDED)
    # (figure.dat may have mismatched coordinates from different model state)
    opd_extractor = ZernikeOPDExtractor(
        op2_file,
        figure_path=None  # Force use of BDF geometry
    )
    print(f"  OPD (BDF geometry): {len(opd_extractor.figure_geometry)} figure nodes")

    print()

    # Get available subcases
    subcases = list(std_extractor.displacements.keys())
    print(f"Available subcases: {subcases}")

    # Results table header
    print("\n" + "=" * 80)
    print(f"{'Subcase':<10} {'Standard':<15} {'Analytic':<15} {'OPD':<15} {'Max Lat (um)':<12}")
    print("-" * 80)

    all_results = {}

    for sc in subcases:
        try:
            # Extract with each method
            std_res = std_extractor.extract_subcase(sc)
            analytic_res = analytic_extractor.extract_subcase(sc)
            opd_res = opd_extractor.extract_subcase(sc)

            std_rms = std_res['filtered_rms_nm']
            analytic_rms = analytic_res['filtered_rms_nm']
            opd_rms = opd_res['filtered_rms_nm']
            max_lat = opd_res.get('max_lateral_displacement_um', 0)

            print(f"{sc:<10} {std_rms:<15.2f} {analytic_rms:<15.2f} {opd_rms:<15.2f} {max_lat:<12.3f}")

            all_results[sc] = {
                'standard': std_res,
                'analytic': analytic_res,
                'opd': opd_res
            }

        except Exception as e:
            print(f"{sc:<10} ERROR: {e}")

    print("-" * 80)

    # Detailed comparison for each subcase
    print("\n" + "=" * 80)
    print("DETAILED COMPARISON")
    print("=" * 80)

    for sc, results in all_results.items():
        print(f"\n{'-' * 40}")
        print(f"SUBCASE {sc}")
        print(f"{'-' * 40}")

        std = results['standard']
        analytic = results['analytic']
        opd = results['opd']

        print(f"\n{'Metric':<25} {'Standard':<15} {'Analytic':<15} {'OPD':<15}")
        print("-" * 70)

        # RMS metrics
        print(f"{'Filtered RMS (nm)':<25} {std['filtered_rms_nm']:<15.2f} {analytic['filtered_rms_nm']:<15.2f} {opd['filtered_rms_nm']:<15.2f}")
        print(f"{'Global RMS (nm)':<25} {std['global_rms_nm']:<15.2f} {analytic['global_rms_nm']:<15.2f} {opd['global_rms_nm']:<15.2f}")

        # Aberrations
        for aberr in ['defocus', 'astigmatism', 'coma', 'trefoil', 'spherical']:
            key = f'{aberr}_nm'
            if key in std and key in analytic and key in opd:
                print(f"{aberr.capitalize():<25} {std[key]:<15.2f} {analytic[key]:<15.2f} {opd[key]:<15.2f}")

        # Node count
        print(f"{'Nodes':<25} {std['n_nodes']:<15} {analytic['n_nodes']:<15} {opd['n_nodes']:<15}")

        # Lateral displacement (only for OPD methods)
        print()
        print(f"Lateral displacement (OPD method):")
        print(f"  Max:  {opd.get('max_lateral_displacement_um', 0):.3f} um")
        print(f"  RMS:  {opd.get('rms_lateral_displacement_um', 0):.3f} um")
        print(f"  Mean: {opd.get('mean_lateral_displacement_um', 0):.3f} um")

        # Differences
        print()
        diff_std_opd = opd['filtered_rms_nm'] - std['filtered_rms_nm']
        diff_analytic_opd = opd['filtered_rms_nm'] - analytic['filtered_rms_nm']

        print(f"Difference from Standard:")
        print(f"  OPD:      {diff_std_opd:+.2f} nm ({100*diff_std_opd/std['filtered_rms_nm']:+.1f}%)")
        print(f"  Analytic: {analytic['filtered_rms_nm'] - std['filtered_rms_nm']:+.2f} nm")
        print()
        print(f"Difference OPD vs Analytic: {diff_analytic_opd:+.2f} nm")

    # Tracking WFE analysis
    if '2' in all_results and '3' in all_results and '4' in all_results:
        print("\n" + "=" * 80)
        print("TRACKING WFE ANALYSIS (elevation changes)")
        print("=" * 80)

        for method_name, method_key in [('Standard', 'standard'), ('Analytic', 'analytic'), ('OPD', 'opd')]:
            print(f"\n{method_name}:")

            z20 = np.array(all_results['2'][method_key].get('coefficients', [0]*36))
            z40 = np.array(all_results['3'][method_key].get('coefficients', [0]*36))
            z60 = np.array(all_results['4'][method_key].get('coefficients', [0]*36))

            if len(z20) > 4:
                # Differential (J5+)
                diff_40_20 = z40 - z20
                diff_60_20 = z60 - z20

                rms_40_20 = np.sqrt(np.sum(diff_40_20[4:]**2))
                rms_60_20 = np.sqrt(np.sum(diff_60_20[4:]**2))

                print(f"  40-20 deg tracking: {rms_40_20:.2f} nm RMS (J5+ filtered)")
                print(f"  60-20 deg tracking: {rms_60_20:.2f} nm RMS (J5+ filtered)")
            else:
                print(f"  (coefficients not available)")

    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)
    print("""
Key findings:
- Standard method: Uses Z-displacement only at original (x,y) - fast but ignores lateral shift
- Analytic method: Accounts for lateral shift using parabola formula (requires focal length)
- OPD method: Uses actual mesh geometry - MOST RIGOROUS, no shape assumption

Recommendation: Use OPD method (ZernikeOPDExtractor) for all mirror optimization.
The Analytic method is useful for comparison against theoretical parabola.
""")


def main():
    import argparse
    parser = argparse.ArgumentParser(description='Compare Zernike extraction methods')
    parser.add_argument('study_dir', type=str, help='Path to study directory')
    args = parser.parse_args()

    study_dir = Path(args.study_dir).resolve()
    if not study_dir.exists():
        print(f"Study directory not found: {study_dir}")
        sys.exit(1)

    run_full_comparison(study_dir)


if __name__ == '__main__':
    main()