#!/usr/bin/env python
"""
Test ZernikeOPDExtractor with validated M1 Mirror optical prescription.

Compares:
1. Standard Zernike (Z-displacement only at original x,y)
2. OPD Zernike with auto-estimated focal length
3. OPD Zernike with correct focal length (1445 mm from prescription)

M1 Mirror Optical Prescription:
- Radius of Curvature: 2890 ± 3 mm
- Conic Constant: -0.987 ± 0.001 (near-parabolic)
- Clear Aperture: 1202 mm
- Central Bore: 271.56 mm
- Focal Length: 1445 mm (R/2)
"""

from pathlib import Path
import sys

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

import numpy as np


def run_comparison(op2_path: Path):
    """Run comparison between standard and OPD Zernike methods."""
    from optimization_engine.extractors.extract_zernike_opd import ZernikeOPDExtractor
    from optimization_engine.extractors.extract_zernike_wfe import ZernikeExtractor

    print("=" * 70)
    print("ZERNIKE METHOD COMPARISON WITH VALIDATED PRESCRIPTION")
    print("=" * 70)
    print(f"\nOP2 file: {op2_path.name}")
    print(f"Optical prescription focal length: 1445 mm")
    print()

    # 1. Standard Zernike (Z-displacement only)
    print("1. STANDARD ZERNIKE (Z-displacement at original x,y)")
    print("-" * 50)
    try:
        std_extractor = ZernikeExtractor(op2_path)
        std_results = std_extractor.extract_all_subcases()

        for sc, data in std_results.items():
            coeffs = data['coefficients']
            rms = data['rms_wfe_nm']
            print(f"   Subcase {sc}: RMS WFE = {rms:.2f} nm")
    except Exception as e:
        print(f"   Error: {e}")
        std_results = None
    print()

    # 2. OPD Zernike with auto-estimated focal length
    print("2. OPD ZERNIKE (auto-estimated focal length)")
    print("-" * 50)
    try:
        opd_auto = ZernikeOPDExtractor(op2_path, concave=True)
        auto_focal = opd_auto.focal_length
        print(f"   Auto-estimated focal length: {auto_focal:.1f} mm")

        opd_auto_results = opd_auto.extract_all_subcases()
        for sc, data in opd_auto_results.items():
            rms = data['rms_wfe_nm']
            lat = data.get('max_lateral_displacement_um', 0)
            print(f"   Subcase {sc}: RMS WFE = {rms:.2f} nm, Max lateral = {lat:.2f} µm")
    except Exception as e:
        print(f"   Error: {e}")
        opd_auto_results = None
    print()

    # 3. OPD Zernike with correct prescription focal length
    print("3. OPD ZERNIKE (prescription focal length = 1445 mm)")
    print("-" * 50)
    try:
        opd_correct = ZernikeOPDExtractor(op2_path, focal_length=1445.0, concave=True)
        print(f"   Using focal length: {opd_correct.focal_length:.1f} mm")

        opd_correct_results = opd_correct.extract_all_subcases()
        for sc, data in opd_correct_results.items():
            rms = data['rms_wfe_nm']
            lat = data.get('max_lateral_displacement_um', 0)
            print(f"   Subcase {sc}: RMS WFE = {rms:.2f} nm, Max lateral = {lat:.2f} µm")
    except Exception as e:
        print(f"   Error: {e}")
        opd_correct_results = None
    print()

    # 4. Comparison summary
    if std_results and opd_correct_results:
        print("=" * 70)
        print("COMPARISON SUMMARY")
        print("=" * 70)
        print()
        print(f"{'Subcase':<10} {'Standard':<15} {'OPD (auto)':<15} {'OPD (1445mm)':<15} {'Diff %':<10}")
        print("-" * 65)

        for sc in std_results.keys():
            std_rms = std_results[sc]['rms_wfe_nm']
            auto_rms = opd_auto_results[sc]['rms_wfe_nm'] if opd_auto_results else 0
            corr_rms = opd_correct_results[sc]['rms_wfe_nm']
            diff_pct = ((corr_rms - std_rms) / std_rms * 100) if std_rms > 0 else 0

            print(f"{sc:<10} {std_rms:<15.2f} {auto_rms:<15.2f} {corr_rms:<15.2f} {diff_pct:>+8.1f}%")

        print()
        print("LATERAL DISPLACEMENT ANALYSIS")
        print("-" * 50)

        for sc, data in opd_correct_results.items():
            lat = data.get('max_lateral_displacement_um', 0)
            severity = "CRITICAL - OPD method required" if lat > 10 else "Low - standard OK" if lat < 1 else "Moderate"
            print(f"   Subcase {sc}: Max lateral = {lat:.2f} µm ({severity})")

        print()

        # Tracking WFE comparison (40-20 and 60-20)
        if 2 in opd_correct_results and 3 in opd_correct_results and 4 in opd_correct_results:
            print("TRACKING WFE (differential between elevations)")
            print("-" * 50)

            # Get coefficients for differential analysis
            z20 = np.array(opd_correct_results[2]['coefficients'])
            z40 = np.array(opd_correct_results[3]['coefficients'])
            z60 = np.array(opd_correct_results[4]['coefficients'])

            # Differential (remove J1-J4: piston, tip, tilt, defocus)
            diff_40_20 = z40 - z20
            diff_60_20 = z60 - z20

            # RMS of filtered differential (J5+)
            rms_40_20 = np.sqrt(np.sum(diff_40_20[4:]**2))  # Skip J1-J4
            rms_60_20 = np.sqrt(np.sum(diff_60_20[4:]**2))

            print(f"   40°-20° tracking WFE: {rms_40_20:.2f} nm RMS (filtered)")
            print(f"   60°-20° tracking WFE: {rms_60_20:.2f} nm RMS (filtered)")

            print()
            print("   Standard method comparison:")
            z20_std = np.array(std_results[2]['coefficients'])
            z40_std = np.array(std_results[3]['coefficients'])
            z60_std = np.array(std_results[4]['coefficients'])

            diff_40_20_std = z40_std - z20_std
            diff_60_20_std = z60_std - z20_std

            rms_40_20_std = np.sqrt(np.sum(diff_40_20_std[4:]**2))
            rms_60_20_std = np.sqrt(np.sum(diff_60_20_std[4:]**2))

            print(f"   40°-20° tracking WFE (std): {rms_40_20_std:.2f} nm RMS")
            print(f"   60°-20° tracking WFE (std): {rms_60_20_std:.2f} nm RMS")

            print()
            print(f"   Difference (OPD vs Standard):")
            print(f"   40°-20°: {rms_40_20 - rms_40_20_std:+.2f} nm ({(rms_40_20/rms_40_20_std - 1)*100:+.1f}%)")
            print(f"   60°-20°: {rms_60_20 - rms_60_20_std:+.2f} nm ({(rms_60_20/rms_60_20_std - 1)*100:+.1f}%)")

    print()
    print("=" * 70)


def main():
    import argparse

    parser = argparse.ArgumentParser(description='Test ZernikeOPD with M1 prescription')
    parser.add_argument('path', nargs='?', default='.',
                       help='Path to OP2 file or study directory')

    args = parser.parse_args()
    path = Path(args.path).resolve()

    # Find OP2 file
    if path.is_file() and path.suffix.lower() == '.op2':
        op2_path = path
    elif path.is_dir():
        # Look for best design or recent iteration
        patterns = [
            '3_results/best_design_archive/**/*.op2',
            '2_iterations/iter1/*.op2',
            '**/*.op2'
        ]
        for pattern in patterns:
            files = list(path.glob(pattern))
            if files:
                op2_path = max(files, key=lambda p: p.stat().st_mtime)
                break
        else:
            print(f"No OP2 file found in {path}")
            sys.exit(1)
    else:
        print(f"Invalid path: {path}")
        sys.exit(1)

    run_comparison(op2_path)


if __name__ == '__main__':
    main()