"""The stepresult module.

Contains capabilities for generating and saving results of a single depletion
timestep.
"""

from collections import OrderedDict
import copy
import warnings

import h5py
import numpy as np

import openmc
from openmc.mpi import comm, MPI
from .reaction_rates import ReactionRates

VERSION_RESULTS = (1, 1)


__all__ = ["StepResult"]


class StepResult:
    """Result of a single depletion timestep

    .. versionchanged:: 0.13.1
        Name changed from ``Results`` to ``StepResult``

    Attributes
    ----------
    k : list of (float, float)
        Eigenvalue and uncertainty for each substep.
    time : list of float
        Time at beginning, end of step, in seconds.
    source_rate : float
        Source rate during timestep in [W] or [neutron/sec]
    n_mat : int
        Number of mats.
    n_nuc : int
        Number of nuclides.
    rates : list of ReactionRates
        The reaction rates for each substep.
    volume : OrderedDict of str to float
        Dictionary mapping mat id to volume.
    index_mat : OrderedDict of str to int
        A dictionary mapping mat ID as string to index.
    index_nuc : OrderedDict of str to int
        A dictionary mapping nuclide name as string to index.
    mat_to_hdf5_ind : OrderedDict of str to int
        A dictionary mapping mat ID as string to global index.
    n_hdf5_mats : int
        Number of materials in entire geometry.
    n_stages : int
        Number of stages in simulation.
    data : numpy.ndarray
        Atom quantity, stored by stage, mat, then by nuclide.
    proc_time : int
        Average time spent depleting a material across all
        materials and processes

    """
    def __init__(self):
        self.k = None
        self.time = None
        self.source_rate = None
        self.rates = None
        self.volume = None
        self.proc_time = None

        self.index_mat = None
        self.index_nuc = None
        self.mat_to_hdf5_ind = None

        self.data = None

    def __repr__(self):
        t = self.time[0]
        dt = self.time[1] - self.time[0]
        return f"<StepResult: t={t}, dt={dt}, source={self.source_rate}>"

    def __getitem__(self, pos):
        """Retrieves an item from results.

        Parameters
        ----------
        pos : tuple
            A three-length tuple containing a stage index, mat index and a nuc
            index.  All can be integers or slices.  The second two can be
            strings corresponding to their respective dictionary.

        Returns
        -------
        float
            The atoms for stage, mat, nuc

        """
        stage, mat, nuc = pos
        if isinstance(mat, openmc.Material):
            mat = str(mat.id)
        if isinstance(mat, str):
            mat = self.index_mat[mat]
        if isinstance(nuc, str):
            nuc = self.index_nuc[nuc]

        return self.data[stage, mat, nuc]

    def __setitem__(self, pos, val):
        """Sets an item from results.

        Parameters
        ----------
        pos : tuple
            A three-length tuple containing a stage index, mat index and a nuc
            index.  All can be integers or slices.  The second two can be
            strings corresponding to their respective dictionary.

        val : float
            The value to set data to.

        """
        stage, mat, nuc = pos
        if isinstance(mat, str):
            mat = self.index_mat[mat]
        if isinstance(nuc, str):
            nuc = self.index_nuc[nuc]

        self.data[stage, mat, nuc] = val

    @property
    def n_mat(self):
        return len(self.index_mat)

    @property
    def n_nuc(self):
        return len(self.index_nuc)

    @property
    def n_hdf5_mats(self):
        return len(self.mat_to_hdf5_ind)

    @property
    def n_stages(self):
        return self.data.shape[0]

    def allocate(self, volume, nuc_list, burn_list, full_burn_list, stages):
        """Allocate memory for depletion step data

        Parameters
        ----------
        volume : dict of str float
            Volumes corresponding to materials in full_burn_dict
        nuc_list : list of str
            A list of all nuclide names. Used for sorting the simulation.
        burn_list : list of int
            A list of all mat IDs to be burned.  Used for sorting the simulation.
        full_burn_list : list of str
            List of all burnable material IDs
        stages : int
            Number of stages in simulation.

        """
        self.volume = copy.deepcopy(volume)
        self.index_nuc = {nuc: i for i, nuc in enumerate(nuc_list)}
        self.index_mat = {mat: i for i, mat in enumerate(burn_list)}
        self.mat_to_hdf5_ind = {mat: i for i, mat in enumerate(full_burn_list)}

        # Create storage array
        self.data = np.zeros((stages, self.n_mat, self.n_nuc))

    def distribute(self, local_materials, ranges):
        """Create a new object containing data for distributed materials

        Parameters
        ----------
        local_materials : iterable of str
            Materials for this process
        ranges : iterable of int
            Slice-like object indicating indicies of ``local_materials``
            in the material dimension of :attr:`data` and each element
            in :attr:`rates`

        Returns
        -------
        StepResult
            New results object
        """
        new = StepResult()
        new.volume = {lm: self.volume[lm] for lm in local_materials}
        new.index_mat = {mat: idx for (idx, mat) in enumerate(local_materials)}

        # Direct transfer
        direct_attrs = ("time", "k", "source_rate", "index_nuc",
                        "mat_to_hdf5_ind", "proc_time")
        for attr in direct_attrs:
            setattr(new, attr, getattr(self, attr))
        # Get applicable slice of data
        new.data = self.data[:, ranges]
        new.rates = [r[ranges] for r in self.rates]
        return new

    def get_material(self, mat_id):
        """Return material object for given depleted composition

        .. versionadded:: 0.13.2

        Parameters
        ----------
        mat_id : str
            Material ID as a string

        Returns
        -------
        openmc.Material
            Equivalent material
        """
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', openmc.IDWarning)
            material = openmc.Material(material_id=int(mat_id))
        vol = self.volume[mat_id]
        for nuc, _ in sorted(self.index_nuc.items(), key=lambda x: x[1]):
            atoms = self[0, mat_id, nuc]
            if atoms < 0.0:
                continue
            atom_per_bcm = atoms / vol * 1e-24
            material.add_nuclide(nuc, atom_per_bcm)
        material.volume = vol
        return material

    def export_to_hdf5(self, filename, step):
        """Export results to an HDF5 file

        Parameters
        ----------
        filename : str
            The filename to write to
        step : int
            What step is this?

        """
        # Write new file if first time step, else add to existing file
        kwargs = {'mode': "w" if step == 0 else "a"}

        if h5py.get_config().mpi and comm.size > 1:
            # Write results in parallel
            kwargs['driver'] = 'mpio'
            kwargs['comm'] = comm
            with h5py.File(filename, **kwargs) as handle:
                self._to_hdf5(handle, step, parallel=True)
        else:
            # Gather results at root process
            all_results = comm.gather(self)

            # Only root process writes results
            if comm.rank == 0:
                with h5py.File(filename, **kwargs) as handle:
                    for res in all_results:
                        res._to_hdf5(handle, step, parallel=False)

    def _write_hdf5_metadata(self, handle):
        """Writes result metadata in HDF5 file

        Parameters
        ----------
        handle : h5py.File or h5py.Group
            An hdf5 file or group type to store this in.

        """
        # Create and save the 5 dictionaries:
        # quantities
        #   self.index_mat -> self.volume (TODO: support for changing volumes)
        #   self.index_nuc
        # reactions
        #   self.rates[0].index_nuc (can be different from above, above is superset)
        #   self.rates[0].index_rx
        # these are shared by every step of the simulation, and should be deduplicated.

        # Store concentration mat and nuclide dictionaries (along with volumes)

        handle.attrs['version'] = np.array(VERSION_RESULTS)
        handle.attrs['filetype'] = np.string_('depletion results')

        mat_list = sorted(self.mat_to_hdf5_ind, key=int)
        nuc_list = sorted(self.index_nuc)
        rxn_list = sorted(self.rates[0].index_rx)

        n_mats = self.n_hdf5_mats
        n_nuc_number = len(nuc_list)
        n_nuc_rxn = len(self.rates[0].index_nuc)
        n_rxn = len(rxn_list)
        n_stages = self.n_stages

        mat_group = handle.create_group("materials")

        for mat in mat_list:
            mat_single_group = mat_group.create_group(mat)
            mat_single_group.attrs["index"] = self.mat_to_hdf5_ind[mat]
            mat_single_group.attrs["volume"] = self.volume[mat]

        nuc_group = handle.create_group("nuclides")

        for nuc in nuc_list:
            nuc_single_group = nuc_group.create_group(nuc)
            nuc_single_group.attrs["atom number index"] = self.index_nuc[nuc]
            if nuc in self.rates[0].index_nuc:
                nuc_single_group.attrs["reaction rate index"] = self.rates[0].index_nuc[nuc]

        rxn_group = handle.create_group("reactions")

        for rxn in rxn_list:
            rxn_single_group = rxn_group.create_group(rxn)
            rxn_single_group.attrs["index"] = self.rates[0].index_rx[rxn]

        # Construct array storage

        handle.create_dataset("number", (1, n_stages, n_mats, n_nuc_number),
                              maxshape=(None, n_stages, n_mats, n_nuc_number),
                              chunks=(1, 1, n_mats, n_nuc_number),
                              dtype='float64')

        if n_nuc_rxn > 0 and n_rxn > 0:
            handle.create_dataset("reaction rates", (1, n_stages, n_mats, n_nuc_rxn, n_rxn),
                                maxshape=(None, n_stages, n_mats, n_nuc_rxn, n_rxn),
                                chunks=(1, 1, n_mats, n_nuc_rxn, n_rxn),
                                dtype='float64')

        handle.create_dataset("eigenvalues", (1, n_stages, 2),
                              maxshape=(None, n_stages, 2), dtype='float64')

        handle.create_dataset("time", (1, 2), maxshape=(None, 2), dtype='float64')

        handle.create_dataset("source_rate", (1, n_stages), maxshape=(None, n_stages),
                              dtype='float64')

        handle.create_dataset(
            "depletion time", (1,), maxshape=(None,),
            dtype="float64")

    def _to_hdf5(self, handle, index, parallel=False):
        """Converts results object into an hdf5 object.

        Parameters
        ----------
        handle : h5py.File or h5py.Group
            An HDF5 file or group type to store this in.
        index : int
            What step is this?
        parallel : bool
            Being called with parallel HDF5?

        """
        if "/number" not in handle:
            if parallel:
                comm.barrier()
            self._write_hdf5_metadata(handle)

        if parallel:
            comm.barrier()

        # Grab handles
        number_dset = handle["/number"]
        has_reactions = ("reaction rates" in handle)
        if has_reactions:
            rxn_dset = handle["/reaction rates"]
        eigenvalues_dset = handle["/eigenvalues"]
        time_dset = handle["/time"]
        source_rate_dset = handle["/source_rate"]
        proc_time_dset = handle["/depletion time"]

        # Get number of results stored
        number_shape = list(number_dset.shape)
        number_results = number_shape[0]

        new_shape = index + 1

        if number_results < new_shape:
            # Extend first dimension by 1
            number_shape[0] = new_shape
            number_dset.resize(number_shape)

            if has_reactions:
                rxn_shape = list(rxn_dset.shape)
                rxn_shape[0] = new_shape
                rxn_dset.resize(rxn_shape)

            eigenvalues_shape = list(eigenvalues_dset.shape)
            eigenvalues_shape[0] = new_shape
            eigenvalues_dset.resize(eigenvalues_shape)

            time_shape = list(time_dset.shape)
            time_shape[0] = new_shape
            time_dset.resize(time_shape)

            source_rate_shape = list(source_rate_dset.shape)
            source_rate_shape[0] = new_shape
            source_rate_dset.resize(source_rate_shape)

            proc_shape = list(proc_time_dset.shape)
            proc_shape[0] = new_shape
            proc_time_dset.resize(proc_shape)

        # If nothing to write, just return
        if len(self.index_mat) == 0:
            return

        # Add data
        # Note, for the last step, self.n_stages = 1, even if n_stages != 1.
        n_stages = self.n_stages
        inds = [self.mat_to_hdf5_ind[mat] for mat in self.index_mat]
        low = min(inds)
        high = max(inds)
        for i in range(n_stages):
            number_dset[index, i, low:high+1] = self.data[i]
            if has_reactions:
                rxn_dset[index, i, low:high+1] = self.rates[i]
            if comm.rank == 0:
                eigenvalues_dset[index, i] = self.k[i]
        if comm.rank == 0:
            time_dset[index] = self.time
            source_rate_dset[index] = self.source_rate
            if self.proc_time is not None:
                proc_time_dset[index] = (
                    self.proc_time / (comm.size * self.n_hdf5_mats)
                )

    @classmethod
    def from_hdf5(cls, handle, step):
        """Loads results object from HDF5.

        Parameters
        ----------
        handle : h5py.File or h5py.Group
            An HDF5 file or group type to load from.
        step : int
            Index for depletion step
        """
        results = cls()

        # Grab handles
        number_dset = handle["/number"]
        eigenvalues_dset = handle["/eigenvalues"]
        time_dset = handle["/time"]
        if "source_rate" in handle:
            source_rate_dset = handle["/source_rate"]
        else:
            # Older versions used "power" instead of "source_rate"
            source_rate_dset = handle["/power"]

        results.data = number_dset[step, :, :, :]
        results.k = eigenvalues_dset[step, :]
        results.time = time_dset[step, :]
        results.source_rate = source_rate_dset[step, 0]

        if "depletion time" in handle:
            proc_time_dset = handle["/depletion time"]
            if step < proc_time_dset.shape[0]:
                results.proc_time = proc_time_dset[step]

        if results.proc_time is None:
            results.proc_time = np.array([np.nan])

        # Reconstruct dictionaries
        results.volume = OrderedDict()
        results.index_mat = OrderedDict()
        results.index_nuc = OrderedDict()
        rxn_nuc_to_ind = OrderedDict()
        rxn_to_ind = OrderedDict()

        for mat, mat_handle in handle["/materials"].items():
            vol = mat_handle.attrs["volume"]
            ind = mat_handle.attrs["index"]

            results.volume[mat] = vol
            results.index_mat[mat] = ind

        for nuc, nuc_handle in handle["/nuclides"].items():
            ind_atom = nuc_handle.attrs["atom number index"]
            results.index_nuc[nuc] = ind_atom

            if "reaction rate index" in nuc_handle.attrs:
                rxn_nuc_to_ind[nuc] = nuc_handle.attrs["reaction rate index"]

        for rxn, rxn_handle in handle["/reactions"].items():
            rxn_to_ind[rxn] = rxn_handle.attrs["index"]

        results.rates = []
        # Reconstruct reactions
        for i in range(results.n_stages):
            rate = ReactionRates(results.index_mat, rxn_nuc_to_ind, rxn_to_ind, True)

            if "reaction rates" in handle:
                rate[:] = handle["/reaction rates"][step, i, :, :, :]
            results.rates.append(rate)

        return results

    @staticmethod
    def save(op, x, op_results, t, source_rate, step_ind, proc_time=None):
        """Creates and writes depletion results to disk

        Parameters
        ----------
        op : openmc.deplete.abc.TransportOperator
            The operator used to generate these results.
        x : list of list of numpy.array
            The prior x vectors.  Indexed [i][cell] using the above equation.
        op_results : list of openmc.deplete.OperatorResult
            Results of applying transport operator
        t : list of float
            Time indices.
        source_rate : float
            Source rate during time step in [W] or [neutron/sec]
        step_ind : int
            Step index.
        proc_time : float or None
            Total process time spent depleting materials. This may
            be process-dependent and will be reduced across MPI
            processes.

        """
        # Get indexing terms
        vol_dict, nuc_list, burn_list, full_burn_list = op.get_results_info()

        stages = len(x)

        # Create results
        results = StepResult()
        results.allocate(vol_dict, nuc_list, burn_list, full_burn_list, stages)

        n_mat = len(burn_list)

        for i in range(stages):
            for mat_i in range(n_mat):
                results[i, mat_i, :] = x[i][mat_i]

        ks = []
        for r in op_results:
            if isinstance(r.k, type(None)):
                ks += [(None, None)]
            else:
                ks += [(r.k.nominal_value, r.k.std_dev)]
        results.k = ks
        results.rates = [r.rates for r in op_results]
        results.time = t
        results.source_rate = source_rate
        results.proc_time = proc_time
        if results.proc_time is not None:
            results.proc_time = comm.reduce(proc_time, op=MPI.SUM)

        results.export_to_hdf5("depletion_results.h5", step_ind)

    def transfer_volumes(self, model):
        """Transfers volumes from depletion results to geometry

        Parameters
        ----------
        model : OpenMC model to be used in a depletion restart
            calculation

        """

        if not model.materials:
            materials = openmc.Materials(
                model.geometry.get_all_materials().values()
            )
        else:
            materials = model.materials

        for material in materials:
            if material.depletable:
                material.volume = self.volume[str(material.id)]
