"""
SCALE/KENO Interface

Interface for generating SCALE input files and parsing output.
Supports KENO criticality calculations.
"""

import logging
import re
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


class GeometryType(str, Enum):
    """Supported geometry types."""

    SPHERE = "sphere"
    CYLINDER = "cylinder"
    CUBOID = "cuboid"
    HEMISPHERE = "hemisphere"


class ReflectionType(str, Enum):
    """Reflection boundary conditions."""

    BARE = "bare"
    WATER = "water"
    CONCRETE = "concrete"
    STEEL = "steel"


@dataclass
class MaterialComposition:
    """Material composition specification."""

    name: str
    isotopes: dict[str, float] = field(default_factory=dict)  # isotope -> atom fraction
    density: float = 0.0  # g/cm³
    temperature: float = 293.0  # K


@dataclass
class GeometrySpec:
    """Geometry specification for SCALE input."""

    geometry_type: GeometryType
    dimensions: dict[str, float] = field(default_factory=dict)
    position: tuple[float, float, float] = (0.0, 0.0, 0.0)
    material_id: int = 1
    region_id: int = 1


@dataclass
class SCALEInput:
    """Complete SCALE input specification."""

    title: str = "SAGE Generated Input"
    materials: list[MaterialComposition] = field(default_factory=list)
    geometry: list[GeometrySpec] = field(default_factory=list)
    reflection: ReflectionType = ReflectionType.BARE
    keno_params: dict[str, Any] = field(default_factory=dict)


@dataclass
class SCALEResult:
    """Parsed SCALE output result."""

    k_eff: float = 0.0
    k_eff_sigma: float = 0.0
    converged: bool = False
    generations: int = 0
    histories_per_generation: int = 0
    energy_of_avg_lethargy: float = 0.0
    warnings: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)
    raw_output: str = ""


class SCALEInterface:
    """
    Interface for SCALE/KENO criticality calculations.

    Generates SCALE input files from material/geometry specifications
    and parses SCALE output to extract k-eff results.

    Example:
        interface = SCALEInterface()

        # Build input
        input_spec = interface.build_input(
            material="U-235",
            enrichment=0.93,
            mass_kg=10.0,
            geometry="sphere",
            reflection="bare"
        )

        # Generate input file
        input_text = interface.generate_input(input_spec)

        # Parse output
        result = interface.parse_output(scale_output_text)
        print(f"k-eff = {result.k_eff} +/- {result.k_eff_sigma}")
    """

    # Standard material library (simplified)
    MATERIAL_LIBRARY: dict[str, dict[str, Any]] = {
        "U-235": {
            "isotopes": {"U-235": 1.0},
            "metal_density": 19.1,
            "oxide_density": 10.97,
        },
        "U-238": {
            "isotopes": {"U-238": 1.0},
            "metal_density": 19.1,
            "oxide_density": 10.97,
        },
        "Pu-239": {
            "isotopes": {"Pu-239": 1.0},
            "metal_density": 19.86,
            "oxide_density": 11.46,
        },
    }

    # Default KENO parameters
    DEFAULT_KENO_PARAMS = {
        "gen": 500,  # Total generations
        "npg": 1000,  # Neutrons per generation
        "nsk": 50,  # Generations to skip
        "sig": 0.0001,  # Target sigma
    }

    def __init__(self, scale_version: str = "6.2") -> None:
        """Initialize SCALE interface.

        Args:
            scale_version: SCALE version for compatibility
        """
        self.scale_version = scale_version
        logger.info(f"Initialized SCALE interface for version {scale_version}")

    def validate_input(self, config: dict[str, Any]) -> list[str]:
        """
        Validate input configuration.

        Args:
            config: Configuration dictionary with material, geometry, etc.

        Returns:
            List of validation issues (empty if valid)
        """
        issues = []

        # Check required fields
        required = ["material", "geometry"]
        for field_name in required:
            if field_name not in config:
                issues.append(f"Missing required field: {field_name}")

        # Validate material
        material = config.get("material", "")
        if material and material not in self.MATERIAL_LIBRARY:
            issues.append(f"Unknown material: {material}")

        # Validate geometry type
        geometry = config.get("geometry", "")
        if geometry:
            try:
                GeometryType(geometry)
            except ValueError:
                issues.append(f"Invalid geometry type: {geometry}")

        # Validate reflection type
        reflection = config.get("reflection", "bare")
        try:
            ReflectionType(reflection)
        except ValueError:
            issues.append(f"Invalid reflection type: {reflection}")

        # Validate enrichment
        enrichment = config.get("enrichment", 0.0)
        if enrichment < 0.0 or enrichment > 1.0:
            issues.append(f"Enrichment must be between 0 and 1, got {enrichment}")

        # Validate mass
        mass_kg = config.get("mass_kg", 0.0)
        if mass_kg < 0:
            issues.append(f"Mass cannot be negative: {mass_kg}")

        return issues

    def build_input(
        self,
        material: str,
        enrichment: float = 0.93,
        mass_kg: float = 1.0,
        geometry: str = "sphere",
        reflection: str = "bare",
        form: str = "metal",
        **kwargs: Any,
    ) -> SCALEInput:
        """
        Build SCALE input from high-level parameters.

        Args:
            material: Fissile material (e.g., "U-235", "Pu-239")
            enrichment: Enrichment fraction (0.0 to 1.0)
            mass_kg: Mass in kilograms
            geometry: Geometry type (sphere, cylinder, etc.)
            reflection: Reflection type (bare, water, etc.)
            form: Physical form (metal, oxide, solution)
            **kwargs: Additional parameters

        Returns:
            SCALEInput specification
        """
        mat_data = self.MATERIAL_LIBRARY.get(material, {})

        # Calculate density based on form
        if form == "metal":
            density = mat_data.get("metal_density", 19.0)
        elif form == "oxide":
            density = mat_data.get("oxide_density", 11.0)
        elif form == "solution":
            concentration = kwargs.get("concentration_g_per_l", 50.0)
            density = 1.0 + concentration * 0.001  # Simplified
        else:
            density = mat_data.get("metal_density", 19.0)

        # Build isotopic composition with enrichment
        isotopes = {}
        if material == "U-235":
            isotopes["U-235"] = enrichment
            isotopes["U-238"] = 1.0 - enrichment
        else:
            isotopes = mat_data.get("isotopes", {material: 1.0})

        mat_comp = MaterialComposition(
            name=material,
            isotopes=isotopes,
            density=density,
        )

        # Calculate dimensions from mass
        volume_cm3 = (mass_kg * 1000) / density
        geom_type = GeometryType(geometry)

        if geom_type == GeometryType.SPHERE:
            # V = 4/3 * pi * r³
            import math
            radius = (3 * volume_cm3 / (4 * math.pi)) ** (1 / 3)
            dimensions = {"radius": radius}
        elif geom_type == GeometryType.CYLINDER:
            # Assume H/D = 1 for cylinder
            import math
            radius = (volume_cm3 / (2 * math.pi)) ** (1 / 3)
            dimensions = {"radius": radius, "height": 2 * radius}
        else:
            # Cube approximation for other geometries
            side = volume_cm3 ** (1 / 3)
            dimensions = {"x": side, "y": side, "z": side}

        geom_spec = GeometrySpec(
            geometry_type=geom_type,
            dimensions=dimensions,
        )

        return SCALEInput(
            title=f"SAGE: {material} {form} {geometry}",
            materials=[mat_comp],
            geometry=[geom_spec],
            reflection=ReflectionType(reflection),
            keno_params={**self.DEFAULT_KENO_PARAMS, **kwargs.get("keno_params", {})},
        )

    def generate_input(self, input_spec: SCALEInput) -> str:
        """
        Generate SCALE input file text.

        Args:
            input_spec: SCALE input specification

        Returns:
            SCALE input file as string
        """
        lines = []

        # Title
        lines.append("=csas6")
        lines.append(input_spec.title)
        lines.append("v7.1-252n")
        lines.append("")

        # Compositions
        lines.append("read composition")
        for i, mat in enumerate(input_spec.materials, 1):
            for iso, frac in mat.isotopes.items():
                lines.append(f"  {iso} {i} 0 {frac} {mat.temperature} end")
        lines.append("end composition")
        lines.append("")

        # Geometry (simplified KENO-VI format)
        lines.append("read geometry")
        for geom in input_spec.geometry:
            if geom.geometry_type == GeometryType.SPHERE:
                radius = geom.dimensions.get("radius", 10.0)
                lines.append("  global unit 1")
                lines.append(f"    sphere 1 {radius:.4f}")
                lines.append("    media 1 1 1")
                lines.append("    boundary 1")
        lines.append("end geometry")
        lines.append("")

        # KENO parameters
        lines.append("read parameters")
        params = input_spec.keno_params
        lines.append(f"  gen={params.get('gen', 500)}")
        lines.append(f"  npg={params.get('npg', 1000)}")
        lines.append(f"  nsk={params.get('nsk', 50)}")
        lines.append("end parameters")
        lines.append("")

        lines.append("end")
        lines.append("")

        return "\n".join(lines)

    def parse_output(self, output_text: str) -> SCALEResult:
        """
        Parse SCALE output to extract results.

        Args:
            output_text: SCALE output file contents

        Returns:
            SCALEResult with parsed k-eff and metadata
        """
        result = SCALEResult(raw_output=output_text)

        # Parse k-eff
        keff_pattern = r"best estimate system k-eff\s+([\d.]+)\s+\+\s+or\s+-\s+([\d.]+)"
        match = re.search(keff_pattern, output_text, re.IGNORECASE)
        if match:
            result.k_eff = float(match.group(1))
            result.k_eff_sigma = float(match.group(2))
            result.converged = True

        # Alternative pattern
        if not result.converged:
            alt_pattern = r"k-effective\s*=\s*([\d.]+)\s*\+/-\s*([\d.]+)"
            match = re.search(alt_pattern, output_text, re.IGNORECASE)
            if match:
                result.k_eff = float(match.group(1))
                result.k_eff_sigma = float(match.group(2))
                result.converged = True

        # Parse generation info
        gen_pattern = r"(\d+)\s+generations\s+were\s+run"
        match = re.search(gen_pattern, output_text, re.IGNORECASE)
        if match:
            result.generations = int(match.group(1))

        # Check for convergence warnings
        if "not converged" in output_text.lower():
            result.converged = False
            result.warnings.append("Calculation did not converge")

        # Check for errors
        error_pattern = r"error[:\s]+(.+?)(?:\n|$)"
        for match in re.finditer(error_pattern, output_text, re.IGNORECASE):
            result.errors.append(match.group(1).strip())

        # Check for warnings
        warning_pattern = r"warning[:\s]+(.+?)(?:\n|$)"
        for match in re.finditer(warning_pattern, output_text, re.IGNORECASE):
            result.warnings.append(match.group(1).strip())

        return result

    def check_convergence(self, result: SCALEResult) -> tuple[bool, list[str]]:
        """
        Check if SCALE result has converged adequately.

        Args:
            result: Parsed SCALE result

        Returns:
            Tuple of (converged, list of issues)
        """
        issues = []

        if not result.converged:
            issues.append("Calculation did not converge")

        if result.k_eff_sigma > 0.001:
            issues.append(f"Standard deviation too high: {result.k_eff_sigma}")

        if result.generations < 100:
            issues.append(f"Too few generations: {result.generations}")

        if result.errors:
            issues.extend([f"Error: {e}" for e in result.errors])

        return len(issues) == 0, issues
