"""Quick test script for displacement field extraction."""
import h5py
import numpy as np
from pathlib import Path

# Test file
h5_path = Path("C:/Users/Antoine/Atomizer/studies/m1_mirror_adaptive_V11/gnn_data/trial_0091/displacement_field.h5")

print(f"Testing: {h5_path}")
print(f"Exists: {h5_path.exists()}")

if h5_path.exists():
    with h5py.File(h5_path, 'r') as f:
        print(f"\nDatasets in file: {list(f.keys())}")

        node_coords = f['node_coords'][:]
        node_ids = f['node_ids'][:]

        print(f"\nTotal nodes: {len(node_ids)}")

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

        # Check nodes in optical surface range (100-650 mm radius)
        surface_mask = (r >= 100) & (r <= 650)
        print(f"Nodes in r=[100, 650]: {np.sum(surface_mask)}")

        # Check subcases
        subcases = [k for k in f.keys() if k.startswith("subcase_")]
        print(f"Subcases: {subcases}")

        if subcases:
            for sc in subcases:
                disp = f[sc][:]
                print(f"  {sc}: Z-disp range [{disp.min():.4f}, {disp.max():.4f}] mm")