"""
Implements loading of nchilada files.

"""

import os
import struct
import warnings
import xml.dom.minidom

import numpy as np

from .. import chunk, family, units
from . import SimSnap, namemapper

_name_map, _rev_name_map = namemapper.setup_name_maps('nchilada-name-mapping')
_translate_array_name = namemapper.name_map_function(_name_map, _rev_name_map)

_type_codes = list(map(np.dtype, [None, 'int8', 'uint8', 'int16', 'uint16',
                             'int32', 'uint32', 'int64', 'uint64',
                             'float32', 'float64']))


_max_buf = 1024 * 512


class NchiladaSnap(SimSnap):
    """Implements loading of nchilada files."""

    def _load_header(self, f):
        # Used to use xdrlib, but that is deprecated in Python 3.11
        # The following code just follows the old xdrlib implementation
        buf= f.read(28)
        pos= 0
        # Read int, walrus assignment updates pos
        assert struct.unpack('>l',buf[pos:(pos:=pos+4)])[0] == 1062053
        # Read double, walrus assignment updates pos
        t= struct.unpack('>d',buf[pos:(pos:=pos+8)])[0]
        # Read four ints, walrus assignment updates pos
        highword, nbod, ndim, code = (
            struct.unpack('>l',buf[pos:(pos:=pos+4)])[0] for i in range(4)
        )
        return t, nbod, ndim, _type_codes[code]

    def _update_loadable_keys(self):
        self._loadable_keys_registry = {}
        d = self._loadable_keys_registry
        fams = self._dom_sim.getElementsByTagName('family')
        for f in sorted(fams, key=lambda x: str(x)):
            fam_name = f.attributes['name'].value
            our_fam = family.get_family(fam_name)
            d_f = {}
            for a in f.getElementsByTagName('attribute'):
                our_name = _translate_array_name(
                    a.attributes['name'].value, reverse=True)
                filename = os.path.join(
                    self._filename, a.attributes['link'].value)
                d_f[our_name] = filename
            d[our_fam] = d_f

    def _setup_slices(self, take=None):
        disk_family_slice = {}
        i = 0
        # for each family, find an array (any array) to determine length and
        # set up a logical map of particles on disk
        for f in sorted(self._loadable_keys_registry.keys()):
            ars = self._loadable_keys_registry[f]
            with open(list(ars.values())[0], 'rb') as tf:
                header_time, nbod, _, _ = self._load_header(tf)
            disk_family_slice[f] = slice(i, i + nbod)
            i += nbod
            self.properties['a'] = header_time
        self._load_control = chunk.LoadControl(
            disk_family_slice, _max_buf, take)
        self._family_slice = self._load_control.mem_family_slice
        self._num_particles = self._load_control.mem_num_particles

    def __init__(self, filename, **kwargs):
        """Load a nchilada file.

        Parameters
        ----------

        filename : str
            The path to the nchilada output to load. Nchilada outputs are directories containing a description.xml
            file and a number of binary files.
        take : np.ndarray, optional
            The array of particles to load. If not specified, all particles are loaded.
        """

        super().__init__()

        must_have_paramfile = kwargs.get('must_have_paramfile', False)
        take = kwargs.get('take', None)

        self._dom_sim = xml.dom.minidom.parse(
            os.path.join(filename, "description.xml")).getElementsByTagName('simulation')[0]
        self._filename = filename
        self._update_loadable_keys()
        self._setup_slices(take=take)
        self._paramfilename = kwargs.get('paramfile', None)

        self._decorate()

        if 'dKpcUnit' not in self._paramfile:
            if must_have_paramfile:
                raise RuntimeError("Could not find .param file for this run. Place it in the run's directory or parent directory.")
            else:
                warnings.warn(
                    "No readable param file in the run directory or parent directory: using defaults.", RuntimeWarning)
                self._file_units_system = [
                    units.Unit(x) for x in ('G', '1 kpc', '1e10 Msol')]


    def loadable_keys(self, fam=None):
        if fam is not None:
            return list(self._loadable_keys_registry[fam].keys())
        else:
            loadable = None
            for f in self._loadable_keys_registry.values():
                if loadable is None:
                    loadable = set(f.keys())
                else:
                    loadable = loadable.intersection(iter(f.keys()))
            return list(loadable)

    def _open_file_for_array(self, fam, array_name):
        fname = self._loadable_keys_registry[fam].get(array_name, None)
        if not fname:
            raise OSError("No such array on disk")
        f = open(fname, 'rb')
        return f

    def _attempt_load_all_families(self, array_name):
        fams = self.families()
        universal_dtype = None
        universal_ndim = None
        if fams==[]:
            return
        for fam in fams:
            # this will raise an IOError propagating upwards if any family doesn't have the appropriate array
            with self._open_file_for_array(fam, array_name) as f:
                _, nbod, ndim, dtype = self._load_header(f)
            if universal_dtype is not None:
                if ndim!=universal_ndim:
                    raise OSError("Mismatching dimensions for array")
                if dtype!=universal_dtype:
                    raise OSError("Mismatching data type for array")
            universal_ndim, universal_dtype = ndim, dtype

        self._create_array(array_name,universal_ndim,universal_dtype,False)

        for fam in fams:
            self._load_array(array_name, fam)


    def _load_array(self, array_name, fam=None):
        if fam is None:
            self._attempt_load_all_families(array_name)
            return

        f = self._open_file_for_array(fam, array_name)

        _, nbod, ndim, dtype = self._load_header(f)
        if array_name not in list(self.keys()):
            self._create_family_array(array_name, fam, ndim=ndim,dtype=dtype)
        r = self[fam][array_name]
        if units.has_units(r):
            r.convert_units(self._default_units_for(array_name))
        else:
            r.set_default_units()

        disk_dtype = dtype.newbyteorder('>')

        buf_shape = _max_buf
        if ndim > 1:
            buf_shape = (_max_buf, ndim)

        b = np.empty(buf_shape)

        # skip over min and max values (see issue #211)
        np.fromfile(f, dtype=disk_dtype, count=2 * ndim)

        for readlen, buf_index, mem_index in self._load_control.iterate(fam, fam):
            b = np.fromfile(f, dtype=disk_dtype, count=readlen * ndim)
            if ndim > 1:
                b = b.reshape((readlen, ndim))

            if mem_index is not None:
                r[mem_index] = b[buf_index]
        f.close()

    @classmethod
    def _can_load(cls, f):
        return os.path.isdir(f) and os.path.exists(os.path.join(f, "description.xml"))
