"""
BDF Mass Extractor
==================

Extract mass from Nastran BDF/DAT file using pyNastran.
This is more reliable than OP2 GRDPNT extraction since the BDF always contains
material properties and element geometry.

Usage:
    extractor = BDFMassExtractor(bdf_file="model.dat")
    mass_kg = extractor.extract_mass()
"""

from pathlib import Path
from typing import Dict, Any
import numpy as np

try:
    from pyNastran.bdf.bdf import read_bdf
except ImportError:
    raise ImportError("pyNastran is required. Install with: pip install pyNastran")


class BDFMassExtractor:
    """
    Extract mass from Nastran BDF/DAT file.

    This extractor calculates mass from:
    - Element geometry (GRID nodes, element definitions)
    - Material properties (MAT1, etc. with density)
    - Physical properties (PSOLID, PSHELL, etc. with thickness)
    """

    def __init__(self, bdf_file: str):
        """
        Args:
            bdf_file: Path to .dat or .bdf file
        """
        self.bdf_file = Path(bdf_file)
        self._bdf_model = None

    def _load_bdf(self):
        """Lazy load BDF file"""
        if self._bdf_model is None:
            if not self.bdf_file.exists():
                raise FileNotFoundError(f"BDF file not found: {self.bdf_file}")
            # Read BDF with xref=True to cross-reference cards
            self._bdf_model = read_bdf(str(self.bdf_file), xref=True, debug=False)
        return self._bdf_model

    def extract_mass(self) -> Dict[str, Any]:
        """
        Calculate total structural mass from BDF.

        Returns:
            dict: {
                'mass_kg': total mass in kg,
                'mass_g': total mass in grams,
                'mass_ton': total mass in metric tons,
                'cg': [x, y, z] center of gravity (if calculable),
                'num_elements': number of elements,
                'units': 'ton-mm-sec (Nastran default)',
                'breakdown': {element_type: mass} breakdown by element type
            }
        """
        bdf = self._load_bdf()

        # Calculate mass manually by iterating through elements
        total_mass_ton = 0.0
        breakdown = {}
        total_elements = 0
        cg_numerator = np.zeros(3)  # For weighted CG calculation

        for eid, elem in bdf.elements.items():
            try:
                # Get element mass
                elem_mass = elem.Mass()  # Returns mass in Nastran units (ton)

                total_mass_ton += elem_mass

                # Track by element type
                elem_type = elem.type
                if elem_type not in breakdown:
                    breakdown[elem_type] = 0.0
                breakdown[elem_type] += elem_mass

                # Get element centroid for CG calculation
                try:
                    centroid = elem.Centroid()
                    cg_numerator += elem_mass * np.array(centroid)
                except:
                    pass  # Some elements may not have centroid method

                total_elements += 1

            except Exception as e:
                # Some elements might not have Mass() method
                # (e.g., rigid elements, SPCs, etc.)
                continue

        # Calculate center of gravity
        if total_mass_ton > 0:
            cg = cg_numerator / total_mass_ton
        else:
            cg = np.zeros(3)

        # Convert units: NX uses kg-mm-s where "ton" = 1 kg (not 1000 kg!)
        # So the mass returned by pyNastran is already in kg
        mass_kg = float(total_mass_ton)  # Already in kg
        mass_g = mass_kg * 1000.0

        # Breakdown is already in kg
        breakdown_kg = {k: float(v) for k, v in breakdown.items()}

        return {
            'mass_kg': mass_kg,
            'mass_g': mass_g,
            'mass_ton': float(total_mass_ton),  # This is actually kg in NX units
            'cg': cg.tolist() if hasattr(cg, 'tolist') else list(cg),
            'num_elements': total_elements,
            'units': 'kg-mm-s (NX default)',
            'breakdown': breakdown,
            'breakdown_kg': breakdown_kg,
        }

    def get_material_info(self) -> Dict[str, Any]:
        """
        Get material property information from BDF.

        Returns:
            dict: Material ID -> {density, E, nu, ...}
        """
        bdf = self._load_bdf()

        materials = {}
        for mid, mat in bdf.materials.items():
            mat_info = {
                'type': mat.type,
                'mid': mid,
            }

            # MAT1 (isotropic)
            if hasattr(mat, 'rho') and mat.rho is not None:
                mat_info['density'] = float(mat.rho)  # ton/mm^3
                mat_info['density_kg_m3'] = float(mat.rho * 1e12)  # Convert to kg/m^3

            if hasattr(mat, 'e'):
                mat_info['E'] = float(mat.e) if mat.e is not None else None

            if hasattr(mat, 'nu'):
                mat_info['nu'] = float(mat.nu) if mat.nu is not None else None

            materials[mid] = mat_info

        return materials


def extract_mass_from_bdf(bdf_file: str) -> float:
    """
    Convenience function to extract mass in kg.

    Args:
        bdf_file: Path to .dat or .bdf file

    Returns:
        Mass in kilograms
    """
    extractor = BDFMassExtractor(bdf_file)
    result = extractor.extract_mass()
    return result['mass_kg']


if __name__ == "__main__":
    # Example usage
    import sys
    if len(sys.argv) > 1:
        bdf_file = sys.argv[1]
        extractor = BDFMassExtractor(bdf_file)

        # Extract mass
        mass_result = extractor.extract_mass()
        print(f"Mass: {mass_result['mass_kg']:.6f} kg ({mass_result['mass_g']:.2f} g)")
        print(f"CG: {mass_result['cg']}")
        print(f"Elements: {mass_result['num_elements']}")
        print(f"Element breakdown: {mass_result['breakdown']}")

        # Get material info
        materials = extractor.get_material_info()
        print(f"\nMaterials:")
        for mid, mat_info in materials.items():
            print(f"  MAT{mid}: {mat_info}")