"""
Displacement Field Extraction for GNN Training
===============================================

This module extracts full displacement fields from Nastran OP2 files for GNN training.
Unlike the Zernike extractors that reduce to coefficients, this preserves the raw
spatial data that GNN needs to learn the physics.

Key Features:
1. Extract Z-displacement for all optical surface nodes
2. Store node coordinates for graph construction
3. Support for multiple subcases (gravity orientations)
4. HDF5 storage for efficient loading during training

Output Format (HDF5):
    /node_ids         - [n_nodes] int array
    /node_coords      - [n_nodes, 3] float array (X, Y, Z)
    /subcase_1        - [n_nodes] Z-displacement for subcase 1
    /subcase_2        - [n_nodes] Z-displacement for subcase 2
    /subcase_3        - [n_nodes] Z-displacement for subcase 3
    /subcase_4        - [n_nodes] Z-displacement for subcase 4
    /metadata         - JSON string with extraction info

Usage:
    from optimization_engine.gnn.extract_displacement_field import extract_displacement_field

    field_data = extract_displacement_field(op2_path, bdf_path)
    save_field_to_hdf5(field_data, output_path)
"""

import json
import numpy as np
from pathlib import Path
from typing import Dict, Any, Optional, List, Tuple
from datetime import datetime

try:
    import h5py
    HAS_H5PY = True
except ImportError:
    HAS_H5PY = False

from pyNastran.op2.op2 import OP2
from pyNastran.bdf.bdf import BDF


def identify_optical_surface_nodes(
    node_coords: Dict[int, np.ndarray],
    r_inner: float = 100.0,
    r_outer: float = 650.0,
    z_tolerance: float = 100.0
) -> Tuple[List[int], np.ndarray]:
    """
    Identify nodes on the optical surface by spatial filtering.

    The optical surface is identified by:
    1. Radial position (between inner and outer radius)
    2. Consistent Z range (nodes on the curved mirror surface)

    Args:
        node_coords: Dictionary mapping node ID to (X, Y, Z) coordinates
        r_inner: Inner radius cutoff (central hole)
        r_outer: Outer radius limit
        z_tolerance: Maximum Z deviation from mean to include

    Returns:
        Tuple of (node_ids list, coordinates array [n, 3])
    """
    # Get all coordinates as arrays
    nids = list(node_coords.keys())
    coords = np.array([node_coords[nid] for nid in nids])

    # Calculate radial position
    r = np.sqrt(coords[:, 0]**2 + coords[:, 1]**2)

    # Initial radial filter
    radial_mask = (r >= r_inner) & (r <= r_outer)

    # Find nodes in radial range
    radial_nids = np.array(nids)[radial_mask]
    radial_coords = coords[radial_mask]

    if len(radial_coords) == 0:
        raise ValueError(f"No nodes found in radial range [{r_inner}, {r_outer}]")

    # The optical surface should have a relatively small Z range
    z_vals = radial_coords[:, 2]
    z_mean = np.mean(z_vals)

    # Filter to nodes within z_tolerance of the mean Z
    z_mask = np.abs(radial_coords[:, 2] - z_mean) < z_tolerance

    surface_nids = radial_nids[z_mask].tolist()
    surface_coords = radial_coords[z_mask]

    return surface_nids, surface_coords


def extract_displacement_field(
    op2_path: Path,
    bdf_path: Optional[Path] = None,
    r_inner: float = 100.0,
    r_outer: float = 650.0,
    subcases: Optional[List[int]] = None,
    verbose: bool = True
) -> Dict[str, Any]:
    """
    Extract full displacement field for GNN training.

    This function extracts Z-displacement from OP2 files for nodes on the optical
    surface (defined by radial position). It builds node coordinates directly from
    the OP2 data matched against BDF geometry, then filters by radial position.

    Args:
        op2_path: Path to OP2 file
        bdf_path: Path to BDF/DAT file (auto-detected if None)
        r_inner: Inner radius for surface identification (mm)
        r_outer: Outer radius for surface identification (mm)
        subcases: List of subcases to extract (default: [1, 2, 3, 4])
        verbose: Print progress messages

    Returns:
        Dictionary containing:
            - node_ids: List of node IDs on optical surface
            - node_coords: Array [n_nodes, 3] of coordinates
            - z_displacement: Dict mapping subcase -> [n_nodes] Z-displacements
            - metadata: Extraction metadata
    """
    op2_path = Path(op2_path)

    # Find BDF file
    if bdf_path is None:
        for ext in ['.dat', '.bdf']:
            candidate = op2_path.with_suffix(ext)
            if candidate.exists():
                bdf_path = candidate
                break
        if bdf_path is None:
            raise FileNotFoundError(f"No .dat or .bdf found for {op2_path}")

    if subcases is None:
        subcases = [1, 2, 3, 4]

    if verbose:
        print(f"[FIELD] Reading geometry from: {bdf_path.name}")

    # Read geometry from BDF
    bdf = BDF()
    bdf.read_bdf(str(bdf_path))
    node_geo = {int(nid): node.get_position() for nid, node in bdf.nodes.items()}

    if verbose:
        print(f"[FIELD] Total nodes in BDF: {len(node_geo)}")

    # Read OP2
    if verbose:
        print(f"[FIELD] Reading displacements from: {op2_path.name}")
    op2 = OP2()
    op2.read_op2(str(op2_path))

    if not op2.displacements:
        raise RuntimeError("No displacement data in OP2")

    # Extract data by iterating through OP2 nodes and matching to BDF geometry
    # This approach works even when node numbering differs between sources
    subcase_data = {}

    for key, darr in op2.displacements.items():
        isub = int(getattr(darr, 'isubcase', key))
        if isub not in subcases:
            continue

        data = darr.data
        dmat = data[0] if data.ndim == 3 else data
        ngt = darr.node_gridtype
        op2_node_ids = ngt[:, 0] if ngt.ndim == 2 else ngt

        # Build arrays of matched data
        nids = []
        X = []
        Y = []
        Z = []
        disp_z = []

        for i, nid in enumerate(op2_node_ids):
            nid_int = int(nid)
            if nid_int in node_geo:
                pos = node_geo[nid_int]
                nids.append(nid_int)
                X.append(pos[0])
                Y.append(pos[1])
                Z.append(pos[2])
                disp_z.append(float(dmat[i, 2]))  # Z component

        X = np.array(X, dtype=np.float32)
        Y = np.array(Y, dtype=np.float32)
        Z = np.array(Z, dtype=np.float32)
        disp_z = np.array(disp_z, dtype=np.float32)
        nids = np.array(nids, dtype=np.int32)

        # Filter to optical surface by radial position
        r = np.sqrt(X**2 + Y**2)
        surface_mask = (r >= r_inner) & (r <= r_outer)

        subcase_data[isub] = {
            'node_ids': nids[surface_mask],
            'coords': np.column_stack([X[surface_mask], Y[surface_mask], Z[surface_mask]]),
            'disp_z': disp_z[surface_mask],
        }

        if verbose:
            print(f"[FIELD] Subcase {isub}: {len(nids)} matched, {np.sum(surface_mask)} on surface")

    # Get common nodes across all subcases (should be the same)
    all_subcase_keys = list(subcase_data.keys())
    if not all_subcase_keys:
        raise RuntimeError("No subcases found in OP2")

    # Use first subcase to define node list
    ref_subcase = all_subcase_keys[0]
    surface_nids = subcase_data[ref_subcase]['node_ids'].tolist()
    surface_coords = subcase_data[ref_subcase]['coords']

    # Build displacement dict for all subcases
    z_displacement = {}
    for isub in subcases:
        if isub in subcase_data:
            z_displacement[isub] = subcase_data[isub]['disp_z']

    if verbose:
        print(f"[FIELD] Final surface: {len(surface_nids)} nodes")
        r_surface = np.sqrt(surface_coords[:, 0]**2 + surface_coords[:, 1]**2)
        print(f"[FIELD]   Radial range: [{r_surface.min():.1f}, {r_surface.max():.1f}] mm")

    # Build metadata
    metadata = {
        'extraction_timestamp': datetime.now().isoformat(),
        'op2_file': str(op2_path.name),
        'bdf_file': str(bdf_path.name),
        'n_nodes': len(surface_nids),
        'r_inner': r_inner,
        'r_outer': r_outer,
        'subcases': list(z_displacement.keys()),
    }

    return {
        'node_ids': surface_nids,
        'node_coords': surface_coords,
        'z_displacement': z_displacement,
        'metadata': metadata,
    }


def save_field_to_hdf5(
    field_data: Dict[str, Any],
    output_path: Path,
    compression: str = 'gzip'
) -> None:
    """
    Save displacement field data to HDF5 file.

    Args:
        field_data: Output from extract_displacement_field()
        output_path: Path to save HDF5 file
        compression: Compression algorithm ('gzip', 'lzf', or None)
    """
    if not HAS_H5PY:
        raise ImportError("h5py required for HDF5 storage: pip install h5py")

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with h5py.File(output_path, 'w') as f:
        # Node data
        f.create_dataset('node_ids', data=np.array(field_data['node_ids'], dtype=np.int32),
                        compression=compression)
        f.create_dataset('node_coords', data=field_data['node_coords'].astype(np.float32),
                        compression=compression)

        # Displacement for each subcase
        for subcase, z_disp in field_data['z_displacement'].items():
            f.create_dataset(f'subcase_{subcase}', data=z_disp.astype(np.float32),
                           compression=compression)

        # Metadata as JSON string
        f.attrs['metadata'] = json.dumps(field_data['metadata'])

    # Report file size
    size_kb = output_path.stat().st_size / 1024
    print(f"[FIELD] Saved to {output_path.name} ({size_kb:.1f} KB)")


def load_field_from_hdf5(hdf5_path: Path) -> Dict[str, Any]:
    """
    Load displacement field data from HDF5 file.

    Args:
        hdf5_path: Path to HDF5 file

    Returns:
        Dictionary with same structure as extract_displacement_field()
    """
    if not HAS_H5PY:
        raise ImportError("h5py required for HDF5 storage: pip install h5py")

    with h5py.File(hdf5_path, 'r') as f:
        node_ids = f['node_ids'][:].tolist()
        node_coords = f['node_coords'][:]

        # Load subcases
        z_displacement = {}
        for key in f.keys():
            if key.startswith('subcase_'):
                subcase = int(key.split('_')[1])
                z_displacement[subcase] = f[key][:]

        metadata = json.loads(f.attrs['metadata'])

    return {
        'node_ids': node_ids,
        'node_coords': node_coords,
        'z_displacement': z_displacement,
        'metadata': metadata,
    }


def save_field_to_npz(
    field_data: Dict[str, Any],
    output_path: Path
) -> None:
    """
    Save displacement field data to compressed NPZ file (fallback if no h5py).

    Args:
        field_data: Output from extract_displacement_field()
        output_path: Path to save NPZ file
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    save_dict = {
        'node_ids': np.array(field_data['node_ids'], dtype=np.int32),
        'node_coords': field_data['node_coords'].astype(np.float32),
        'metadata_json': np.array([json.dumps(field_data['metadata'])]),
    }

    # Add subcases
    for subcase, z_disp in field_data['z_displacement'].items():
        save_dict[f'subcase_{subcase}'] = z_disp.astype(np.float32)

    np.savez_compressed(output_path, **save_dict)

    size_kb = output_path.stat().st_size / 1024
    print(f"[FIELD] Saved to {output_path.name} ({size_kb:.1f} KB)")


def load_field_from_npz(npz_path: Path) -> Dict[str, Any]:
    """
    Load displacement field data from NPZ file.

    Args:
        npz_path: Path to NPZ file

    Returns:
        Dictionary with same structure as extract_displacement_field()
    """
    data = np.load(npz_path, allow_pickle=True)

    node_ids = data['node_ids'].tolist()
    node_coords = data['node_coords']
    metadata = json.loads(str(data['metadata_json'][0]))

    # Load subcases
    z_displacement = {}
    for key in data.keys():
        if key.startswith('subcase_'):
            subcase = int(key.split('_')[1])
            z_displacement[subcase] = data[key]

    return {
        'node_ids': node_ids,
        'node_coords': node_coords,
        'z_displacement': z_displacement,
        'metadata': metadata,
    }


# =============================================================================
# Convenience functions
# =============================================================================

def save_field(field_data: Dict[str, Any], output_path: Path) -> None:
    """Save field data using best available format (HDF5 preferred)."""
    output_path = Path(output_path)
    if HAS_H5PY and output_path.suffix == '.h5':
        save_field_to_hdf5(field_data, output_path)
    else:
        if output_path.suffix != '.npz':
            output_path = output_path.with_suffix('.npz')
        save_field_to_npz(field_data, output_path)


def load_field(path: Path) -> Dict[str, Any]:
    """Load field data from HDF5 or NPZ file."""
    path = Path(path)
    if path.suffix == '.h5':
        return load_field_from_hdf5(path)
    else:
        return load_field_from_npz(path)


# =============================================================================
# CLI
# =============================================================================

if __name__ == '__main__':
    import sys
    import argparse

    parser = argparse.ArgumentParser(
        description='Extract displacement field from Nastran OP2 for GNN training'
    )
    parser.add_argument('op2_path', type=Path, help='Path to OP2 file')
    parser.add_argument('-o', '--output', type=Path, help='Output path (default: same dir as OP2)')
    parser.add_argument('--r-inner', type=float, default=100.0, help='Inner radius (mm)')
    parser.add_argument('--r-outer', type=float, default=650.0, help='Outer radius (mm)')
    parser.add_argument('--format', choices=['h5', 'npz'], default='h5',
                       help='Output format (default: h5)')

    args = parser.parse_args()

    # Extract field
    field_data = extract_displacement_field(
        args.op2_path,
        r_inner=args.r_inner,
        r_outer=args.r_outer,
    )

    # Determine output path
    if args.output:
        output_path = args.output
    else:
        output_path = args.op2_path.parent / f'displacement_field.{args.format}'

    # Save
    save_field(field_data, output_path)

    # Print summary
    print("\n" + "="*60)
    print("EXTRACTION SUMMARY")
    print("="*60)
    print(f"Nodes: {len(field_data['node_ids'])}")
    print(f"Subcases: {list(field_data['z_displacement'].keys())}")
    for sc, disp in field_data['z_displacement'].items():
        print(f"  Subcase {sc}: Z range [{disp.min():.4f}, {disp.max():.4f}] mm")