"""
Compare V8 Best Candidate: OPD Method vs Standard Z-Only Method

This script extracts WFE using both methods and compares the results
to quantify the difference between using full X,Y,Z displacement (OPD)
vs just Z displacement (Standard).
"""
import sys
sys.path.insert(0, '.')

import sqlite3
from pathlib import Path
import json
import numpy as np

# Find V8 best trial
db_path = Path('studies/M1_Mirror/m1_mirror_cost_reduction_V8/3_results/study.db')
print(f"V8 Database: {db_path}")
print(f"Exists: {db_path.exists()}")

if not db_path.exists():
    print("ERROR: V8 database not found!")
    sys.exit(1)

conn = sqlite3.connect(db_path)
c = conn.cursor()

# Get all completed trials with their objectives
c.execute('''
    SELECT t.trial_id
    FROM trials t
    WHERE t.state = 'COMPLETE'
    ORDER BY t.trial_id
''')
completed_ids = [row[0] for row in c.fetchall()]
print(f"Completed trials: {len(completed_ids)}")

# Build trial data and find best
trials = []
for tid in completed_ids:
    c.execute('''
        SELECT key, value_json
        FROM trial_user_attributes
        WHERE trial_id = ?
    ''', (tid,))
    attrs = {row[0]: json.loads(row[1]) for row in c.fetchall()}

    # V8 used different objective names - check what's available
    rms_40 = attrs.get('rel_filtered_rms_40_vs_20', attrs.get('filtered_rms_40_20', None))
    rms_60 = attrs.get('rel_filtered_rms_60_vs_20', attrs.get('filtered_rms_60_20', None))
    mfg_90 = attrs.get('mfg_90_optician_workload', attrs.get('optician_workload_90', None))
    ws = attrs.get('weighted_sum', None)

    if rms_40 is not None:
        trials.append({
            'id': tid,
            'rms_40': rms_40,
            'rms_60': rms_60 if rms_60 else 999,
            'mfg_90': mfg_90 if mfg_90 else 999,
            'ws': ws if ws else 999,
        })

conn.close()

if not trials:
    # Check what keys are available
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('SELECT DISTINCT key FROM trial_user_attributes LIMIT 20')
    keys = [row[0] for row in c.fetchall()]
    print(f"\nAvailable attribute keys: {keys}")
    conn.close()
    print("ERROR: No trials found with expected objective names!")
    sys.exit(1)

# Calculate weighted sum if not present
for t in trials:
    if t['ws'] == 999:
        w40 = 100 / 4.0
        w60 = 50 / 10.0
        w90 = 20 / 20.0
        t['ws'] = w40 * t['rms_40'] + w60 * t['rms_60'] + w90 * t['mfg_90']

# Find best
trials.sort(key=lambda x: x['ws'])
best = trials[0]

print(f"\n{'='*70}")
print("V8 Best Trial Summary")
print('='*70)
print(f"Best Trial: #{best['id']}")
print(f"  40-20 RMS: {best['rms_40']:.2f} nm")
print(f"  60-20 RMS: {best['rms_60']:.2f} nm")
print(f"  MFG 90:    {best['mfg_90']:.2f} nm")
print(f"  WS:        {best['ws']:.1f}")

# Now compare both extraction methods on this trial
iter_path = Path(f'studies/M1_Mirror/m1_mirror_cost_reduction_V8/2_iterations/iter{best["id"]}')
op2_path = iter_path / 'assy_m1_assyfem1_sim1-solution_1.op2'
geo_path = iter_path / 'assy_m1_assyfem1_sim1-solution_1.dat'

print(f"\nIteration path: {iter_path}")
print(f"OP2 exists: {op2_path.exists()}")
print(f"Geometry exists: {geo_path.exists()}")

if not op2_path.exists():
    print("ERROR: OP2 file not found for best trial!")
    sys.exit(1)

print(f"\n{'='*70}")
print("Comparing Zernike Methods: OPD vs Standard")
print('='*70)

# Import extractors
from optimization_engine.extractors.extract_zernike_figure import ZernikeOPDExtractor
from optimization_engine.extractors.extract_zernike import ZernikeExtractor

# Standard method (Z-only)
print("\n1. STANDARD METHOD (Z-only displacement)")
print("-" * 50)
try:
    std_extractor = ZernikeExtractor(
        op2_path,
        bdf_path=geo_path,
        n_modes=50,
        filter_orders=4
    )

    # Extract relative WFE
    std_40_20 = std_extractor.extract_relative(target_subcase='3', reference_subcase='2')
    std_60_20 = std_extractor.extract_relative(target_subcase='4', reference_subcase='2')

    # MFG uses J1-J3 filter - need new extractor instance
    std_extractor_mfg = ZernikeExtractor(
        op2_path,
        bdf_path=geo_path,
        n_modes=50,
        filter_orders=3  # J1-J3 for manufacturing
    )
    std_90 = std_extractor_mfg.extract_subcase(subcase_label='1')

    print(f"  40-20 Relative RMS: {std_40_20['relative_filtered_rms_nm']:.2f} nm")
    print(f"  60-20 Relative RMS: {std_60_20['relative_filtered_rms_nm']:.2f} nm")
    print(f"  90 MFG (J1-J3):     {std_90['filtered_rms_nm']:.2f} nm")

    std_results = {
        '40_20': std_40_20['relative_filtered_rms_nm'],
        '60_20': std_60_20['relative_filtered_rms_nm'],
        '90_mfg': std_90['filtered_rms_nm'],
    }
except Exception as e:
    print(f"  ERROR: {e}")
    import traceback
    traceback.print_exc()
    std_results = None

# OPD method (X,Y,Z displacement with mesh interpolation)
print("\n2. OPD METHOD (X,Y,Z displacement with mesh interpolation)")
print("-" * 50)
try:
    opd_extractor = ZernikeOPDExtractor(
        op2_path,
        bdf_path=geo_path,
        n_modes=50,
        filter_orders=4
    )

    # Extract relative WFE using OPD method
    opd_40_20 = opd_extractor.extract_relative(target_subcase='3', reference_subcase='2')
    opd_60_20 = opd_extractor.extract_relative(target_subcase='4', reference_subcase='2')

    # MFG uses J1-J3 filter
    opd_extractor_mfg = ZernikeOPDExtractor(
        op2_path,
        bdf_path=geo_path,
        n_modes=50,
        filter_orders=3  # J1-J3 for manufacturing
    )
    opd_90 = opd_extractor_mfg.extract_subcase(subcase_label='1')

    print(f"  40-20 Relative RMS: {opd_40_20['relative_filtered_rms_nm']:.2f} nm")
    print(f"  60-20 Relative RMS: {opd_60_20['relative_filtered_rms_nm']:.2f} nm")
    print(f"  90 MFG (J1-J3):     {opd_90['filtered_rms_nm']:.2f} nm")

    # Also get lateral displacement info
    print(f"\n  Lateral Displacement (40° vs 20°):")
    print(f"    Max: {opd_40_20.get('max_lateral_displacement_um', 'N/A')} µm")
    print(f"    RMS: {opd_40_20.get('rms_lateral_displacement_um', 'N/A')} µm")

    opd_results = {
        '40_20': opd_40_20['relative_filtered_rms_nm'],
        '60_20': opd_60_20['relative_filtered_rms_nm'],
        '90_mfg': opd_90['filtered_rms_nm'],
        'lateral_max': opd_40_20.get('max_lateral_displacement_um', 0),
        'lateral_rms': opd_40_20.get('rms_lateral_displacement_um', 0),
    }
except Exception as e:
    print(f"  ERROR: {e}")
    import traceback
    traceback.print_exc()
    opd_results = None

# Comparison
if std_results and opd_results:
    print(f"\n{'='*70}")
    print("COMPARISON: OPD vs Standard Method")
    print('='*70)

    print(f"\n{'Metric':<25} {'Standard':<12} {'OPD':<12} {'Delta':<12} {'Delta %':<10}")
    print("-" * 70)

    for key, label in [('40_20', '40-20 RMS (nm)'), ('60_20', '60-20 RMS (nm)'), ('90_mfg', '90 MFG (nm)')]:
        std_val = std_results[key]
        opd_val = opd_results[key]
        delta = opd_val - std_val
        delta_pct = 100 * delta / std_val if std_val > 0 else 0

        print(f"{label:<25} {std_val:<12.2f} {opd_val:<12.2f} {delta:+<12.2f} {delta_pct:+.1f}%")

    print(f"\n{'='*70}")
    print("INTERPRETATION")
    print('='*70)

    delta_40 = opd_results['40_20'] - std_results['40_20']
    delta_60 = opd_results['60_20'] - std_results['60_20']

    print(f"""
The OPD method accounts for lateral (X,Y) displacement when computing WFE.

For telescope mirrors with lateral supports:
- Gravity causes the mirror to shift laterally (X,Y) as well as sag (Z)
- The Standard method ignores this lateral shift
- The OPD method interpolates the ideal surface at deformed (x+dx, y+dy) positions

Key observations:
- 40-20 difference: {delta_40:+.2f} nm ({100*delta_40/std_results['40_20']:+.1f}%)
- 60-20 difference: {delta_60:+.2f} nm ({100*delta_60/std_results['60_20']:+.1f}%)
- Lateral displacement: Max {opd_results['lateral_max']:.3f} µm, RMS {opd_results['lateral_rms']:.3f} µm

Significance:
""")

    if abs(delta_40) < 0.5 and abs(delta_60) < 0.5:
        print("  -> SMALL DIFFERENCE: For this design, lateral displacement is minimal.")
        print("     Both methods give similar results.")
    else:
        print("  -> SIGNIFICANT DIFFERENCE: Lateral displacement affects WFE computation.")
        print("     OPD method is more physically accurate for this geometry.")

    if opd_results['lateral_rms'] > 0.1:
        print(f"\n  WARNING: Lateral RMS {opd_results['lateral_rms']:.3f} µm is notable.")
        print("           OPD method recommended for accurate optimization.")
