import os.path
import weakref
import warnings
import numpy as np

from . import HaloCatalogue, Halo
from .. import units
from ..array import SimArray


class SubfindCatalogue(HaloCatalogue):

    """Class to handle catalogues produced by the SubFind halo finder.

    By default, the FoF groups are imported, but the subhalos can also be imported by setting subs=True
    when constructing.

    """


    def __init__(self, sim, subs=False, ordered=None, make_grp=None, verbose=False, **kwargs):
        """Initialise a SubFind catalogue

        **kwargs** :

        *subs*: if True, load the subhalos; otherwise, load FoF groups (default)
        *ordered*: if True, the snapshot must have sequential iords. If False, the iords might be out-of-sequence, and
                 therefore reordering will take place. If None (default), the iords are examined to check whether
                 re-ordering is required or not. Note that re-ordering can be undesirable because it also destroys
                 subfind's order-by-binding-energy
        *make_grp*: if True, create a 'grp' array with the halo id to which each particle is assigned
        *verbose*: if True, produces printed output to track progress while making the 'grp' array
        """

        if 'order' in kwargs:
            warnings.warn("The keyword 'order' has been renamed to 'ordered'", DeprecationWarning)
            ordered = kwargs.pop('order')

        self._base = weakref.ref(sim)
        self._subs=subs

        if ordered is not None:
            self._ordered = ordered
        else:
            self._ordered = bool((self.base['iord']==np.arange(len(self.base))).all())

        self._halos = {}
        HaloCatalogue.__init__(self,sim)
        self.dtype_int = sim['iord'].dtype
        self.dtype_flt='float32' #SUBFIND data is apparently always single precision???
        self.halodir = self._name_of_catalogue(sim)
        self.header = self._readheader()
        if subs is True:
            if self.header[6]==0:
                raise ValueError("This file does not contain subhalos")
            if make_grp is True:
                raise ValueError("subs=True and make_grp=True are not compatible.")
        self._tasks = self.header[4]
        self.ids = self._read_ids()
        self._keys={}
        self._halodat, self._subhalodat=self._read_groups()
        #self.data_len, self.data_off = self._read_groups()
        if make_grp:
            self.make_grp(v=verbose)

    def make_grp(self, name='grp', v=False):
        """Creates a 'grp' array which labels each particle according to its parent halo.

        This can take quite some time! Option: v=True prints out 'progress' in terms of total number of groups.
        """
        self.base[name] = self.get_group_array(v=v)

    def get_group_array(self, v=False):
        ar = np.zeros(len(self.base), dtype=int)-1
        for i in range(0, self.header[1]): #total number of groups
            if v:
                print("Halo #", i , "of", self.header[1])
            halo=self[i]
            ar[halo.get_index_list(self.base)] = halo._halo_id
        return ar

    def get_halo_properties(self, i, with_unit=True):
        if with_unit:
            extract = units.get_item_with_unit
        else:
            extract = lambda array, element: array[element]
        properties = {}
        if self._subs is False:
            for key in self._keys:
                properties[key] = extract(self._halodat[key], i)
            if self.header[6] > 0:
                properties['children'] = np.where(self._subhalodat['sub_groupNr'] == i)[
                    0]  # this is the FIRST level of substructure, sub-subhalos (etc) can be accessed via the subs=True output (below)
        else:
            for key in self._keys:
                properties[key] = extract(self._subhalodat[key], i)
            properties['children'] = np.where(self._subhalodat['sub_parent'] == i)[
                0]  # this goes down one level in the hierarchy, i.e. a subhalo will have all its sub-subhalos listed, but not its sub-sub-subhalos (those will be listed in each sub-subhalo)
        return properties

    def _get_halo(self,i):
        if self._ordered is False:
            if self._subs is True:
                #this needs to be tested again on a snapshot that is not ordered!
                x = Halo(i, self, self.base, np.where(np.in1d(self.base['iord'], self.ids[self._subhalodat['sub_off'][i]:self._subhalodat['sub_off'][i]+self._subhalodat['sub_len'][i]]  )))
            else:
                x = Halo(i, self, self.base, np.where(np.in1d(self.base['iord'], self.ids[self._halodat['group_off'][i]:self._halodat['group_off'][i]+self._halodat['group_len'][i]]  )))

        else:
            if self._subs is False: #to use groups as halos:
                x = Halo(i, self, self.base, self.ids[self._halodat['group_off'][i]:self._halodat['group_off'][i]+self._halodat['group_len'][i]] )
            else:
                x=Halo(i, self, self.base, self.ids[self._subhalodat['sub_off'][i]:self._subhalodat['sub_off'][i]+self._subhalodat['sub_len'][i]] )

        x._descriptor = "halo_"+str(i)
        x.properties.update(self.get_halo_properties(i))
        return x

    def _readheader(self):
        header = np.array([], dtype='int32')
        filename = self.halodir + "/subhalo_tab_" + \
            self.halodir.split("_")[-1] + ".0"
        fd = open(filename, "rb")
        # read header: this is strange but it works: there is an extra value in
        # header which we delete in the next step
        header1 = np.fromfile(fd, dtype='int32', sep="", count=8)
        header = np.delete(header1, 4)
        fd.close()
        return header  # [4]

    def _read_ids(self):
        data_ids = np.array([], dtype=self.dtype_int)
        for n in range(0, self._tasks):
            filename = self.halodir + "/subhalo_ids_" + \
                self.halodir.split("_")[-1] + "." + str(n)
            fd = open(filename, "rb")
            # for some reason there is an extra value in header which we delete
            # in the next step
            header1 = np.fromfile(fd, dtype='int32', sep="", count=7)
            header = np.delete(header1, 4)
            # TODO: include a check if both headers agree (they better)
            ids = np.fromfile(fd, dtype=self.dtype_int, sep="", count=-1)
            fd.close()
            data_ids = np.append(data_ids, ids)
        return data_ids

    def _read_groups(self):
        halodat={}
        keys_flt=['mass', 'pos', 'mmean_200', 'rmean_200', 'mcrit_200', 'rcrit_200', 'mtop_200', 'rtop_200', 'contmass']
        keys_int=['group_len', 'group_off',  'first_sub', 'Nsubs', 'cont_count', 'mostboundID']
        for key in keys_flt:
            halodat[key]=np.array([], dtype=self.dtype_flt)
        for key in keys_int:
            halodat[key]=np.array([], dtype='int32')

        subhalodat={}
        subkeys_int=['sub_len', 'sub_off', 'sub_parent', 'sub_mostboundID', 'sub_groupNr']
        subkeys_flt=['sub_pos', 'sub_vel', 'sub_CM', 'sub_mass', 'sub_spin', 'sub_veldisp', 'sub_VMax', 'sub_VMaxRad', 'sub_HalfMassRad', ]
        for key in subkeys_int:
            subhalodat[key]=np.array([], dtype='int32')
        subhalodat['sub_mostboundID']=np.array([], dtype=self.dtype_int)
        #subhalodat['sub_groupNr']=np.array([], dtype=self.dtype_int) #these are special
        for key in subkeys_flt:
            subhalodat[key]=np.array([], dtype=self.dtype_flt)

        self._keys=keys_flt+keys_int
        if self._subs is True:
            self._keys=subkeys_flt+subkeys_int

        for n in range(0,self._tasks):
            filename=self.halodir+"/subhalo_tab_"+self.halodir.split("_")[-1]+"." +str(n)
            fd=open(filename, "rb")
            header1=np.fromfile(fd, dtype='int32', sep="", count=8)
            header=np.delete(header1,4)
            #read groups
            if header[0]>0:
                halodat['group_len']=np.append(halodat['group_len'], np.fromfile(fd, dtype='int32', sep="", count=header[0]))
                halodat['group_off']=np.append(halodat['group_off'], np.fromfile(fd, dtype='int32', sep="", count=header[0]))
                halodat['mass']=np.append(halodat['mass'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['pos']=np.append(halodat['pos'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=3*header[0]) )
                halodat['mmean_200']=np.append(halodat['mmean_200'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['rmean_200']=np.append(halodat['rmean_200'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['mcrit_200']=np.append(halodat['mcrit_200'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['rcrit_200']=np.append(halodat['rcrit_200'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['mtop_200']=np.append(halodat['mtop_200'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['rtop_200']=np.append(halodat['rtop_200'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['cont_count']=np.append(halodat['cont_count'], np.fromfile(fd, dtype='int32', sep="", count=header[0]))
                halodat['contmass']=np.append(halodat['contmass'],np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[0]))
                halodat['Nsubs']=np.append(halodat['Nsubs'],np.fromfile(fd, dtype='int32', sep="", count=header[0]))
                halodat['first_sub']=np.append(halodat['first_sub'],np.fromfile(fd, dtype='int32', sep="", count=header[0]))
            #read subhalos only if expected to exist from header
            if header[5]>0:
                subhalodat['sub_len']=np.append(subhalodat['sub_len'], np.fromfile(fd, dtype='int32', sep="", count=header[5]))
                subhalodat['sub_off']=np.append(subhalodat['sub_off'], np.fromfile(fd, dtype='int32', sep="", count=header[5]))
                subhalodat['sub_parent']=np.append(subhalodat['sub_parent'], np.fromfile(fd, dtype='int32', sep="", count=header[5]))
                subhalodat['sub_mass']=np.append(subhalodat['sub_mass'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[5]))
                subhalodat['sub_pos']=np.append(subhalodat['sub_pos'],np.fromfile(fd, dtype=self.dtype_flt, sep="", count=3*header[5]))
                subhalodat['sub_vel']=np.append(subhalodat['sub_vel'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=3*header[5]))
                subhalodat['sub_CM']=np.append(subhalodat['sub_CM'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=3*header[5]))
                subhalodat['sub_spin']=np.append(subhalodat['sub_spin'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=3*header[5]))
                subhalodat['sub_veldisp']=np.append(subhalodat['sub_veldisp'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[5]))
                subhalodat['sub_VMax']=np.append(subhalodat['sub_VMax'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[5]))
                subhalodat['sub_VMaxRad']=np.append(subhalodat['sub_VMaxRad'],np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[5]))
                subhalodat['sub_HalfMassRad']=np.append(subhalodat['sub_HalfMassRad'], np.fromfile(fd, dtype=self.dtype_flt, sep="", count=header[5]))
                subhalodat['sub_mostboundID']=np.append(subhalodat['sub_mostboundID'], np.fromfile(fd, dtype=self.dtype_int, sep="", count=header[5]))
                subhalodat['sub_groupNr']=np.append(subhalodat['sub_groupNr'], np.fromfile(fd, dtype='int32', sep="", count=header[5]))
            fd.close()

        halodat['pos']=np.reshape(halodat['pos'], (header[1],3))
        if header[6]>0:
            #some voodoo because some SubFind files may have (at least?) one extra entry which is not really a subhalo
            real_ones=np.where(halodat['first_sub']<header[6])[0]
            fake_ones=np.where(halodat['first_sub']>=header[6])[0]
            halodat['mostboundID']=np.zeros(len(halodat['Nsubs']),dtype=self.dtype_int)-1
            halodat['mostboundID'][real_ones]=subhalodat['sub_mostboundID'][halodat['first_sub'][real_ones]]  #useful for the case of unordered snapshot IDs

            subhalodat['sub_pos']=np.reshape(subhalodat['sub_pos'], (header[6],3))
            subhalodat['sub_vel']=np.reshape(subhalodat['sub_vel'], (header[6],3))
            subhalodat['sub_CM']=np.reshape(subhalodat['sub_CM'], (header[6],3))
            subhalodat['sub_spin']=np.reshape(subhalodat['sub_spin'], (header[6],3))

        ar_names = 'mass', 'pos', 'mmean_200', 'rmean_200', 'mcrit_200', 'rcrit_200', 'mtop_200', 'rtop_200', \
                   'sub_mass', 'sub_pos', 'sub_vel', 'sub_CM', 'sub_veldisp', 'sub_VMax', 'sub_VMaxRad', 'sub_HalfMassRad'
        ar_dimensions = 'kg', 'm', 'kg', 'm', 'kg', 'm', 'kg', 'm', \
                    'kg', 'm', 'm s^-1', 'm', 'm s^-1', 'm s^-1', 'm', 'm'

        for name, dimension in zip(ar_names, ar_dimensions):
            if name in halodat:
                halodat[name] = SimArray(halodat[name], self.base.infer_original_units(dimension))
            if name in subhalodat:
                subhalodat[name] = SimArray(subhalodat[name], self.base.infer_original_units(dimension))

        return halodat, subhalodat

    def __len__(self):
        if self._subs:
            return len(self._subhalodat['sub_pos'])
        else:
            return len(self._halodat['pos'])

    @staticmethod
    def _name_of_catalogue(sim):
        # standard path for multiple snapshot files
        snapnum = os.path.basename(
            os.path.dirname(sim.filename)).split("_")[-1]
        parent_dir = os.path.dirname(os.path.dirname(sim.filename))
        dir_path=os.path.join(parent_dir,"groups_" + snapnum)

        if os.path.exists(dir_path):
            return dir_path
        # alternative path if snapshot is single file
        else:
            snapnum = os.path.basename(sim.filename).split("_")[-1]
            parent_dir = os.path.dirname(sim.filename)
            return os.path.join(parent_dir,"groups_" + snapnum)

    @property
    def base(self):
        return self._base()

    @staticmethod
    def _can_load(sim, **kwargs):
        if os.path.exists(SubfindCatalogue._name_of_catalogue(sim)):
            return True
        else:
            return False