materials_for_mc 0.2.0

use crate::nuclide::Nuclide;
use crate::reaction::Reaction;
#[cfg(feature = "pyo3")]
use pyo3::prelude::*;
#[cfg(feature = "pyo3")]
use pyo3::types::PyDict;
use std::collections::HashMap;



#[cfg(feature = "pyo3")]
/// Nuclide data container exposed to Python.
///
/// Create a new (optionally named) nuclide instance.
///
/// Args:
///     name (Optional[str]): Optional nuclide identifier (e.g. "Li6", "Fe56"). If not
///         supplied you must pass `path` to `read_nuclide_from_json` later.
///
/// Notes:
///     Individual fields (e.g. `name`, `atomic_number`, `available_temperatures`,
///     `loaded_temperatures`) are exposed as read-only attributes via PyO3 getters.
///     Detailed descriptions appear once each in the generated documentation—this
///     summary omits a full per-attribute list to avoid duplication.
#[pyclass(name = "Nuclide")]
#[derive(Clone, Default)]
pub struct PyNuclide {
    pub name: Option<String>,
    pub element: Option<String>,
    pub atomic_symbol: Option<String>,
    pub atomic_number: Option<u32>,
    pub neutron_number: Option<u32>,
    pub mass_number: Option<u32>,
    pub library: Option<String>,
    pub energy: Option<HashMap<String, Vec<f64>>>,
    pub reactions: HashMap<String, HashMap<i32, Reaction>>,
    pub fissionable: bool,
    pub available_temperatures: Vec<String>,
    pub loaded_temperatures: Vec<String>,
    pub data_path: Option<String>,
}

#[cfg(feature = "pyo3")]
#[pymethods]
impl PyNuclide {
    /// Name / identifier for the nuclide (e.g. "Li6", "Fe56").
    ///
    /// Returns:
    ///     Optional[str]: Nuclide name or None if not yet set.
    #[getter]
    pub fn name(&self) -> Option<String> {
        self.name.clone()
    }

    /// Chemical element symbol (e.g. "Fe").
    ///
    /// Returns:
    ///     Optional[str]: Element symbol or None if data not loaded.
    #[getter]
    pub fn element(&self) -> Option<String> {
        self.element.clone()
    }

    /// Atomic symbol (currently same as element symbol).
    ///
    /// Returns:
    ///     Optional[str]: Atomic symbol string.
    #[getter]
    pub fn atomic_symbol(&self) -> Option<String> {
        self.atomic_symbol.clone()
    }

    /// Proton number Z.
    ///
    /// Returns:
    ///     Optional[int]: Atomic number.
    #[getter]
    pub fn atomic_number(&self) -> Option<u32> {
        self.atomic_number
    }

    /// Neutron number N.
    ///
    /// Returns:
    ///     Optional[int]: Neutron count.
    #[getter]
    pub fn neutron_number(&self) -> Option<u32> {
        self.neutron_number
    }

    /// Mass number A = Z + N.
    ///
    /// Returns:
    ///     Optional[int]: Mass number.
    #[getter]
    pub fn mass_number(&self) -> Option<u32> {
        self.mass_number
    }

    /// Originating nuclear data library identifier.
    ///
    /// Returns:
    ///     Optional[str]: Library name/code.
    #[getter]
    pub fn library(&self) -> Option<String> {
        self.library.clone()
    }

    /// Whether the nuclide is fissionable.
    ///
    /// Returns:
    ///     bool: True if fissionable.
    #[getter]
    pub fn fissionable(&self) -> bool {
        self.fissionable
    }

    /// All temperatures present in the source data file.
    ///
    /// Returns:
    ///     List[str]: Temperature labels (e.g. ["293K"]).
    #[getter]
    pub fn available_temperatures(&self) -> Vec<String> {
        self.available_temperatures.clone()
    }

    /// Temperatures actually loaded into memory (subset of available_temperatures).
    ///
    /// Returns:
    ///     List[str]: Loaded temperatures.
    #[getter]
    pub fn loaded_temperatures(&self) -> Vec<String> {
        self.loaded_temperatures.clone()
    }

    /// Path to the data file used to populate this nuclide (if known).
    ///
    /// Returns:
    ///     Optional[str]: Filesystem path or None.
    #[getter]
    pub fn data_path(&self) -> Option<String> {
        self.data_path.clone()
    }
    /// Create a new (optionally named) nuclide.
    ///
    /// Args:
    ///     name (Optional[str]): Optional nuclide identifier (e.g. "Li6", "Fe56"). If not
    ///         supplied you must pass `path` to `read_nuclide_from_json` later.
    ///
    /// Returns:
    ///     Nuclide: A nuclide object with no data loaded yet.
    #[new]
    #[pyo3(text_signature = "(name=None)")]
    pub fn new(name: Option<String>) -> Self {
        PyNuclide {
            name,
            element: None,
            atomic_symbol: None,
            atomic_number: None,
            neutron_number: None,
            mass_number: None,
            library: None,
            energy: None,
            reactions: HashMap::new(),
            fissionable: false,
            available_temperatures: Vec::new(),
            loaded_temperatures: Vec::new(),
            data_path: None,
        }
    }

    /// Load nuclear data from a JSON file.
    ///
    /// You can either provide a JSON file path explicitly via `path` or rely on
    /// the `name` given at construction and the global configuration to resolve it.
    /// Providing an explicit `path` will override any global configuration.
    ///
    /// When `temperatures` is provided only those temperatures are loaded while
    /// `available_temperatures` always lists every temperature present in the
    /// file. The subset actually loaded is stored in `loaded_temperatures`.
    ///
    /// Args:
    ///     path (Optional[str]): Optional path to the nuclide JSON file, keyword 
    ///         (e.g. "tendl-21", "fendl-3.2c"), or filesystem path. If provided,
    ///         this overrides any global configuration for this nuclide. If omitted,
    ///         the constructor `name` is used to look up the path from global config.
    ///     temperatures (Optional[List[str]]): Temperature strings (e.g. ["293K"]).
    ///         If given only these temperatures are loaded.
    ///
    /// Returns:
    ///     None
    ///
    /// Raises:
    ///     ValueError: If neither `path` nor `name` is available, if the nuclide
    ///         name is not found in global configuration (when path not provided),
    ///         or if the JSON cannot be read / parsed.
    ///
    /// Example:
    ///     Compare the same nuclide from different data sources:
    ///     
    ///     >>> # Set global default
    ///     >>> m4mc.Config.set_cross_sections("tendl-21")
    ///     >>> 
    ///     >>> # Load from global config (will use TENDL)
    ///     >>> li6_tendl = m4mc.Nuclide("Li6")
    ///     >>> li6_tendl.read_nuclide_from_json()
    ///     >>> 
    ///     >>> # Override to use FENDL for comparison
    ///     >>> li6_fendl = m4mc.Nuclide("Li6")
    ///     >>> li6_fendl.read_nuclide_from_json("fendl-3.2c")
    ///     >>> 
    ///     >>> # Use custom local file
    ///     >>> li6_custom = m4mc.Nuclide("Li6")
    ///     >>> li6_custom.read_nuclide_from_json("path/to/custom_Li6.json")
    #[pyo3(signature = (path=None, temperatures=None), text_signature = "(self, path=None, temperatures=None)")]
    pub fn read_nuclide_from_json(
        &mut self,
        path: Option<String>,
        temperatures: Option<Vec<String>>,
    ) -> PyResult<()> {
        use std::collections::HashSet;
        
        let temps_set: Option<HashSet<String>> = temperatures.map(|v| v.into_iter().collect());
        
        // Use the new Rust backend method that handles all the complex logic
        let nuclide = crate::nuclide::load_nuclide_for_python(
            path.as_deref(),
            self.name.as_deref(),
            temps_set.as_ref(),
        ).map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        
        // Simple field assignment from the loaded nuclide
        *self = PyNuclide::from(nuclide);
        Ok(())
    }

    /// Mapping of temperature -> MT number -> reaction data.
    ///
    /// Returns:
    ///     Dict[str, Dict[int, Dict[str, Any]]]: Nested dictionary. The innermost
    ///     dictionary has these keys:
    ///
    ///         - cross_section (List[float])
    ///         - threshold_idx (int)
    ///         - interpolation (List[int])
    ///         - energy (Optional[List[float]]): Present when reaction has its own grid
    #[getter]
    pub fn reactions(&self, py: Python) -> PyResult<PyObject> {
        let py_dict = PyDict::new(py);

        // Create a dictionary of temperature -> mt -> reaction
        for (temp, mt_map) in &self.reactions {
            let mt_dict = PyDict::new(py);
            for (mt, reaction) in mt_map {
                let reaction_dict = PyDict::new(py);
                reaction_dict.set_item("cross_section", &reaction.cross_section)?;
                reaction_dict.set_item("threshold_idx", reaction.threshold_idx)?;
                reaction_dict.set_item("interpolation", &reaction.interpolation)?;
                if !reaction.energy.is_empty() {
                    reaction_dict.set_item("energy", &reaction.energy)?;
                }
                mt_dict.set_item(mt, reaction_dict)?;
            }
            py_dict.set_item(temp, mt_dict)?;
        }

        Ok(py_dict.into())
    }

    /// List of MT numbers available for the (first) loaded temperature.
    ///
    /// Returns:
    ///     Optional[List[int]]: List of MT identifiers or None if no data.
    #[getter]
    pub fn reaction_mts(&self) -> Option<Vec<i32>> {
        Nuclide::from(self.clone()).reaction_mts()
    }

    /// Energy grids by temperature.
    ///
    /// Returns:
    ///     Optional[Dict[str, List[float]]]: Map of temperature key to energy grid
    ///     or None if no energy data loaded.
    #[getter]
    pub fn energy(&self, py: Python) -> PyResult<Option<PyObject>> {
        if let Some(energy_map) = &self.energy {
            let py_dict = PyDict::new(py);
            for (temp_key, energy_grid) in energy_map.iter() {
                py_dict.set_item(temp_key, energy_grid)?;
            }
            Ok(Some(py_dict.into()))
        } else {
            Ok(None)
        }
    }

    /// Get the energy grid for a specific temperature.
    ///
    /// Args:
    ///     temperature (str): Temperature key (e.g. "293K").
    ///
    /// Returns:
    ///     Optional[List[float]]: The energy grid or None if not present.
    pub fn energy_grid(&self, temperature: &str) -> Option<Vec<f64>> {
        let nuclide = Nuclide::from(self.clone());
        nuclide.energy_grid(temperature).cloned()
    }

    /// Get energy grid for a specific temperature and MT number.
    ///
    /// Args:
    ///     temperature (str): Temperature to use for reaction data.
    ///     mt (int): ENDF/MT number for the reaction channel.
    ///
    /// Returns:
    ///     Optional[List[float]]: Reaction energy grid if present.
    pub fn get_reaction_energy_grid(&self, temperature: &str, mt: i32) -> Option<Vec<f64>> {
        if let Some(temp_reactions) = self.reactions.get(temperature) {
            if let Some(reaction) = temp_reactions.get(&mt) {
                if !reaction.energy.is_empty() {
                    return Some(reaction.energy.clone());
                }
            }
        }
        None
    }

    /// Get microscopic cross section data for a specific reaction and temperature.
    ///
    /// Args:
    ///     reaction (Union[int, str]): Either an ENDF/MT number (int) or reaction name (str) 
    ///         like "(n,gamma)", "(n,elastic)", "fission", etc.
    ///     temperature (Optional[str]): Temperature to use. If None, uses the single
    ///         loaded temperature if only one is available.
    ///
    /// Returns:
    ///     Tuple[List[float], List[float]]: A tuple of (cross_section_values, energy_grid).
    ///
    /// Raises:
    ///     Exception: If temperature not found, reaction not found, multiple temperatures loaded
    ///         without specifying one, or no data available.
    pub fn microscopic_cross_section(
        &self,
        reaction: &PyAny,
        temperature: Option<&str>,
    ) -> PyResult<(Vec<f64>, Vec<f64>)> {
        let mut nuclide: Nuclide = self.clone().into();
        
        // Handle both integer and string inputs
        let result = if let Ok(mt_num) = reaction.extract::<i32>() {
            nuclide.microscopic_cross_section(mt_num, temperature)
        } else if let Ok(reaction_name) = reaction.extract::<String>() {
            nuclide.microscopic_cross_section(reaction_name, temperature)
        } else {
            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                "reaction must be either an integer (MT number) or string (reaction name)"
            ));
        };
        
        match result {
            Ok((cross_section, energy)) => Ok((cross_section, energy)),
            Err(e) => Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string())),
        }
    }

    /// Sample a reaction based on cross sections at a given energy and temperature.
    ///
    /// This method randomly selects a nuclear reaction channel based on the relative 
    /// cross sections at the specified neutron energy. It uses Monte Carlo sampling
    /// to select between absorption, elastic scattering, fission (if fissionable), 
    /// and non-elastic reactions according to their probabilities.
    ///
    /// Args:
    ///     energy (float): Neutron energy in eV.
    ///     temperature (str): Temperature to use for reaction data (e.g. "294", "300K").
    ///     seed (Optional[int]): Random seed for reproducible sampling. If None, 
    ///         uses system random state.
    ///
    /// Returns:
    ///     Optional[Dict[str, Any]]: Dictionary containing the sampled reaction data:
    ///         - mt_number (int): ENDF/MT number of the sampled reaction
    ///         - cross_section (List[float]): Cross section values in barns
    ///         - threshold_idx (int): Index where reaction becomes active
    ///         - interpolation (List[int]): Interpolation flags
    ///         - energy (List[float]): Reaction energy grid
    ///     Returns None if no reaction could be sampled (e.g., zero total cross section).
    ///
    /// Raises:
    ///     ValueError: If temperature not found or no reaction data available.
    ///
    /// Example:
    ///     >>> nuclide = Nuclide("Li6")
    ///     >>> nuclide.read_nuclide_from_json()
    ///     >>> reaction = nuclide.sample_reaction(1e-3, "294", seed=42)
    ///     >>> if reaction:
    ///     ...     print(f"Sampled MT {reaction['mt_number']}")
    #[pyo3(signature = (energy, temperature, seed=None), text_signature = "(self, energy, temperature, seed=None)")]
    pub fn sample_reaction(
        &self,
        energy: f64,
        temperature: &str,
        seed: Option<u64>,
    ) -> PyResult<Option<PyObject>> {
        use rand::{Rng, SeedableRng};
        use rand::rngs::StdRng;
        use pyo3::types::PyDict;
        use pyo3::Python;

        let nuclide: Nuclide = self.clone().into();
        
        // Create random number generator with optional seed
        let mut rng = if let Some(seed_val) = seed {
            StdRng::seed_from_u64(seed_val)
        } else {
            StdRng::from_entropy()
        };

        // Sample the reaction
        let sampled_reaction = nuclide.sample_reaction(energy, temperature, &mut rng);

        if let Some(reaction) = sampled_reaction {
            // Convert the reaction to a Python dictionary
            Python::with_gil(|py| {
                let reaction_dict = PyDict::new(py);
                reaction_dict.set_item("mt_number", reaction.mt_number)?;
                reaction_dict.set_item("cross_section", &reaction.cross_section)?;
                reaction_dict.set_item("threshold_idx", reaction.threshold_idx)?;
                reaction_dict.set_item("interpolation", &reaction.interpolation)?;
                reaction_dict.set_item("energy", &reaction.energy)?;
                Ok(Some(reaction_dict.into()))
            })
        } else {
            Ok(None)
        }
    }
}

#[cfg(feature = "pyo3")]
impl From<Nuclide> for PyNuclide {
    fn from(n: Nuclide) -> Self {
        PyNuclide {
            name: n.name,
            element: n.element,
            atomic_symbol: n.atomic_symbol,
            atomic_number: n.atomic_number,
            neutron_number: n.neutron_number,
            mass_number: n.mass_number,
            library: n.library,
            energy: n.energy,
            reactions: n.reactions,
            fissionable: n.fissionable,
            available_temperatures: n.available_temperatures,
            loaded_temperatures: n.loaded_temperatures,
            data_path: n.data_path,
        }
    }
}

impl From<PyNuclide> for Nuclide {
    fn from(py: PyNuclide) -> Self {
        Nuclide {
            name: py.name,
            element: py.element,
            atomic_symbol: py.atomic_symbol,
            atomic_number: py.atomic_number,
            neutron_number: py.neutron_number,
            mass_number: py.mass_number,
            library: py.library,
            energy: py.energy,
            reactions: py.reactions,
            fissionable: py.fissionable,
            available_temperatures: py.available_temperatures,
            loaded_temperatures: py.loaded_temperatures,
            data_path: py.data_path,
        }
    }
}

#[cfg(feature = "pyo3")]
#[pyfunction]
/// Read a nuclide JSON file and return a `Nuclide` instance.
///
/// Args:
///     path (str): Path to nuclide JSON file or keyword like "tendl-21".
///
/// Returns:
///     Nuclide: A fully populated `Nuclide` object with all available temperatures loaded.
///
/// Raises:
///     OSError: If the file cannot be opened or parsed.
#[pyo3(text_signature = "(path)")]
pub fn py_read_nuclide_from_json(path: &str) -> PyResult<PyNuclide> {
    let nuclide = crate::nuclide::load_nuclide_from_path_or_keyword(path)
        .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;
    Ok(PyNuclide::from(nuclide))
}

#[cfg(feature = "pyo3")]
#[pyfunction]
/// Clear any internally cached nuclide data.
///
/// This forces subsequent reads to re-parse JSON files.
///
/// Returns:
///     None
#[pyo3(text_signature = "()")]
pub fn clear_nuclide_cache() {
    crate::nuclide::clear_nuclide_cache();
}