import os

from .read_dataset import read_dataset_collective, read_dataset_distributed
from .read_header import read_header
from .select_region import _get_filename, select_region
from .split_selection import split_selection_collective, split_selection_distributed


class _SwiftSnapshotParams:
    def __init__(self, fname, comm, verbose, mpi_read_format, max_concur_io):
        """
        Class to store parameters for pyread_swift.

        Also stores the mpi4py communicator.

        Parameters
        ----------
        fname : string
            The path to the reference Swift snapshot(part)
        comm : mpi4py communicator
        verbose : bool
            Print output?
        mpi_read_format : string
            "distributed":
                The file parts are split between the ranks, each rank loads an
                entire file part(s)
            "collective":
                The cells within a file are split between the ranks, each file
                is loaded collectivley using all ranks (needs parallel-hdf5
                installed)
        max_concur_io : int
            Max number of HDF5 files that can be open at once

        Attributes
        ----------
        fname/comm/verbose/mpi_read_format/max_concur_io : See above
		region_selected_on : int
			Parttype region was selected with
		split_selection_called : bool
			True is split_selection() has been called
		max_size_to_read_at_once : float
			Maximum number of bytes that can be read at one from HDF5 file
		min_in_file_collective : int
			Minimum number of particles in file to allow "collective" reading
		min_in_tot_collective : int
			Minimum number of particles in total to allow "collective" reading
		comm_size : int
			Number of MPI ranks
		comm_rank : int
			ID of this MPI rank
        """

        # Reference snapshot(part).
        self.fname = fname
        assert os.path.isfile(self.fname), f"File {fname} does not exist."

        self.verbose = verbose
        self.mpi_read_format = mpi_read_format
        self.max_concur_io = max_concur_io

        # Has select region or split selection been called?
        self.region_selected_on = -1
        self.split_selection_called = False

        # HDF5 single read limit (2Gb in bytes)
        self.max_size_to_read_at_once = 2 * 1024.0 * 1024 * 1024

        # MPI communicator
        self.comm = comm
        if self.comm is None:
            self.comm_rank = 0
            self.comm_size = 1
        else:
            self.comm_rank = self.comm.Get_rank()
            self.comm_size = self.comm.Get_size()

        # Minimum number of particles in single file to read in collective mode.
        # A single rank will read files below this count.
        self.min_in_file_collective = 1e3

        # Minimum number of particles to read in collective mode.
        # If total selected particles is below this number, one rank reads.
        self.min_in_tot_collective = self.min_in_file_collective * self.comm_size

    def message(self, msg, only_zero=False):
        """
		Print message to stdout, noting the rank it came from.

        Parameters
        ----------
        msg : string
            Message to print
        only_zero : bool
            True for only rank 0 to print the message
		"""

        if self.verbose:
            if only_zero:
                if self.comm_rank == 0:
                    print(f"[Rank {self.comm_rank}] {msg}")
            else:
                print(f"[Rank {self.comm_rank}] {msg}")


class SwiftSnapshot:
    def __init__(
        self,
        fname,
        comm=None,
        verbose=False,
        mpi_read_format="collective",
        max_concur_io=64,
    ):
        """
        Class to read particle data from Swift snapshots.

        Parameters
        ----------
        fname : string
            The path to the reference Swift snapshot(part)
        comm : mpi4py communicator
        verbose : bool
            Print output?
        mpi_read_format : string
            "distributed":
                The file parts are split between the ranks, each rank loads an
                entire file part(s)
            "collective":
                The cells within a file are split between the ranks, each file
                is loaded collectivley using all ranks (needs parallel-hdf5
                installed)
        max_concur_io : int
            Max number of HDF5 files that can be open at once

        Attributes
        ----------
        params : _SwiftSnapshotParams object
            Stores the pyread_swift options
        header : dict
            Stores the header information from the snapshot
        region_data : dict
            Stores the indexes for the particles in the selected region (same
            for all ranks). Generated by self.select_region().
        index_data : dict
            Stores the indexes for the particles this rank will load (unique to
            each rank). Generated by self.split_selection().
        """

        # Store params.
        self.params = _SwiftSnapshotParams(
            fname, comm, verbose, mpi_read_format, max_concur_io
        )

        # Get information from the header.
        self.header = read_header(self.params)

    def split_selection(self):
        """
        Splits the selected particles from self.select_region() between the
        ranks.

        When self.mpi_read_format == "distributed" the HDF5 file part(s) are
        split between the ranks, each rank loads an entire file part(s).

        When self.mpi_read_format = "collective" the top level cells within a
        file are split between the ranks, each file is loaded collectivley
        using all ranks (needs parallel-hdf5 installed).

        Generates "index_data" dict, which is essentially "region_data", but
        just for this rank.
        """

        self.params.message(
            f"Split selection is {self.params.mpi_read_format}.", only_zero=True
        )

        # Make sure self.select_region() has been called first.
        assert (
            self.params.region_selected_on != -1
        ), "Need to call select region before split_selection"

        # Split files between ranks.
        if self.params.mpi_read_format == "distributed":
            self.index_data = split_selection_distributed(self.region_data, self.params)

        # Splot top level cells between ranks.
        else:
            self.index_data = split_selection_collective(self.region_data, self.params)

        # Record we now know what ranks will load what data.
        self.params.split_selection_called = True

    def select_region(self, part_type, x_min, x_max, y_min, y_max, z_min, z_max):
        """
        Find the snapshot files and top level cells that contain particles
        of a given type within a selected cuboidal region.

        Selection is based off the position of the top level cells. Any top
        level cells that intercet the load region will be selected, and their
        particles indexed for loading.

        Generates a "region_data" dict that stores the HDF5 array indexs to
        load from each file. This is the same for all ranks, until
        split_selection is called.

        Parameters
        ----------
        part_type : int
            Parttype to select on
        x_min/x_max : float
            Minimum and maximum bounds in x-dim to select
        y_min/y_max : float
            Minimum and maximum bounds in y-dim to select
        z_min/z_max : float
            Minimum and maximum bounds in z-dim to select
        """

        # Region we are selecting.
        self.params.message(
            "Selecting region x=%.4f->%.4f y=%.4f->%.4f z=%.4f->%.4f PT=%i"
            % (x_min, x_max, y_min, y_max, z_min, z_max, part_type),
            only_zero=True,
        )

        # Find the particles in the files.
        self.region_data = select_region(
            self.params,
            self.header,
            part_type,
            x_min,
            x_max,
            y_min,
            y_max,
            z_min,
            z_max,
        )

        assert (
            self.region_data["total_num_to_load"] > 0
        ), "Found no particles in selected region"
        self.params.region_selected_on = part_type

    def read_dataset(self, parttype, att):
        """
        Read particles from snapshot.

        Parameters
        ----------
        parttype : int
            Parttype of dataset we are loading
        att : string
            Name of dataset we are loading

        Returns
        -------
        data : ndarray
            The particles in the selected region
        """

        assert (
            self.header["NumPart_Total"][parttype] > 0
        ), "No particles of PT=%i found in %s. [%s]" % (
            parttype,
            self.fname,
            self.header["NumPart_Total"],
        )
        assert self.params.split_selection_called, "Need to call split selection first"
        assert (
            self.params.region_selected_on == parttype
        ), "Selected region on PT=%i but trying to read PT=%i" % (
            self.params.region_selected_on,
            parttype,
        )

        if self.params.mpi_read_format == "distributed":
            data = read_dataset_distributed(
                parttype, att, self.params, self.header, self.index_data
            )
        else:
            data = read_dataset_collective(
                parttype,
                att,
                self.params,
                self.header,
                self.region_data,
                self.index_data,
            )

        self.params.message(f"{att} - shape:{data.shape} - dtype:{data.dtype}")

        # Apply coordinates offset.
        if att == "Coordinates" and len(data) > 0:
            if "CoordinatesOffset" in self.header.keys():
                data += self.header["CoordinatesOffset"]

        return data
