"""
Temperature Extractor for Thermal Analysis Results
===================================================

Phase 3 Task 3.1 - NX Open Automation Roadmap

Extracts nodal temperature results from Nastran OP2 files for thermal optimization.

Usage:
    from optimization_engine.extractors.extract_temperature import extract_temperature

    result = extract_temperature("path/to/thermal.op2", subcase=1)
    print(f"Max temperature: {result['max_temperature']} K")

Supports:
    - SOL 153 (Steady-State Heat Transfer)
    - SOL 159 (Transient Heat Transfer)
    - Thermal subcases in coupled analyses
"""

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

logger = logging.getLogger(__name__)


def extract_temperature(
    op2_file: Union[str, Path],
    subcase: int = 1,
    nodes: Optional[List[int]] = None,
    return_field: bool = False
) -> Dict[str, Any]:
    """
    Extract nodal temperatures from thermal analysis OP2 file.

    Args:
        op2_file: Path to the OP2 results file
        subcase: Subcase number to extract (default: 1)
        nodes: Optional list of specific node IDs to extract.
               If None, extracts all nodes.
        return_field: If True, include full temperature field in result

    Returns:
        dict: {
            'success': bool,
            'max_temperature': float (K or °C depending on model units),
            'min_temperature': float,
            'avg_temperature': float,
            'max_node_id': int (node with max temperature),
            'min_node_id': int (node with min temperature),
            'node_count': int,
            'temperatures': dict (node_id: temp) - only if return_field=True,
            'unit': str ('K' or 'C'),
            'subcase': int,
            'error': str or None
        }

    Example:
        >>> result = extract_temperature("thermal_analysis.op2", subcase=1)
        >>> if result['success']:
        ...     print(f"Max temp: {result['max_temperature']:.1f} K at node {result['max_node_id']}")
        ...     print(f"Temperature range: {result['min_temperature']:.1f} - {result['max_temperature']:.1f} K")
    """
    op2_path = Path(op2_file)

    if not op2_path.exists():
        return {
            'success': False,
            'error': f"OP2 file not found: {op2_path}",
            'max_temperature': None,
            'min_temperature': None,
            'avg_temperature': None,
            'max_node_id': None,
            'min_node_id': None,
            'node_count': 0,
            'subcase': subcase
        }

    try:
        from pyNastran.op2.op2 import read_op2

        # Read OP2 with minimal output
        op2 = read_op2(str(op2_path), load_geometry=False, debug=False, log=None)

        # Check for temperature results
        # pyNastran stores temperatures in different attributes depending on analysis type
        temperatures = None

        # Method 1: Check temperatures attribute (SOL 153/159)
        if hasattr(op2, 'temperatures') and op2.temperatures:
            if subcase in op2.temperatures:
                temp_data = op2.temperatures[subcase]
                temperatures = _extract_from_table(temp_data, nodes)
                logger.debug(f"Found temperatures in op2.temperatures[{subcase}]")

        # Method 2: Check thermal load results (alternative storage)
        if temperatures is None and hasattr(op2, 'thermal_load_vectors'):
            if subcase in op2.thermal_load_vectors:
                temp_data = op2.thermal_load_vectors[subcase]
                temperatures = _extract_from_table(temp_data, nodes)
                logger.debug(f"Found temperatures in op2.thermal_load_vectors[{subcase}]")

        # Method 3: Check displacement as temperature (some solvers store temp in disp)
        # In thermal analysis, "displacement" can actually be temperature
        if temperatures is None and hasattr(op2, 'displacements') and op2.displacements:
            if subcase in op2.displacements:
                disp_data = op2.displacements[subcase]
                # Check if this is actually temperature data
                # Temperature data typically has only 1 DOF (scalar field)
                if hasattr(disp_data, 'data'):
                    data = disp_data.data
                    if len(data.shape) >= 2:
                        # For thermal, we expect scalar temperature at each node
                        # Column 0 of translational data contains temperature
                        temps_array = data[0, :, 0] if len(data.shape) == 3 else data[:, 0]
                        node_ids = disp_data.node_gridtype[:, 0]
                        temperatures = {int(nid): float(temps_array[i])
                                       for i, nid in enumerate(node_ids)
                                       if nodes is None or nid in nodes}
                        logger.debug(f"Found temperatures in op2.displacements[{subcase}] (thermal mode)")

        if temperatures is None or len(temperatures) == 0:
            # List available subcases for debugging
            available_subcases = []
            if hasattr(op2, 'temperatures') and op2.temperatures:
                available_subcases.extend(list(op2.temperatures.keys()))

            return {
                'success': False,
                'error': f"No temperature results found for subcase {subcase}. "
                        f"Available subcases: {available_subcases}",
                'max_temperature': None,
                'min_temperature': None,
                'avg_temperature': None,
                'max_node_id': None,
                'min_node_id': None,
                'node_count': 0,
                'subcase': subcase
            }

        # Compute statistics
        temp_values = np.array(list(temperatures.values()))
        temp_nodes = np.array(list(temperatures.keys()))

        max_idx = np.argmax(temp_values)
        min_idx = np.argmin(temp_values)

        result = {
            'success': True,
            'error': None,
            'max_temperature': float(temp_values[max_idx]),
            'min_temperature': float(temp_values[min_idx]),
            'avg_temperature': float(np.mean(temp_values)),
            'std_temperature': float(np.std(temp_values)),
            'max_node_id': int(temp_nodes[max_idx]),
            'min_node_id': int(temp_nodes[min_idx]),
            'node_count': len(temperatures),
            'unit': 'K',  # Nastran typically uses Kelvin
            'subcase': subcase
        }

        if return_field:
            result['temperatures'] = temperatures

        return result

    except ImportError:
        return {
            'success': False,
            'error': "pyNastran not installed. Install with: pip install pyNastran",
            'max_temperature': None,
            'min_temperature': None,
            'avg_temperature': None,
            'max_node_id': None,
            'min_node_id': None,
            'node_count': 0,
            'subcase': subcase
        }
    except Exception as e:
        logger.exception(f"Error extracting temperature from {op2_path}")
        return {
            'success': False,
            'error': str(e),
            'max_temperature': None,
            'min_temperature': None,
            'avg_temperature': None,
            'max_node_id': None,
            'min_node_id': None,
            'node_count': 0,
            'subcase': subcase
        }


def _extract_from_table(temp_data, nodes: Optional[List[int]] = None) -> Dict[int, float]:
    """Extract temperature values from a pyNastran result table."""
    temperatures = {}

    if hasattr(temp_data, 'data') and hasattr(temp_data, 'node_gridtype'):
        data = temp_data.data
        node_ids = temp_data.node_gridtype[:, 0]

        # Data shape is typically (ntimes, nnodes, ncomponents)
        # For temperature, we want the first component
        if len(data.shape) == 3:
            temp_values = data[0, :, 0]  # First time step, all nodes, first component
        elif len(data.shape) == 2:
            temp_values = data[:, 0]
        else:
            temp_values = data

        for i, nid in enumerate(node_ids):
            if nodes is None or nid in nodes:
                temperatures[int(nid)] = float(temp_values[i])

    return temperatures


def extract_temperature_gradient(
    op2_file: Union[str, Path],
    subcase: int = 1,
    method: str = 'nodal_difference'
) -> Dict[str, Any]:
    """
    Extract temperature gradients from thermal analysis.

    Computes temperature gradients based on nodal temperature differences.
    This is useful for identifying thermal stress hot spots.

    Args:
        op2_file: Path to the OP2 results file
        subcase: Subcase number
        method: Gradient calculation method:
                - 'nodal_difference': Max temperature difference between adjacent nodes
                - 'element_based': Gradient within elements (requires mesh connectivity)

    Returns:
        dict: {
            'success': bool,
            'max_gradient': float (K/mm or temperature units/length),
            'avg_gradient': float,
            'temperature_range': float (max - min temperature),
            'gradient_location': tuple (node_id_hot, node_id_cold),
            'error': str or None
        }

    Note:
        For accurate gradients, element-based calculation requires mesh connectivity
        which may not be available in all OP2 files. The nodal_difference method
        provides an approximation based on temperature range.
    """
    # First extract temperatures
    temp_result = extract_temperature(op2_file, subcase=subcase, return_field=True)

    if not temp_result['success']:
        return {
            'success': False,
            'error': temp_result['error'],
            'max_gradient': None,
            'avg_gradient': None,
            'temperature_range': None,
            'gradient_location': None
        }

    temperatures = temp_result.get('temperatures', {})

    if len(temperatures) < 2:
        return {
            'success': False,
            'error': "Need at least 2 nodes to compute gradient",
            'max_gradient': None,
            'avg_gradient': None,
            'temperature_range': None,
            'gradient_location': None
        }

    # Compute temperature range (proxy for max gradient without mesh)
    temp_range = temp_result['max_temperature'] - temp_result['min_temperature']

    # For nodal_difference method, we report the temperature range
    # True gradient computation would require mesh connectivity
    return {
        'success': True,
        'error': None,
        'max_gradient': temp_range,  # Simplified - actual gradient needs mesh
        'avg_gradient': temp_range / 2,  # Rough estimate
        'temperature_range': temp_range,
        'gradient_location': (temp_result['max_node_id'], temp_result['min_node_id']),
        'max_temperature': temp_result['max_temperature'],
        'min_temperature': temp_result['min_temperature'],
        'note': "Gradient approximated from temperature range. "
                "Accurate gradient requires mesh connectivity."
    }


def extract_heat_flux(
    op2_file: Union[str, Path],
    subcase: int = 1,
    element_type: str = 'all'
) -> Dict[str, Any]:
    """
    Extract element heat flux from thermal analysis OP2 file.

    Args:
        op2_file: Path to the OP2 results file
        subcase: Subcase number
        element_type: Element type to extract ('all', 'ctetra', 'chexa', etc.)

    Returns:
        dict: {
            'success': bool,
            'max_heat_flux': float (W/mm² or model units),
            'min_heat_flux': float,
            'avg_heat_flux': float,
            'max_element_id': int,
            'element_count': int,
            'unit': str,
            'error': str or None
        }
    """
    op2_path = Path(op2_file)

    if not op2_path.exists():
        return {
            'success': False,
            'error': f"OP2 file not found: {op2_path}",
            'max_heat_flux': None,
            'min_heat_flux': None,
            'avg_heat_flux': None,
            'max_element_id': None,
            'element_count': 0,
            'subcase': subcase
        }

    try:
        from pyNastran.op2.op2 import read_op2

        op2 = read_op2(str(op2_path), load_geometry=False, debug=False, log=None)

        # Check for heat flux results
        # pyNastran stores thermal flux in chbdyg_thermal_load or similar
        heat_flux_data = None

        # Check various thermal flux attributes
        flux_attrs = [
            'chbdyg_thermal_load',
            'chbdye_thermal_load',
            'chbdyp_thermal_load',
            'thermalLoad_CONV',
            'thermalLoad_CHBDY'
        ]

        for attr in flux_attrs:
            if hasattr(op2, attr):
                data = getattr(op2, attr)
                if data and subcase in data:
                    heat_flux_data = data[subcase]
                    logger.debug(f"Found heat flux in op2.{attr}[{subcase}]")
                    break

        if heat_flux_data is None:
            return {
                'success': False,
                'error': f"No heat flux results found for subcase {subcase}. "
                        "Heat flux output may not be requested in the analysis.",
                'max_heat_flux': None,
                'min_heat_flux': None,
                'avg_heat_flux': None,
                'max_element_id': None,
                'element_count': 0,
                'subcase': subcase
            }

        # Extract flux values
        if hasattr(heat_flux_data, 'data'):
            data = heat_flux_data.data
            element_ids = heat_flux_data.element if hasattr(heat_flux_data, 'element') else []

            # Flux magnitude
            if len(data.shape) == 3:
                flux_values = np.linalg.norm(data[0, :, :], axis=1)
            else:
                flux_values = np.abs(data[:, 0]) if len(data.shape) == 2 else data

            max_idx = np.argmax(flux_values)

            return {
                'success': True,
                'error': None,
                'max_heat_flux': float(np.max(flux_values)),
                'min_heat_flux': float(np.min(flux_values)),
                'avg_heat_flux': float(np.mean(flux_values)),
                'max_element_id': int(element_ids[max_idx]) if len(element_ids) > max_idx else None,
                'element_count': len(flux_values),
                'unit': 'W/mm²',
                'subcase': subcase
            }

        return {
            'success': False,
            'error': "Could not parse heat flux data format",
            'max_heat_flux': None,
            'min_heat_flux': None,
            'avg_heat_flux': None,
            'max_element_id': None,
            'element_count': 0,
            'subcase': subcase
        }

    except Exception as e:
        logger.exception(f"Error extracting heat flux from {op2_path}")
        return {
            'success': False,
            'error': str(e),
            'max_heat_flux': None,
            'min_heat_flux': None,
            'avg_heat_flux': None,
            'max_element_id': None,
            'element_count': 0,
            'subcase': subcase
        }


# Convenience function for optimization constraints
def get_max_temperature(op2_file: Union[str, Path], subcase: int = 1) -> float:
    """
    Get maximum temperature from OP2 file.

    Convenience function for use in optimization constraints.
    Returns inf if extraction fails.

    Args:
        op2_file: Path to OP2 file
        subcase: Subcase number

    Returns:
        float: Maximum temperature or inf on failure
    """
    result = extract_temperature(op2_file, subcase=subcase)
    if result['success']:
        return result['max_temperature']
    else:
        logger.warning(f"Temperature extraction failed: {result['error']}")
        return float('inf')


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        op2_file = sys.argv[1]
        subcase = int(sys.argv[2]) if len(sys.argv) > 2 else 1

        print(f"Extracting temperature from: {op2_file}")
        result = extract_temperature(op2_file, subcase=subcase)

        if result['success']:
            print(f"\n=== Temperature Results (Subcase {subcase}) ===")
            print(f"Max temperature: {result['max_temperature']:.2f} {result['unit']} (node {result['max_node_id']})")
            print(f"Min temperature: {result['min_temperature']:.2f} {result['unit']} (node {result['min_node_id']})")
            print(f"Avg temperature: {result['avg_temperature']:.2f} {result['unit']}")
            print(f"Node count: {result['node_count']}")
        else:
            print(f"Error: {result['error']}")
    else:
        print("Usage: python extract_temperature.py <op2_file> [subcase]")