"""
Modal Mass Extractor for Dynamic Analysis Results
==================================================

Phase 3 Task 3.4 - NX Open Automation Roadmap

Extracts modal effective mass and participation factors from Nastran F06 files.
This is essential for dynamic optimization where mode shapes affect system response.

Usage:
    from optimization_engine.extractors.extract_modal_mass import extract_modal_mass

    result = extract_modal_mass("path/to/modal.f06", mode=1)
    print(f"Modal mass X: {result['modal_mass_x']} kg")

Supports:
    - SOL 103 (Normal Modes / Eigenvalue Analysis)
    - SOL 111 (Modal Frequency Response)
    - Modal effective mass table parsing
    - Participation factor extraction
"""

import re
import numpy as np
from pathlib import Path
from typing import Dict, Any, Optional, List, Union, Tuple
import logging

logger = logging.getLogger(__name__)


def extract_modal_mass(
    f06_file: Union[str, Path],
    mode: Optional[int] = None,
    direction: str = 'all'
) -> Dict[str, Any]:
    """
    Extract modal effective mass from F06 file.

    Modal effective mass indicates how much of the total mass participates
    in each mode for each direction. Important for:
    - Base excitation problems
    - Seismic analysis
    - Random vibration

    Args:
        f06_file: Path to the F06 results file
        mode: Mode number to extract (1-indexed). If None, returns all modes.
        direction: Direction(s) to extract:
                   'x', 'y', 'z' - single direction
                   'all' - all directions (default)
                   'total' - sum of all directions

    Returns:
        dict: {
            'success': bool,
            'modes': list of mode data (if mode=None),
            'modal_mass_x': float (kg) - effective mass in X,
            'modal_mass_y': float (kg) - effective mass in Y,
            'modal_mass_z': float (kg) - effective mass in Z,
            'modal_mass_rx': float (kg·m²) - rotational mass about X,
            'modal_mass_ry': float (kg·m²) - rotational mass about Y,
            'modal_mass_rz': float (kg·m²) - rotational mass about Z,
            'participation_x': float (0-1) - participation factor X,
            'participation_y': float (0-1) - participation factor Y,
            'participation_z': float (0-1) - participation factor Z,
            'frequency': float (Hz) - natural frequency,
            'cumulative_mass_x': float - cumulative mass fraction X,
            'cumulative_mass_y': float - cumulative mass fraction Y,
            'cumulative_mass_z': float - cumulative mass fraction Z,
            'total_mass': float (kg) - total model mass,
            'error': str or None
        }

    Example:
        >>> result = extract_modal_mass("modal_analysis.f06", mode=1)
        >>> if result['success']:
        ...     print(f"Mode 1 frequency: {result['frequency']:.2f} Hz")
        ...     print(f"X participation: {result['participation_x']*100:.1f}%")
    """
    f06_path = Path(f06_file)

    if not f06_path.exists():
        return {
            'success': False,
            'error': f"F06 file not found: {f06_path}",
            'modes': [],
            'modal_mass_x': None,
            'modal_mass_y': None,
            'modal_mass_z': None,
            'frequency': None
        }

    try:
        with open(f06_path, 'r', errors='ignore') as f:
            content = f.read()

        # Parse modal effective mass table
        modes_data = _parse_modal_effective_mass(content)

        if not modes_data:
            # Try alternative format (participation factors)
            modes_data = _parse_participation_factors(content)

        if not modes_data:
            return {
                'success': False,
                'error': "No modal effective mass or participation factor table found in F06. "
                        "Ensure MEFFMASS case control is present.",
                'modes': [],
                'modal_mass_x': None,
                'modal_mass_y': None,
                'modal_mass_z': None,
                'frequency': None
            }

        # Extract total mass from F06
        total_mass = _extract_total_mass(content)

        if mode is not None:
            # Return specific mode
            if mode < 1 or mode > len(modes_data):
                return {
                    'success': False,
                    'error': f"Mode {mode} not found. Available modes: 1-{len(modes_data)}",
                    'modes': modes_data,
                    'modal_mass_x': None,
                    'modal_mass_y': None,
                    'modal_mass_z': None,
                    'frequency': None
                }

            mode_data = modes_data[mode - 1]
            return {
                'success': True,
                'error': None,
                'mode_number': mode,
                'frequency': mode_data.get('frequency'),
                'modal_mass_x': mode_data.get('mass_x'),
                'modal_mass_y': mode_data.get('mass_y'),
                'modal_mass_z': mode_data.get('mass_z'),
                'modal_mass_rx': mode_data.get('mass_rx'),
                'modal_mass_ry': mode_data.get('mass_ry'),
                'modal_mass_rz': mode_data.get('mass_rz'),
                'participation_x': mode_data.get('participation_x'),
                'participation_y': mode_data.get('participation_y'),
                'participation_z': mode_data.get('participation_z'),
                'cumulative_mass_x': mode_data.get('cumulative_x'),
                'cumulative_mass_y': mode_data.get('cumulative_y'),
                'cumulative_mass_z': mode_data.get('cumulative_z'),
                'total_mass': total_mass,
                'modes': modes_data
            }
        else:
            # Return all modes summary
            return {
                'success': True,
                'error': None,
                'mode_count': len(modes_data),
                'modes': modes_data,
                'total_mass': total_mass,
                'frequencies': [m.get('frequency') for m in modes_data],
                'dominant_mode_x': _find_dominant_mode(modes_data, 'x'),
                'dominant_mode_y': _find_dominant_mode(modes_data, 'y'),
                'dominant_mode_z': _find_dominant_mode(modes_data, 'z'),
            }

    except Exception as e:
        logger.exception(f"Error extracting modal mass from {f06_path}")
        return {
            'success': False,
            'error': str(e),
            'modes': [],
            'modal_mass_x': None,
            'modal_mass_y': None,
            'modal_mass_z': None,
            'frequency': None
        }


def _parse_modal_effective_mass(content: str) -> List[Dict[str, Any]]:
    """Parse modal effective mass table from F06 content."""
    modes = []

    # Pattern for modal effective mass table header
    # This varies by Nastran version, so we try multiple patterns

    # Pattern 1: Standard MEFFMASS output
    # MODE  FREQUENCY  T1  T2  T3  R1  R2  R3
    pattern1 = re.compile(
        r'MODAL\s+EFFECTIVE\s+MASS.*?'
        r'MODE\s+FREQUENCY.*?'
        r'((?:\s*\d+\s+[\d.E+-]+\s+[\d.E+-]+.*?\n)+)',
        re.IGNORECASE | re.DOTALL
    )

    # Pattern 2: Alternative format
    pattern2 = re.compile(
        r'EFFECTIVE\s+MASS\s+FRACTION.*?'
        r'((?:\s*\d+\s+[\d.E+-]+.*?\n)+)',
        re.IGNORECASE | re.DOTALL
    )

    # Try to find modal effective mass table
    match = pattern1.search(content)
    if not match:
        match = pattern2.search(content)

    if match:
        table_text = match.group(1)
        lines = table_text.strip().split('\n')

        for line in lines:
            # Parse each mode line
            # Expected format: MODE FREQ T1 T2 T3 R1 R2 R3 (or subset)
            parts = line.split()
            if len(parts) >= 3:
                try:
                    mode_num = int(parts[0])
                    frequency = float(parts[1])

                    mode_data = {
                        'mode': mode_num,
                        'frequency': frequency,
                        'mass_x': float(parts[2]) if len(parts) > 2 else 0.0,
                        'mass_y': float(parts[3]) if len(parts) > 3 else 0.0,
                        'mass_z': float(parts[4]) if len(parts) > 4 else 0.0,
                        'mass_rx': float(parts[5]) if len(parts) > 5 else 0.0,
                        'mass_ry': float(parts[6]) if len(parts) > 6 else 0.0,
                        'mass_rz': float(parts[7]) if len(parts) > 7 else 0.0,
                    }
                    modes.append(mode_data)
                except (ValueError, IndexError):
                    continue

    return modes


def _parse_participation_factors(content: str) -> List[Dict[str, Any]]:
    """Parse modal participation factors from F06 content."""
    modes = []

    # Pattern for eigenvalue/frequency table with participation
    # This is the more common output format
    eigenvalue_pattern = re.compile(
        r'R E A L\s+E I G E N V A L U E S.*?'
        r'MODE\s+NO\.\s+EXTRACTION\s+ORDER\s+EIGENVALUE\s+RADIANS\s+CYCLES\s+GENERALIZED\s+GENERALIZED\s*\n'
        r'\s+MASS\s+STIFFNESS\s*\n'
        r'((?:\s*\d+\s+\d+\s+[\d.E+-]+\s+[\d.E+-]+\s+[\d.E+-]+\s+[\d.E+-]+\s+[\d.E+-]+.*?\n)+)',
        re.IGNORECASE | re.DOTALL
    )

    match = eigenvalue_pattern.search(content)
    if match:
        table_text = match.group(1)
        lines = table_text.strip().split('\n')

        for line in lines:
            parts = line.split()
            if len(parts) >= 5:
                try:
                    mode_num = int(parts[0])
                    # parts[1] = extraction order
                    # parts[2] = eigenvalue
                    # parts[3] = radians
                    frequency = float(parts[4])  # cycles (Hz)
                    gen_mass = float(parts[5]) if len(parts) > 5 else 1.0
                    gen_stiff = float(parts[6]) if len(parts) > 6 else 0.0

                    mode_data = {
                        'mode': mode_num,
                        'frequency': frequency,
                        'generalized_mass': gen_mass,
                        'generalized_stiffness': gen_stiff,
                        # Participation factors may need separate parsing
                        'mass_x': 0.0,
                        'mass_y': 0.0,
                        'mass_z': 0.0,
                    }
                    modes.append(mode_data)
                except (ValueError, IndexError):
                    continue

    # Also try to find participation factor table
    participation_pattern = re.compile(
        r'MODAL\s+PARTICIPATION\s+FACTORS.*?'
        r'((?:\s*\d+\s+[\d.E+-]+\s+[\d.E+-]+\s+[\d.E+-]+.*?\n)+)',
        re.IGNORECASE | re.DOTALL
    )

    match = participation_pattern.search(content)
    if match and modes:
        table_text = match.group(1)
        lines = table_text.strip().split('\n')

        for i, line in enumerate(lines):
            parts = line.split()
            if len(parts) >= 4 and i < len(modes):
                try:
                    modes[i]['participation_x'] = float(parts[1])
                    modes[i]['participation_y'] = float(parts[2])
                    modes[i]['participation_z'] = float(parts[3])
                except (ValueError, IndexError):
                    pass

    return modes


def _extract_total_mass(content: str) -> Optional[float]:
    """Extract total model mass from F06 content."""
    # Pattern for total mass in OLOAD resultant or GPWG
    patterns = [
        r'TOTAL\s+MASS\s*=\s*([\d.E+-]+)',
        r'MASS\s+TOTAL\s*:\s*([\d.E+-]+)',
        r'GPWG.*?MASS\s*=\s*([\d.E+-]+)',
        r'MASS\s+CENTER\s+OF\s+GRAVITY.*?MASS\s*=\s*([\d.E+-]+)',
    ]

    for pattern in patterns:
        match = re.search(pattern, content, re.IGNORECASE)
        if match:
            try:
                return float(match.group(1))
            except ValueError:
                continue

    return None


def _find_dominant_mode(modes: List[Dict], direction: str) -> Dict[str, Any]:
    """Find the mode with highest participation in given direction."""
    if not modes:
        return {'mode': None, 'participation': 0.0}

    key = f'mass_{direction}' if direction in 'xyz' else f'participation_{direction}'

    max_participation = 0.0
    dominant_mode = None

    for mode in modes:
        participation = mode.get(key, 0.0) or 0.0
        if abs(participation) > max_participation:
            max_participation = abs(participation)
            dominant_mode = mode.get('mode')

    return {
        'mode': dominant_mode,
        'participation': max_participation,
        'frequency': modes[dominant_mode - 1].get('frequency') if dominant_mode else None
    }


def extract_frequencies(
    f06_file: Union[str, Path],
    n_modes: Optional[int] = None
) -> Dict[str, Any]:
    """
    Extract natural frequencies from modal analysis F06 file.

    Simpler version of extract_modal_mass that just gets frequencies.

    Args:
        f06_file: Path to F06 file
        n_modes: Number of modes to extract (default: all)

    Returns:
        dict: {
            'success': bool,
            'frequencies': list of frequencies in Hz,
            'mode_count': int,
            'first_frequency': float,
            'error': str or None
        }
    """
    result = extract_modal_mass(f06_file, mode=None)

    if not result['success']:
        return {
            'success': False,
            'error': result['error'],
            'frequencies': [],
            'mode_count': 0,
            'first_frequency': None
        }

    frequencies = result.get('frequencies', [])

    if n_modes is not None:
        frequencies = frequencies[:n_modes]

    return {
        'success': True,
        'error': None,
        'frequencies': frequencies,
        'mode_count': len(frequencies),
        'first_frequency': frequencies[0] if frequencies else None
    }


def get_first_frequency(f06_file: Union[str, Path]) -> float:
    """
    Get first natural frequency from F06 file.

    Convenience function for optimization constraints.
    Returns 0 if extraction fails.

    Args:
        f06_file: Path to F06 file

    Returns:
        float: First natural frequency in Hz, or 0 on failure
    """
    result = extract_modal_mass(f06_file, mode=1)

    if result['success'] and result.get('frequency'):
        return result['frequency']
    else:
        logger.warning(f"Frequency extraction failed: {result.get('error')}")
        return 0.0


def get_modal_mass_ratio(
    f06_file: Union[str, Path],
    direction: str = 'z',
    n_modes: int = 10
) -> float:
    """
    Get cumulative modal mass ratio for first n modes.

    This indicates what fraction of total mass participates in the
    first n modes. Important for determining if enough modes are included.

    Args:
        f06_file: Path to F06 file
        direction: Direction ('x', 'y', or 'z')
        n_modes: Number of modes to include

    Returns:
        float: Cumulative mass ratio (0-1), or 0 on failure
    """
    result = extract_modal_mass(f06_file, mode=None)

    if not result['success']:
        return 0.0

    modes = result.get('modes', [])[:n_modes]
    key = f'mass_{direction}'

    total_participation = sum(abs(m.get(key, 0.0) or 0.0) for m in modes)
    total_mass = result.get('total_mass')

    if total_mass and total_mass > 0:
        return total_participation / total_mass
    else:
        return total_participation


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        f06_file = sys.argv[1]
        mode = int(sys.argv[2]) if len(sys.argv) > 2 else None

        print(f"Extracting modal mass from: {f06_file}")

        if mode:
            result = extract_modal_mass(f06_file, mode=mode)
            if result['success']:
                print(f"\n=== Mode {mode} Results ===")
                print(f"Frequency: {result['frequency']:.4f} Hz")
                print(f"Modal mass X: {result['modal_mass_x']}")
                print(f"Modal mass Y: {result['modal_mass_y']}")
                print(f"Modal mass Z: {result['modal_mass_z']}")
            else:
                print(f"Error: {result['error']}")
        else:
            result = extract_modal_mass(f06_file)
            if result['success']:
                print(f"\n=== Modal Analysis Results ===")
                print(f"Mode count: {result['mode_count']}")
                print(f"Total mass: {result['total_mass']}")
                print(f"\nFrequencies (Hz):")
                for i, freq in enumerate(result['frequencies'][:10], 1):
                    print(f"  Mode {i}: {freq:.4f} Hz")
                if len(result['frequencies']) > 10:
                    print(f"  ... and {len(result['frequencies']) - 10} more modes")
            else:
                print(f"Error: {result['error']}")
    else:
        print("Usage: python extract_modal_mass.py <f06_file> [mode_number]")