"""

nchilada
========

Implements classes and functions for handling nchilada files.  You rarely
need to access this module directly as it will be invoked
automatically via pynbody.load.

"""

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):

    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):
        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()

    """
    def _write_array(self, array_name, fam=None) :
        if fam is None :
            fam = self.families()

        for f in fam :
            fname = self._loadable_keys_registry[fam][array_name]
            # to do: sort out what happens when this doesn't exist
            ar = self[fam][array_name]

            _, nbod, ndim, dtype = self._load_header(f)
            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 :
                    b[buf_index] = ar[mem_index]
                    f.seek(-readlen*ndim*disk_dtype.itemsize,1)
    """

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