"""Test the fixed extraction function directly on OP2."""
import sys
sys.path.insert(0, "C:/Users/Antoine/Atomizer")

import numpy as np
from pathlib import Path
from optimization_engine.gnn.extract_displacement_field import extract_displacement_field

# Test direct extraction from OP2
op2_path = Path("C:/Users/Antoine/Atomizer/studies/m1_mirror_adaptive_V11/2_iterations/iter91/assy_m1_assyfem1_sim1-solution_1.op2")

print(f"Testing extraction from: {op2_path.name}")
print(f"Exists: {op2_path.exists()}")

if op2_path.exists():
    field_data = extract_displacement_field(op2_path, r_inner=100.0, r_outer=650.0)

    print(f"\n=== EXTRACTION RESULT ===")
    print(f"Total surface nodes: {len(field_data['node_ids'])}")

    coords = field_data['node_coords']
    r = np.sqrt(coords[:, 0]**2 + coords[:, 1]**2)
    print(f"Radial range: [{r.min():.1f}, {r.max():.1f}] mm")
    print(f"Z range: [{coords[:, 2].min():.1f}, {coords[:, 2].max():.1f}] mm")

    print(f"\nSubcases: {list(field_data['z_displacement'].keys())}")
    for sc, disp in field_data['z_displacement'].items():
        nan_count = np.sum(np.isnan(disp))
        if nan_count == 0:
            print(f"  Subcase {sc}: Z-disp range [{disp.min():.6f}, {disp.max():.6f}] mm")
        else:
            valid = disp[~np.isnan(disp)]
            print(f"  Subcase {sc}: {nan_count}/{len(disp)} NaN values, valid range: [{valid.min():.6f}, {valid.max():.6f}]")
else:
    print("OP2 file not found!")