"""
Test: J1 coefficient vs mean(WFE) for each subcase individually.

Hypothesis: At 90 deg (zenith), gravity is axially symmetric, so J1 should
closely match mean(WFE). At other angles (20, 40, 60), lateral gravity
components break symmetry, potentially causing J1 != mean.
"""

import numpy as np
from pathlib import Path
import sys

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

from optimization_engine.extractors.extract_zernike import (
    ZernikeExtractor,
    compute_zernike_coefficients,
    DEFAULT_N_MODES,
)


def test_j1_vs_mean_per_subcase():
    """Test J1 vs mean for each subcase in real data."""
    print("=" * 70)
    print("J1 vs mean(WFE) PER SUBCASE")
    print("=" * 70)

    # Find OP2 files
    studies_dir = Path(__file__).parent.parent / "studies"

    # Look for M1 mirror studies specifically
    op2_files = list(studies_dir.rglob("**/m1_mirror*/**/*.op2"))
    if not op2_files:
        op2_files = list(studies_dir.rglob("*.op2"))

    if not op2_files:
        print("No OP2 files found!")
        return

    # Use the first one
    op2_file = op2_files[0]
    print(f"\nUsing: {op2_file.relative_to(studies_dir.parent)}")

    try:
        extractor = ZernikeExtractor(op2_file)
        subcases = list(extractor.displacements.keys())
        print(f"Available subcases: {subcases}")

        print(f"\n{'Subcase':<10} {'Mean(WFE)':<15} {'J1 Coeff':<15} {'|Diff|':<12} {'Diff %':<10}")
        print("-" * 70)

        results = []
        for sc in sorted(subcases, key=lambda x: int(x) if x.isdigit() else 0):
            try:
                X, Y, WFE = extractor._build_dataframe(sc)

                # Compute Zernike coefficients
                coeffs, R_max = compute_zernike_coefficients(X, Y, WFE, DEFAULT_N_MODES)

                j1 = coeffs[0]
                wfe_mean = np.mean(WFE)
                diff = abs(j1 - wfe_mean)
                pct_diff = 100 * diff / abs(wfe_mean) if abs(wfe_mean) > 1e-6 else 0

                print(f"{sc:<10} {wfe_mean:<15.4f} {j1:<15.4f} {diff:<12.4f} {pct_diff:<10.4f}")

                results.append({
                    'subcase': sc,
                    'mean_wfe': wfe_mean,
                    'j1': j1,
                    'diff': diff,
                    'pct_diff': pct_diff
                })

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

        # Also check RELATIVE WFE (e.g., 20 vs 90, 40 vs 90, 60 vs 90)
        print(f"\n" + "=" * 70)
        print("RELATIVE WFE (vs reference subcase)")
        print("=" * 70)

        # Find 90 or use first subcase as reference
        ref_sc = '90' if '90' in subcases else subcases[0]
        print(f"Reference subcase: {ref_sc}")

        print(f"\n{'Relative':<15} {'Mean(WFE_rel)':<15} {'J1 Coeff':<15} {'|Diff|':<12} {'Diff %':<10}")
        print("-" * 70)

        ref_data = extractor.displacements[ref_sc]
        ref_node_to_idx = {int(nid): i for i, nid in enumerate(ref_data['node_ids'])}

        for sc in sorted(subcases, key=lambda x: int(x) if x.isdigit() else 0):
            if sc == ref_sc:
                continue

            try:
                target_data = extractor.displacements[sc]

                X_rel, Y_rel, WFE_rel = [], [], []

                for i, nid in enumerate(target_data['node_ids']):
                    nid = int(nid)
                    if nid not in ref_node_to_idx:
                        continue
                    ref_idx = ref_node_to_idx[nid]
                    geo = extractor.node_geometry.get(nid)
                    if geo is None:
                        continue

                    X_rel.append(geo[0])
                    Y_rel.append(geo[1])

                    target_wfe = target_data['disp'][i, 2] * extractor.wfe_factor
                    ref_wfe = ref_data['disp'][ref_idx, 2] * extractor.wfe_factor
                    WFE_rel.append(target_wfe - ref_wfe)

                X_rel = np.array(X_rel)
                Y_rel = np.array(Y_rel)
                WFE_rel = np.array(WFE_rel)

                # Compute Zernike on relative WFE
                coeffs_rel, _ = compute_zernike_coefficients(X_rel, Y_rel, WFE_rel, DEFAULT_N_MODES)

                j1_rel = coeffs_rel[0]
                wfe_rel_mean = np.mean(WFE_rel)
                diff_rel = abs(j1_rel - wfe_rel_mean)
                pct_diff_rel = 100 * diff_rel / abs(wfe_rel_mean) if abs(wfe_rel_mean) > 1e-6 else 0

                label = f"{sc} vs {ref_sc}"
                print(f"{label:<15} {wfe_rel_mean:<15.4f} {j1_rel:<15.4f} {diff_rel:<12.4f} {pct_diff_rel:<10.4f}")

            except Exception as e:
                print(f"{sc} vs {ref_sc}: ERROR: {e}")

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()


def test_symmetry_analysis():
    """Analyze why J1 != mean for different subcases."""
    print(f"\n" + "=" * 70)
    print("SYMMETRY ANALYSIS")
    print("=" * 70)
    print("""
Theory: J1 (piston) should equal mean(WFE) when:
1. The aperture is circular/annular AND
2. The sampling is uniform in angle AND
3. The WFE has no bias correlated with position

At 90 deg (zenith):
- Gravity acts purely in Z direction
- Deformation should be axially symmetric
- J1 should closely match mean(WFE)

At 20/40/60 deg:
- Gravity has lateral (X,Y) components
- Deformation may have asymmetric patterns
- Tip/tilt (J2,J3) will be large
- But J1 vs mean should still be close IF sampling is uniform

The difference J1-mean comes from:
- Non-uniform radial sampling (mesh density varies)
- Correlation between WFE and position (asymmetric loading)
""")


if __name__ == "__main__":
    test_j1_vs_mean_per_subcase()
    test_symmetry_analysis()