from xml.etree import ElementTree
import numpy as np
import os
from collections import namedtuple
from pyvlasiator.vlsv.variables import units_predefined


class VMeshInfo:
    def __init__(
        self,
        vblocks: np.ndarray,
        vblock_size: np.ndarray,
        vmin: np.ndarray,
        vmax: np.ndarray,
        dv: np.ndarray,
    ) -> None:
        self.vblocks = vblocks
        self.vblock_size = vblock_size
        self.vmin = vmin
        self.vmax = vmax
        self.dv = dv
        self.cellwithVDF = np.empty(0, dtype=np.uint64)
        self.nblock_C = np.empty(0, dtype=np.int64)

        self.vblocks.flags.writeable = False
        self.vblock_size.flags.writeable = False
        self.vmin.flags.writeable = False
        self.vmax.flags.writeable = False
        self.dv.flags.writeable = False


"Variable information from the VLSV footer."
VarInfo = namedtuple(
    "VarInfo", ["unit", "unitLaTeX", "variableLaTeX", "unitConversion"]
)


class Vlsv:
    def __init__(self, filename: str):
        self.dir, self.name = os.path.split(filename)
        self.fid = open(filename, "rb")
        self.xmlroot = ElementTree.fromstring("<VLSV></VLSV>")
        self.celldict = {}
        self.maxamr = -1
        self.vg_indexes_on_fg = np.array([])  # SEE: map_vg_onto_fg(self)

        self._read_xml_footer()

        if self.has_parameter(name="time"):  # Vlasiator 5.0+
            self.time = self.read(name="time", tag="PARAMETER")
        elif self.has_parameter(name="t"):
            self.time = self.read(name="t", tag="PARAMETER")
        else:
            self.time = -1.0

        # Check if the file is using new or old VLSV format
        # Read parameters
        meshName = "SpatialGrid"
        bbox = self.read(tag="MESH_BBOX", mesh=meshName)
        if bbox is None:
            try:
                # Vlasiator 4- files where the mesh is defined with parameters
                self.ncells = np.array(
                    (
                        self.read_parameter("xcells_ini"),
                        self.read_parameter("ycells_ini"),
                        self.read_parameter("zcells_ini"),
                    ),
                    dtype=int,
                )
                self.block_size = np.array((1, 1, 1), dtype=int)
                self.coordmin = np.array(
                    (
                        self.read_parameter("xmin"),
                        self.read_parameter("ymin"),
                        self.read_parameter("zmin"),
                    ),
                    dtype=float,
                )
                self.coordmax = np.array(
                    (
                        self.read_parameter("xmax"),
                        self.read_parameter("ymax"),
                        self.read_parameter("zmax"),
                    ),
                    dtype=float,
                )
            except:  # dummy values
                self.ncells = np.array((1, 1, 1), dtype=int)
                self.block_size = np.array((1, 1, 1), dtype=int)
                self.coordmin = np.array((0.0, 0.0, 0.0), dtype=float)
                self.coordmax = np.array((1.0, 1.0, 1.0), dtype=float)

        else:
            # Vlasiator 5+ file
            nodeX = self.read(tag="MESH_NODE_CRDS_X", mesh=meshName)
            nodeY = self.read(tag="MESH_NODE_CRDS_Y", mesh=meshName)
            nodeZ = self.read(tag="MESH_NODE_CRDS_Z", mesh=meshName)
            self.ncells = np.fromiter((i for i in bbox[0:3]), dtype=int)
            self.block_size = np.fromiter((i for i in bbox[3:6]), dtype=int)
            self.coordmin = np.array((nodeX[0], nodeY[0], nodeZ[0]), dtype=float)
            self.coordmax = np.array((nodeX[-1], nodeY[-1], nodeZ[-1]), dtype=float)

        self.dcoord = np.fromiter(
            ((self.coordmax[i] - self.coordmin[i]) / self.ncells[i] for i in range(3)),
            dtype=float,
        )

        self.ncells.flags.writeable = False
        self.block_size.flags.writeable = False
        self.coordmin.flags.writeable = False
        self.coordmax.flags.writeable = False
        self.dcoord.flags.writeable = False

        self.meshes = {}

        # Iterate through the XML tree, find all populations
        self.species = []

        for child in self.xmlroot.findall("BLOCKIDS"):
            if "name" in child.attrib:
                popname = child.attrib["name"]
            else:
                popname = "avgs"

            if not popname in self.species:
                self.species.append(popname)

            bbox = self.read(tag="MESH_BBOX", mesh=popname)
            if bbox is None:
                if self.read_parameter("vxblocks_ini") is not None:
                    # Vlasiator 4- files where the mesh is defined with parameters
                    vblocks = np.array(
                        (
                            self.read_parameter("vxblocks_ini"),
                            self.read_parameter("vyblocks_ini"),
                            self.read_parameter("vzblocks_ini"),
                        ),
                        dtype=int,
                    )
                    vblock_size = np.array((4, 4, 4), dtype=int)
                    vmin = np.array(
                        (
                            self.read_parameter("vxmin"),
                            self.read_parameter("vymin"),
                            self.read_parameter("vzmin"),
                        ),
                        dtype=float,
                    )
                    vmax = np.array(
                        (
                            self.read_parameter("vxmax"),
                            self.read_parameter("vymax"),
                            self.read_parameter("vzmax"),
                        ),
                        dtype=float,
                    )

                else:  # no velocity space
                    vblocks = np.array((0, 0, 0), dtype=int)
                    vblock_size = np.array((4, 4, 4), dtype=int)
                    vmin = np.array((0.0, 0.0, 0.0), dtype=float)
                    vmax = np.array((0.0, 0.0, 0.0), dtype=float)
                    dv = np.array((1.0, 1.0, 1.0), dtype=float)

            else:  # Vlasiator 5+ file with bounding box
                nodeX = self.read(tag="MESH_NODE_CRDS_X", mesh=popname)
                nodeY = self.read(tag="MESH_NODE_CRDS_Y", mesh=popname)
                nodeZ = self.read(tag="MESH_NODE_CRDS_Z", mesh=popname)

                vblocks = np.array((*bbox[0:3],), dtype=int)
                vblock_size = np.array((*bbox[3:6],), dtype=int)
                vmin = np.array((nodeX[0], nodeY[0], nodeZ[0]), dtype=float)
                vmax = np.array((nodeX[-1], nodeY[-1], nodeZ[-1]), dtype=float)

            dv = np.fromiter(
                ((vmax[i] - vmin[i]) / vblocks[i] / vblock_size[i] for i in range(3)),
                dtype=float,
            )

            self.meshes[popname] = VMeshInfo(vblocks, vblock_size, vmin, vmax, dv)

            # Precipitation energy bins
            i = 0
            energybins = []
            binexists = True
            while binexists:
                binexists = self.has_parameter(
                    f"{popname}_PrecipitationCentreEnergy{i}"
                )
                if binexists:
                    binvalue = self.read_parameter(
                        f"{popname}_PrecipitationCentreEnergy{i}"
                    )
                    energybins.append(binvalue)
                i += 1
            if i > 1:
                self.precipitationenergybins[popname] = energybins

        self.variable = [
            node.attrib["name"] for node in self.xmlroot.findall("VARIABLE")
        ]

        cellid = self.read(mesh="SpatialGrid", name="CellID", tag="VARIABLE")
        self.cellindex = np.argsort(cellid)
        self.celldict = {cid: i for (i, cid) in enumerate(cellid)}

        self.maxamr = self.getmaxrefinement(cellid)

        self.nodecellwithVDF = self.xmlroot.findall("CELLSWITHBLOCKS")

        if len(self.nodecellwithVDF) == 0:
            self.hasvdf = False
        else:
            self.hasvdf = self.nodecellwithVDF[0].attrib["arraysize"] != "0"

    def __repr__(self) -> str:
        str = (
            f"File       : {self.name}\n"
            f"Time       : {self.time:.4f}\n"
            f"Dimension  : {self.ndims()}\n"
            f"Max AMR lvl: {self.maxamr}\n"
            f"Has VDF    : {self.hasvdf}\n"
            f"Variables  : {self.variable}\n"
        )
        return str

    def ndims(self) -> int:
        """Get the spatial dimension of data."""
        return sum(i > 1 for i in self.ncells)

    def _read_xml_footer(self) -> None:
        """Read the XML footer of the VLSV file."""
        fid = self.fid
        # first 8 bytes indicate endianness
        endianness_offset = 8
        fid.seek(endianness_offset)
        # offset of the XML file
        uint64_byte_amount = 8
        offset = int.from_bytes(fid.read(uint64_byte_amount), "little", signed=True)
        fid.seek(offset)
        xmlstring = fid.read()
        self.xmlroot = ElementTree.fromstring(xmlstring)

    def read(
        self,
        name: str = "",
        tag: str = "",
        mesh: str = "",
        cellids=-1,
    ) -> np.ndarray:
        """
        Read data of name, tag, and mesh from the vlsv file.

        This is the general reading function for all types of variables in VLSV files.

        Parameters
        ----------
        cellids : int or list of int
            If -1 then all data is read. If nonzero then only the vector for the specified
            cell id or cellids is read.
        Returns
        -------
        numpy.ndarray
        """
        if not tag and not name:
            raise ValueError("Must provide either tag or name")

        name = name.lower()

        fid = self.fid

        if "/" in name:
            popname, varname = name.split("/")
        else:
            popname, varname = "pop", name

        # TODO: add data reducers

        for child in self.xmlroot:
            if tag and child.tag != tag:
                continue
            if name and "name" in child.attrib and child.attrib["name"].lower() != name:
                continue
            if mesh and "mesh" in child.attrib and child.attrib["mesh"] != mesh:
                continue
            if child.tag == tag:
                vsize = int(child.attrib["vectorsize"])
                asize = int(child.attrib["arraysize"])
                dsize = int(child.attrib["datasize"])
                dtype = child.attrib["datatype"]
                variable_offset = int(child.text)

                # Select efficient method to read data based on number of cells
                if hasattr(cellids, "__len__"):
                    ncellids = len(cellids)
                    # Read multiple specified cells
                    # For reading a large amount of single cells, it'll be faster to
                    # read all data from the file and sort afterwards.
                    arraydata = []
                    if ncellids > 5000:
                        result_size = ncellids
                        read_size = asize
                        read_offsets = [0]
                    else:  # Read multiple cell ids one-by-one
                        result_size = ncellids
                        read_size = 1
                        read_offsets = [
                            self.celldict[cid] * dsize * vsize for cid in cellids
                        ]
                else:
                    if cellids < 0:  # all cells
                        result_size = asize
                        read_size = asize
                        read_offsets = [0]
                    else:  # parameter or single cell
                        result_size = 1
                        read_size = 1
                        read_offsets = [self.celldict[cellids] * dsize * vsize]

                for r_offset in read_offsets:
                    use_offset = int(variable_offset + r_offset)
                    fid.seek(use_offset)

                    if dtype == "float" and dsize == 4:
                        dtype = np.float32
                    elif dtype == "float" and dsize == 8:
                        dtype = np.float64
                    elif dtype == "int" and dsize == 4:
                        dtype = np.int32
                    elif dtype == "int" and dsize == 8:
                        dtype = np.int64
                    elif dtype == "uint" and dsize == 4:
                        dtype = np.uint32
                    elif dtype == "uint" and dsize == 8:
                        dtype = np.uint64

                    data = np.fromfile(fid, dtype, count=vsize * read_size)

                    if len(read_offsets) != 1:
                        arraydata.append(data)

                if len(read_offsets) == 1 and result_size < read_size:
                    # Many single cell IDs requested
                    # Pick the elements corresponding to the requested cells
                    for cid in cellids:
                        append_offset = self.celldict[cid] * vsize
                        arraydata.append(data[append_offset : append_offset + vsize])
                    data = np.squeeze(np.array(arraydata))
                elif len(read_offsets) != 1:
                    # Not so many single cell IDs requested
                    data = np.squeeze(np.array(arraydata))

                if vsize > 1:
                    data = data.reshape(result_size, vsize)

                if result_size == 1:
                    return data[0]
                else:
                    return data

        if name:
            raise NameError(
                name
                + "/"
                + tag
                + "/"
                + mesh
                + " not found in .vlsv file or in data reducers!"
            )

    def read_variable(
        self, name: str, cellids: int | list[int] | np.ndarray = -1, sorted: bool = True
    ) -> np.ndarray:
        """
        Read variables as numpy arrays from the open vlsv file.

        Parameters
        ----------
        cellids : int or list[int] or np.ndarray
            If -1 then all data is read. If nonzero then only the vector for the specified
            cell id or cellids is read.
        sorted : bool
            If the returned array is sorted by cell IDs. Only applied for full arrays.
        Returns
        -------
        numpy.ndarray
        """

        if self.has_variable(name) and name.startswith("fg_"):
            if not cellids == -1:
                raise ValueError("CellID requests not supported for FSgrid.")
            return self.read_fg_variable(name=name)

        if self.has_variable(name) and name.startswith("ig_"):
            if not cellids == -1:
                raise ValueError("CellID requests not supported for ionosphere.")
            return self.read_ionosphere_variable(name=name)

        raw = self.read(
            mesh="SpatialGrid",
            name=name,
            tag="VARIABLE",
            cellids=cellids,
        )
        if hasattr(cellids, "__len__"):  # part of cells requested
            return np.float32(raw)
        elif cellids == -1:
            if sorted:
                if raw.ndim == 1:
                    v = raw[self.cellindex]
                else:
                    v = raw[self.cellindex, :]
                if v.dtype == np.float64:  # 32-bit is enough for analysis
                    v = np.float32(v)
            else:
                v = raw
        else:
            v = raw

        return v

    def read_fg_variable(self, name: str):
        raw = self.read(
            mesh="fsgrid",
            name=name,
            tag="VARIABLE",
        )

        bbox = tuple(ncell * 2**self.maxamr for ncell in self.ncells)

        # Determine fsgrid domain decomposition
        nIORanks = self.read_parameter("numWritingRanks")  # Int32

        if raw.ndim > 1:
            v = np.empty((*bbox, raw.shape[-1]), dtype=np.float32)
        else:
            v = np.empty(bbox, dtype=np.float32)

        def getDomainDecomposition(globalsize, nproc: int) -> list[int]:
            """Obtain decomposition of this grid over the given number of processors.
            Reference: fsgrid.hpp
            """
            domainDecomp = (1, 1, 1)
            minValue = 1e20
            for i in range(1, min(nproc, globalsize[0]) + 1):
                iBox = max(globalsize[0] / i, 1.0)
                for j in range(1, min(nproc, globalsize[1]) + 1):
                    if i * j > nproc:
                        break
                    jBox = max(globalsize[1] / j, 1.0)
                    for k in range(1, min(nproc, globalsize[2]) + 1):
                        if i * j * k > nproc:
                            continue
                        kBox = max(globalsize[2] / k, 1.0)
                        v = (
                            10 * iBox * jBox * kBox
                            + ((jBox * kBox) if i > 1 else 0)
                            + ((iBox * kBox) if j > 1 else 0)
                            + ((iBox * jBox) if k > 1 else 0)
                        )
                        if i * j * k == nproc and v < minValue:
                            minValue = v
                            domainDecomp = (i, j, k)

            return domainDecomp

        def calcLocalStart(globalCells, nprocs: int, lcells: int) -> int:
            ncells = globalCells // nprocs
            remainder = globalCells % nprocs
            lstart = (
                lcells * (ncells + 1)
                if lcells < remainder
                else lcells * ncells + remainder
            )

            return lstart

        def calcLocalSize(globalCells, nprocs: int, lcells: int) -> int:
            ncells = globalCells // nprocs
            remainder = globalCells % nprocs
            lsize = ncells + 1 if lcells < remainder else ncells

            return lsize

        fgDecomposition = getDomainDecomposition(bbox, nIORanks)

        offsetnow = 0

        for i in range(nIORanks):
            xyz = (
                i // fgDecomposition[2] // fgDecomposition[1],
                i // fgDecomposition[2] % fgDecomposition[1],
                i % fgDecomposition[2],
            )

            lsize = tuple(
                map(
                    lambda i: calcLocalSize(bbox[i], fgDecomposition[i], xyz[i]),
                    range(0, 3),
                )
            )

            lstart = tuple(
                map(
                    lambda i: calcLocalStart(bbox[i], fgDecomposition[i], xyz[i]),
                    range(0, 3),
                )
            )

            offsetnext = offsetnow + np.prod(lsize)
            lend = tuple(st + si for st, si in zip(lstart, lsize))

            # Reorder data
            if raw.ndim > 1:
                v[lstart[0] : lend[0], lstart[1] : lend[1], lstart[2] : lend[2], :] = (
                    raw[offsetnow:offsetnext, :].reshape(
                        *lsize, raw.shape[-1], order="F"
                    )
                )
            else:
                v[lstart[0] : lend[0], lstart[1] : lend[1], lstart[2] : lend[2]] = raw[
                    offsetnow:offsetnext
                ].reshape(*lsize, order="F")

            offsetnow = offsetnext

        return v

    def read_variable_meta(self, var: str):
        unit, unitLaTeX, variableLaTeX, unitConversion = "", "", "", ""

        if var in units_predefined:
            unit, variableLaTeX, unitLaTeX = units_predefined[var]
        elif self.has_variable(var):  # For Vlasiator 5 files, MetaVLSV is included
            for child in self.xmlroot:
                if "name" in child.attrib and child.attrib["name"] == var:
                    if not "unit" in child.attrib:
                        break
                    else:
                        unit = child.attrib["unit"]
                        unitLaTeX = child.attrib["unitLaTeX"]
                        variableLaTeX = child.attrib["variableLaTeX"]
                        unitConversion = child.attrib["unitConversion"]

        return VarInfo(unit, unitLaTeX, variableLaTeX, unitConversion)

    def read_parameter(self, name: str):
        return self.read(name=name, tag="PARAMETER")

    def read_vcells(self, cellid: int, species: str = "proton"):
        """Read raw velocity block data.

        Parameters
        ----------
        cellid :
            Cell ID of the cell whose velocity blocks are read.
        species : str
            Population required.

        Returns
        -------
        numpy.ndarray
            A numpy array with block ids and their data.
        """

        fid = self.fid
        mesh = self.meshes[species]
        vblock_size = mesh.vblock_size

        self.init_cellswithVDF(species)

        # Check that cells have VDF stored
        try:
            cellWithVDFIndex = np.where(mesh.cellwithVDF == cellid)[0][0]
            nblocks = mesh.nblock_C[cellWithVDFIndex]
        except:
            raise ValueError(f"Cell ID {cellid} does not store VDF!")
        # Offset position to vcell storage
        offset_v = np.sum(mesh.nblock_C[0:cellWithVDFIndex], initial=0)

        # Read raw VDF
        for node in self.xmlroot.findall("BLOCKVARIABLE"):
            if node.attrib["name"] == species:
                dsize = int(node.attrib["datasize"])
                offset = int(node.text)
            break

        bsize = np.prod(vblock_size)
        fid.seek(offset_v * bsize * dsize + offset)
        T = np.float32 if dsize == 4 else np.float64
        data = np.fromfile(
            fid,
            dtype=T,
            count=bsize * nblocks,
        ).reshape(nblocks, bsize)

        # Read block IDs
        for node in self.xmlroot.findall("BLOCKIDS"):
            if node.attrib["name"] == species:
                dsize = int(node.attrib["datasize"])
                offset = int(node.text)
            break

        fid.seek(offset_v * dsize + offset)
        T = np.int32 if dsize == 4 else np.int64
        blockIDs = np.fromfile(fid, dtype=T, count=nblocks)

        # Velocity cell IDs and distributions (ordered by blocks)
        vcellids = np.empty(bsize * nblocks, dtype=np.int32)
        vcellf = np.empty(bsize * nblocks, dtype=np.float32)

        for i, bid in enumerate(blockIDs):
            for j in range(bsize):
                index_ = i * bsize + j
                vcellids[index_] = j + bsize * bid
                vcellf[index_] = data[i, j]

        return vcellids, vcellf

    def init_cellswithVDF(self, species: str = "proton") -> None:
        fid = self.fid
        mesh = self.meshes[species]
        if not np.any(mesh.cellwithVDF):
            for node in self.nodecellwithVDF:
                if node.attrib["name"] == species:
                    asize = int(node.attrib["arraysize"])
                    offset = int(node.text)
                    fid.seek(offset)
                    cellwithVDF = np.fromfile(fid, dtype=np.uint64, count=asize)
                    mesh.cellwithVDF = cellwithVDF
                    break

            for node in self.xmlroot.findall("BLOCKSPERCELL"):
                if node.attrib["name"] == species:
                    asize = int(node.attrib["arraysize"])
                    dsize = int(node.attrib["datasize"])
                    offset = int(node.text)
                    fid.seek(offset)
                    T = np.int32 if dsize == 4 else np.int64
                    nblock_C = np.fromfile(fid, dtype=T, count=asize).astype(np.int64)
                    mesh.nblock_C = nblock_C
                    break

            mesh.cellwithVDF = np.delete(mesh.cellwithVDF, np.where(mesh.nblock_C == 0))

    def _has_attribute(self, attribute: str, name: str) -> bool:
        """Check if a given attribute exists in the xml."""
        for child in self.xmlroot:
            if child.tag == attribute and "name" in child.attrib:
                if child.attrib["name"].lower() == name.lower():
                    return True
        return False

    def has_variable(self, name: str) -> bool:
        return self._has_attribute("VARIABLE", name)

    def has_parameter(self, name: str) -> bool:
        return self._has_attribute("PARAMETER", name)

    def getmaxrefinement(self, cellid: np.ndarray):
        """Get the maximum spatial refinement level."""
        ncell = np.prod(self.ncells)
        maxamr, cid = 0, ncell
        while cid < max(cellid):
            maxamr += 1
            cid += ncell * 8**maxamr

        return maxamr

    def getcell(self, loc: np.ndarray | tuple[int, ...] | list[int]):
        coordmin, coordmax = self.coordmin, self.coordmax
        dcoord = self.dcoord
        ncells = self.ncells
        celldict = self.celldict
        maxamr = self.maxamr

        for i in range(3):
            if not coordmin[i] < loc[i] < coordmax[i]:
                raise ValueError(f"{i} coordinate out of bound!")

        indices = np.fromiter(
            ((loc[i] - coordmin[i]) // dcoord[i] for i in range(3)), dtype=int
        )

        cid = (
            indices[0] + indices[1] * ncells[0] + indices[2] * ncells[0] * ncells[1] + 1
        )

        ncells_lowerlevel = 0
        ncell = np.prod(ncells)

        for ilevel in range(maxamr):
            if cid in celldict:
                break
            ncells_lowerlevel += (8**ilevel) * ncell
            ratio = 2 ** (ilevel + 1)
            indices = np.fromiter(
                (
                    np.floor((loc[i] - coordmin[i]) / dcoord[i] * ratio)
                    for i in range(3)
                ),
                dtype=int,
            )
            cid = (
                ncells_lowerlevel
                + indices[0]
                + ratio * ncells[0] * indices[1]
                + ratio**2 * ncells[0] * ncells[1] * indices[2]
                + 1
            )

        return cid

    def getvcellcoordinates(
        self, vcellids: np.ndarray, species: str = "proton"
    ) -> np.ndarray:
        mesh = self.meshes[species]
        vblocks = mesh.vblocks
        vblock_size = mesh.vblock_size
        dv = mesh.dv
        vmin = mesh.vmin

        bsize = np.prod(vblock_size)
        blockid = np.fromiter((cid // bsize for cid in vcellids), dtype=int)
        # Get block coordinates
        blockInd = [
            np.array(
                (
                    bid % vblocks[0],
                    bid // vblocks[0] % vblocks[1],
                    bid // (vblocks[0] * vblocks[1]),
                ),
                dtype=int,
            )
            for bid in blockid
        ]
        blockCoord = [
            np.array(
                (
                    bInd[0] * dv[0] * vblock_size[0] + vmin[0],
                    bInd[1] * dv[1] * vblock_size[1] + vmin[1],
                    bInd[2] * dv[2] * vblock_size[2] + vmin[2],
                ),
                dtype=float,
            )
            for bInd in blockInd
        ]
        # Get cell indices
        vcellblockids = np.fromiter((vid % bsize for vid in vcellids), dtype=int)
        cellidxyz = np.array(
            [
                np.array(
                    (
                        cid % vblock_size[0],
                        cid // vblock_size[0] % vblock_size[1],
                        cid // (vblock_size[0] * vblock_size[1]),
                    ),
                    dtype=int,
                )
                for cid in vcellblockids
            ]
        )
        # Get cell coordinates
        cellCoords = np.array(
            [
                np.fromiter(
                    (
                        blockCoord[i][j] + (cellidxyz[i][j] + 0.5) * dv[j]
                        for j in range(3)
                    ),
                    dtype=float,
                )
                for i in range(len(vcellids))
            ]
        )

        return cellCoords

    def getnearestcellwithvdf(self, id: int, species: str = "proton"):
        self.init_cellswithVDF(species)
        cells = self.meshes[species].cellwithVDF
        if not np.any(cells):
            raise ValueError(f"No distribution saved in {self.name}")
        coords_orig = self.getcellcoordinates(id)
        coords = [self.getcellcoordinates(cid) for cid in cells]
        min_ = np.argmin(np.sum(np.square(coords - coords_orig), axis=1))

        return cells[min_]

    def getcellcoordinates(self, cid: int):
        ncells = self.ncells
        coordmin, coordmax = self.coordmin, self.coordmax
        cid -= 1  # for easy divisions

        ncells_refmax = list(ncells)
        reflevel = 0
        subtraction = np.prod(ncells) * (2**reflevel) ** 3
        # sizes on the finest level
        while cid >= subtraction:
            cid -= subtraction
            reflevel += 1
            subtraction *= 8
            ncells_refmax[0] *= 2
            ncells_refmax[1] *= 2
            ncells_refmax[2] *= 2

        indices = np.array(
            (
                cid % ncells_refmax[0],
                cid // ncells_refmax[0] % ncells_refmax[1],
                cid // (ncells_refmax[0] * ncells_refmax[1]),
            ),
            dtype=int,
        )

        coords = np.fromiter(
            (
                coordmin[i]
                + (indices[i] + 0.5) * (coordmax[i] - coordmin[i]) / ncells_refmax[i]
                for i in range(3)
            ),
            dtype=float,
        )

        return coords

    def getslicecell(
        self, sliceoffset: float, dir: int, minCoord: float, maxCoord: float
    ):
        if not dir in (0, 1, 2):
            raise ValueError(f"Unknown slice direction {dir}")

        ncells, maxamr, celldict = self.ncells, self.maxamr, self.celldict
        nsize = ncells[dir]
        sliceratio = sliceoffset / (maxCoord - minCoord)
        if not (0.0 <= sliceratio <= 1.0):
            raise ValueError("slice plane index out of bound!")

        # Find the ids
        nlen = 0
        ncell = np.prod(ncells)
        # number of cells up to each refinement level
        lvlC = np.fromiter((ncell * 8**ilvl for ilvl in range(maxamr + 1)), dtype=int)
        lvlAccum = np.add.accumulate(lvlC)
        nStart = np.insert(lvlAccum, 0, 0)

        indexlist = np.empty(0, dtype=int)
        idlist = np.empty(0, dtype=int)

        cellidsorted = np.fromiter(celldict.keys(), dtype=int)
        cellidsorted.sort()

        for ilvl in range(maxamr + 1):
            nLow, nHigh = nStart[ilvl], nStart[ilvl + 1]
            idfirst_ = np.searchsorted(cellidsorted, nLow + 1)
            idlast_ = np.searchsorted(cellidsorted, nHigh, side="right")

            ids = cellidsorted[idfirst_:idlast_]

            ix, iy, iz = getindexes(ilvl, ncells[0], ncells[1], nLow, ids)

            if dir == 0:
                coords = ix
            elif dir == 1:
                coords = iy
            else:
                coords = iz

            # Find the cut plane index for each refinement level (0-based)
            depth = int(np.floor(sliceratio * nsize * 2**ilvl))
            # Find the needed elements to create the cut and save the results
            elements = coords == depth
            indexlist = np.append(indexlist, np.arange(nlen, nlen + len(ids))[elements])
            idlist = np.append(idlist, ids[elements])

            nlen += len(ids)

        return idlist, indexlist

    def refineslice(self, idlist: np.ndarray, data: np.ndarray, normal: int):
        ncells, maxamr = self.ncells, self.maxamr

        dims = _getdim2d(ncells, maxamr, normal)
        # meshgrid-like 2D input for matplotlib
        dpoints = np.empty((dims[1], dims[0]), dtype=data.dtype)

        # Create the plot grid
        ncell = np.prod(ncells)
        nHigh, nLow = ncell, 0

        for i in range(maxamr + 1):
            idfirst_ = np.searchsorted(idlist, nLow + 1)
            idlast_ = np.searchsorted(idlist, nHigh, side="right")

            ids = idlist[idfirst_:idlast_]
            d = data[idfirst_:idlast_]

            ix, iy, iz = getindexes(i, ncells[0], ncells[1], nLow, ids)

            # Get the correct coordinate values and the widths for the plot
            if normal == 0:
                a, b = iy, iz
            elif normal == 1:
                a, b = ix, iz
            elif normal == 2:
                a, b = ix, iy

            # Insert the data values into dpoints
            refineRatio = 2 ** (maxamr - i)
            iRange = range(refineRatio)
            X, Y = np.meshgrid(iRange, iRange, indexing="ij")
            coords = np.empty((len(a), 2 ** (2 * (maxamr - i)), 2), dtype=int)

            for ic, (ac, bc) in enumerate(zip(a, b)):
                for ir in range(2 ** (2 * (maxamr - i))):
                    index_ = np.unravel_index(ir, (refineRatio, refineRatio))
                    coords[ic, ir] = [
                        ac * refineRatio + X[index_],
                        bc * refineRatio + Y[index_],
                    ]

            for ic, dc in enumerate(d):
                for ir in range(2 ** (2 * (maxamr - i))):
                    dpoints[coords[ic, ir, 1], coords[ic, ir, 0]] = dc

            nLow = nHigh
            nHigh += ncell * 8 ** (i + 1)

        return dpoints


def _getdim2d(ncells: tuple, maxamr: int, normal: int):
    ratio = 2**maxamr
    if normal == 0:
        i1, i2 = 1, 2
    elif normal == 1:
        i1, i2 = 0, 2
    elif normal == 2:
        i1, i2 = 0, 1
    dims = (ncells[i1] * ratio, ncells[i2] * ratio)

    return dims


def getindexes(
    ilevel: int, xcells: int, ycells: int, nCellUptoLowerLvl: int, ids: np.ndarray
):
    ratio = 2**ilevel
    slicesize = xcells * ycells * ratio**2

    iz = (ids - nCellUptoLowerLvl - 1) // slicesize
    iy = np.zeros_like(iz)
    ix = np.zeros_like(iz)

    # number of ids up to the coordinate z in the refinement level ilevel
    idUpToZ = iz * slicesize + nCellUptoLowerLvl
    iy = (ids - idUpToZ - 1) // (xcells * ratio)
    ix = ids - idUpToZ - iy * xcells * ratio - 1

    return ix, iy, iz
