"""Contains routines to read in data in SMT format and write it out in HDF5 format"""

import h5py
import sys
import numpy
import glob
import re

import myglobals
import tree

#---------------------------------------------------------------------------------------------------------------------

def loadmergertree(file):

# Loads in a sussing format merger tree file
# Input:
#   file - file containing merger tree
# Ouput:
#   progenitor_dict - dictionary containing a list of progenitors versus HaloID

    progenitor_dict = {}
    with open(file) as f:
        # first few lines are fixed
        format = f.readline()
        name = f.readline()
        count = f.readline()
        while(f):
            line = f.readline()
            if not line or line.startswith("END"): break
            hid, count = map(int, line.split())
            children = []
            for i in range(0,count):
                children.extend(map(int, f.readline().split()))
            progenitor_dict[hid] = children
    # # Check for duplicates (should not be any)
    # # According to the profiler, this line takes about 1/6 of the time
    # l=sum(progenitor_dict.values(),[])
    # if myglobals.verbosity >=2: print(len(l),'progenitors found')
    # # According to the profiler, this next line takes most (over half) of the time
    # progenitor_duplicates=set([x for x in l if l.count(x) > 1])
    # if len(progenitor_duplicates)==0:
    #     if myglobals.verbosity >=2: print('No duplicate progenitors')
    # else:
    #     print('progenitor_duplicates=',progenitor_duplicates)
    #     raise ValueError('Duplicate progenitors')

    return progenitor_dict

#---------------------------------------------------------------------------------------------------------------------

def loadsnapfiles(place, pattern):

# loads in snapshot halo files - fields are read from python module haloprop.py in the data directory.
# Input:
#   place   - where to store the files in the HDF5 hierarchy
#   pattern - load matching files in a glob: e.g., "datadir/*.haloes"
# Output:
#   mass_dict - dictionary of Mvir (in this case) versus HaloID
#   hostHalo_dict - dictionary of hostHalo versus HaloID

    # import information about the halo fields
    sys.path.append(pattern.rpartition('/')[0])
    import haloprop

    # extract those columns that we want to store
    snaphaloprop=haloprop.prop[['name','description','units']]

    # The following does not work because HDF5 doesn't support numpy.str_ type;
    snaphalopropdtype=numpy.dtype([(x,haloprop.dtype[x]) for x in snaphaloprop.dtype.names])
    # So specify by hand a variable length string type
    snaphalopropdtype=numpy.dtype([(x,h5py.special_dtype(vlen=str)) for x in snaphaloprop.dtype.names])
    # The above is very clever but C/Fortran tables can only handle fixed-length strings
    snaphalopropdtype=numpy.dtype([(x,'S100') for x in snaphaloprop.dtype.names])
    # Nor does the following work as for some reason h5py rejects '<i8' etc even though the documentation says not
    #snaphalodtype=numpy.dtype({'names': [x for x in haloprop.prop['name']], 
    #                           'formats': [numpy.dtype(x).type for x in haloprop.prop['type']]})
    # So use fudge to convert dtype into correct format
    conversion={'i8':numpy.int64,'i4':numpy.int32,'f8':numpy.float64,'f4':numpy.float32}
    types=haloprop.prop['type']
    snaphalodtype = [ (haloprop.prop['name'][i], conversion[types[i]]) for i in range(len(types)) ]

    mass_dict = {}
    hostHalo_dict = {}
    count = 0
    # loop over all matching input files
    for fname in glob.iglob(pattern):
        try:
            # Matches / followed by one or more digits and saves as group1
            if myglobals.test: matchdata = re.search('/(\d+)', fname)
            else: matchdata = re.search('dm_(\d+)', fname)
            if not matchdata: matchdata = re.search('ng_(\d+)', fname)
            # Assuming looks like a snapshot number then generate name Snapshot#
            if not matchdata: sys.exit('No snap number found in '+fname)
            else:
                snapplace = 'Snapshot%d'%int(matchdata.group(1))
            # Create group to hold snapshot data
            splace = place.create_group(snapplace)
            # Load in halo data and write out as a table
            print("fname=",fname)
            snaphalo = numpy.loadtxt(fname, dtype=snaphalodtype)
            # Thanks to Andrew Benson for the following line.  
            # loadtxt returns a single tuple when only one line present -
            # we need to force an array of tuples to get consistent behaviour in h5py
            if not snaphalo.shape: snaphalo = snaphalo.reshape((-1,))
            NSnapHalo = snaphalo['HaloID'].size
            if myglobals.verbosity >=1: print('fname, NSnapHalo =',fname, NSnapHalo)
            if NSnapHalo > 0:
                splace.attrs['NSnapHalo'] = NSnapHalo
                splace.create_dataset('SnapHalo', data=snaphalo,
                                      dtype=snaphalodtype)
                snaphaloprop=numpy.array([tuple(snaphaloprop[i][j] for j in range(3)) 
                                          for i in range(len(snaphaloprop))],dtype=snaphalopropdtype)
                splace.create_dataset('SnapHaloProp', data=snaphaloprop, 
                                      dtype=snaphalopropdtype)
                # Create dictionary of hostHalos and masses (Mvir in this case) for later use 
                ids = snaphalo['HaloID']
                hostHalos = snaphalo['hostHalo']
                masses = snaphalo['Mvir']
                # For some unfathomable reason, arrays of length 1 cannot be enumerated
                if NSnapHalo == 1:
                    hostHalo_dict[int(ids)]=int(hostHalos)
                    mass_dict[int(ids)]=int(masses)
                else: 
                    for ind,hid in enumerate(ids):
                        hostHalo_dict[hid] = hostHalos[ind]
                        mass_dict[hid] = masses[ind]
            else:
                splace.attrs['NSnapHalo'] = 0
            count += 1
        except ValueError as e:
            print('Ignore',fname,e)
    place.attrs['NSnap'] = count
    return hostHalo_dict,mass_dict

#---------------------------------------------------------------------------------------------------------------------

def loadsnaptable(place, fname):

# Load the list of snapshots and populate the snapshot tables
# Input:
#   place   - where to store the files in the HDF5 hierarchy
#   fname - file containing list of snapshots and their properties
# Output:

    # Define the structure of the SnapProp table.  
    # h5py documentation recommends use of S type for strings.
    snappropdtype = numpy.dtype([('name', 'S100'),
                                 ('description', 'S100'),
                                 ('Units', 'S100')])
    # Array creation does not seem to work (in Python 3 at least) unless we use byte strings b'string'
    snapprop = numpy.array([(b'snapnum',b'Snapshot number',b'None'),
                            (b'a',b'Expansion factor',b'None'),
                            (b'z',b'Redshift',b'None'),
                            (b't',b'Age of Universe in units of the current age',b't0'),
                            (b't0',b'Current age of the Universe in years',b'yr')],dtype=snappropdtype)
    # This line actually creates the h5py dataset
    place.create_dataset('SnapProp', data=snapprop, dtype=snappropdtype)
    # Define the structure of the Snap table
    snapdtype = numpy.dtype([('snapnum', numpy.int32),
                             ('a', numpy.float32),
                             ('z', numpy.float32),
                             ('t', numpy.float32),
                             ('t0', numpy.float32)])
    # load the data from a plain ascii file
    sdata = numpy.loadtxt(fname, dtype=snapdtype, usecols=(0,1,2,3,4))
    # Create the h5py dataset
    place.create_dataset('Snap', data=sdata, dtype=snapdtype)

#---------------------------------------------------------------------------------------------------------------------

def id2indx(id, haloids):

# Why don't I create a dictionary?????
# Just create a dictionary of haloids to index

# Finds the location in the list of haloIDs where a particular haloid resides.
# This feels like it could be slow, but could be sped up using the fact that haloids is sorted.
# Input:
#   id - desired haloID
#   haloids - list of haloIDs
# Output:
#   indx - index of id in haloids
    indx = -1
    if id != -1:
        # This next line consumes 1/5 of the run-time
        #indx = numpy.where(haloids == id)[0]
        # if len(indx) == 0: indx  = -1
        # else: indx = indx[0]
        # This searches a sorted array so is presumably quicker (but does not guarantee a match)
        indx = numpy.searchsorted(haloids,id)
        if haloids[indx] != id: indx  = -1
    return indx

#---------------------------------------------------------------------------------------------------------------------

def populatemergertree(mergerTree, haloids, progenitor_dict, mass_dict, hosthalo_dict):

# 

    # # The datatype for the Halo table:
    # mtreedtype = numpy.dtype([('HaloID', numpy.int64),
    #                           ('Snapshot', numpy.int32),
    #                           ('Density', numpy.float32),
    #                           ('Mass', numpy.float32),
    #                           ('FirstSubhaloIndex', numpy.int32),
    #                           ('NeighbourIndex', numpy.int32),
    #                           ('HostHaloIndex', numpy.int32),
    #                           ('FirstProgenitorIndex', numpy.int32),
    #                           ('NextSiblingIndex', numpy.int32),
    #                           ('DescendantIndex', numpy.int32),
    #                           ('LastSubhaloIndex', numpy.int32),
    #                           ('LastProgenitorIndex', numpy.int32),
    #                           ])

    # The overdensity of halos:
    dhalo=180.

    # The last snapshot number
    lastSnap=61
    
    # The number of halos
    nhalos=haloids.size
    if myglobals.verbosity >=1: print('nhalos =',nhalos)

    # Populate halo data
    snapids = numpy.array(haloids / 1e12, dtype=numpy.int32)
    density = numpy.full(nhalos, dhalo, dtype=numpy.float32)
    mass = numpy.array([mass_dict[haloid] for haloid in haloids], dtype=numpy.float32)    

    # # make an array to hold the halo data and populate it
    # halodset = numpy.recarray((nhalos,), dtype=mtreedtype)
    # halodset['HaloID'] = haloids
    # halodset['Snapshot'] = snapids
    # halodset['Density'] = density
    # halodset['Mass'] = mass

    # Now compute the local relative pointers
    # Use "ID" for original halo IDs (non-consecutive) and omit the ID for relative position [0:nhalos]
    firstSubhalo = numpy.full(nhalos, -1, dtype=numpy.int32)
    neighbour = numpy.full(nhalos, -1, dtype=numpy.int32)
    hostHalo = numpy.full(nhalos, -1, dtype=numpy.int32)
    firstProgenitor = numpy.full(nhalos, -1, dtype=numpy.int32)
    nextSibling = numpy.full(nhalos, -1, dtype=numpy.int32)
    descendant = numpy.full(nhalos, -1, dtype=numpy.int32)

    # Create halo pointers
    for halo in range(nhalos):
        if myglobals.verbosity >=3: print('halo:',halo)
        haloID=haloids[halo]
        # Spatial pointers.
        # We have a host halo for each halo.
        host=id2indx(hosthalo_dict[haloID],haloids)
        if host != -1:
            hostHalo[halo]=host
            if firstSubhalo[host] == -1: 
                firstSubhalo[host]=halo
            else: 
                neighbour[halo]=neighbour[firstSubhalo[host]]
                neighbour[firstSubhalo[host]]=halo
        # Temporal pointers.
        # We have a list of progentiors for each halo, with the first being the main progenitor
        lastindex = -1
        # look at each progenitor
        for progID in progenitor_dict[haloID]:
            # convert haloID to position in halids list
            prog=id2indx(progID,haloids)
            # the first id is the child index
            if firstProgenitor[halo] == -1: firstProgenitor[halo] = prog
            # subsequent ones are siblings
            else: nextSibling[lastindex] = prog
            lastindex = prog
            # all progenitors have halo as their descendant
            descendant[prog]=halo

    # # Next we run through the tree structure looking for halos with no descendants.
    # # If they have a host with a descendant, then patch things to allow them to use that instead, 
    # # else place in a list of "barren" branches
    # # Create a list of halos with no descendents
    # nfixed=0
    # nbarren=0 # Number of halos leading to a barren one
    # barren=[]
    # for halo in range(nhalos):
    #     root=halo
    #     while True:
    #         while descendant[root] != -1: root=descendant[root]
    #         # If this root halo is in the last snapshot, then everything is fine.
    #         if snapids[root] == lastSnap: 
    #             break
    #         else:
    #             # If no descendant, attempt to associate a halo with the descendant of its host halo and try again
    #             if descendant[hostHalo[root]] != -1:
    #                 nfixed += 1
    #                 descendant[root]=descendant[hostHalo[root]]
    #             # Else add to list of barren halos and exit
    #             else:
    #                 nbarren +=1
    #                 barren.append(root)
    #                 break
    # # Convert to a set and back to remove duplicates
    # barren=list(set(barren))
    # print('Fixed',nfixed,'halos with no descendants')
    # print('Left with',nbarren,'halos on',len(barren),'barren branches')
    # # Some debugging of trees
    # if myglobals.debug:
    #     print('There are',nbarren,'halos on',len(barren),'barren branches')
    #     print('The minimum halo mass at the end-point of a barren branch is',mass[barren].min())
    #     print('The maximum halo mass at the end-point of a barren branch is',mass[barren].max())
    #     halo=barren[mass[barren].argmax()]
    #     print('Max halo is',halo,haloids[halo],'of mass',mass[halo])
    #     print(len(numpy.where(mass[barren]>3e11)[0]),'barren halos have mass >3e11 Msun')
    #     print(len(numpy.where(mass[barren]>1e11)[0]),'barren halos have mass >1e11 Msun')
    #     print(len(numpy.where(mass[barren]>5e10)[0]),'barren halos have mass >5e10 Msun')
    #     print(len(numpy.where(mass[barren]>3e10)[0]),'barren halos have mass >3e10 Msun')
    #     print(len(numpy.where(mass[barren]>2e10)[0]),'barren halos have mass >2e10 Msun')
    #     print(len(numpy.where(mass[barren]>1.9e10)[0]),'barren halos have mass >1.9e10 Msun')
    #     print(len(numpy.where(mass[barren]>1.8e10)[0]),'barren halos have mass >1.8e10 Msun')

    # Initialise tree nodes...
    # Trick: create an extra non-existent node to receive references to index -1 (= last index)
    nodes=[tree.Tree(halo) for halo in range(nhalos+1)]
    nodes[nhalos]=None
    # ...and populate with data
    for halo in range(nhalos): nodes[halo].load(nodes[firstSubhalo[halo]],
                                                nodes[neighbour[halo]],
                                                nodes[firstProgenitor[halo]],
                                                nodes[nextSibling[halo]])

    # To find the lastProgenitor and lastSubhalo we must use depth-first ordering in both time and space.
    # Subhalos of a halo always have the same snapshot, whereas progenitors of a halo may have any overdensity,
    # so that means that the outer search direction has to be spatial and the inner one temporal.
    # In principle all halos should have decendants, but in practice some do not: we label these "barren" halos
    # and keep track of them separately.
    # Create a list of halos with no descendents
    nbarren=0 # Number of halos leading to a barren one
    barren=[]
    # Create a list of root halos
    roots=[]
    for halo in range(nhalos):
        root=halo
        # First find the last descendant
        while descendant[root] != -1: root=descendant[root]
        if snapids[root] != lastSnap: 
            # If last descendant is not located in the final snapshot, add to list of barren halos.
            nbarren +=1
            barren.append(root)
        else:
            # Locate the main host halo and add to list of roots.
            while hostHalo[root] != -1: root=hostHalo[root]
            roots.append(root)
    # Convert to a set and back to remove duplicates
    barren=list(set(barren))
    roots=list(set(roots))
    if myglobals.debug:
        print('There are',nbarren,'halos on',len(barren),'barren branches')
        print('The maximum halo mass at the end-point of a barren branch is',mass[barren].max())

    # Run over trees
    # Table to hold new values.  Make this one larger than necessary to hold the value -1
    halo_sorted=numpy.full(nhalos+1, -1, numpy.int32)
    halonew=-1
    # First run over the tree proper
    for root in roots:
        if myglobals.verbosity >= 3: print('New tree:')
        for halo in nodes[root].next():
            halonew +=1
            halo_sorted[halo]=halonew
            if myglobals.verbosity >= 3: print(halonew,halo,haloids[halo])
    if myglobals.verbosity >=2: print(halonew,'halos found in main tree search')
    # Next add in the barren branches
    for root in barren:
        if myglobals.verbosity >= 3: print('New barren branch:')
        for halo in nodes[root].next(treeType='temporal'):
            halonew +=1
            halo_sorted[halo]=halonew
            if myglobals.verbosity >= 3: print(halonew,halo,haloids[halo])        
    if myglobals.verbosity >=2: print(halonew,'halos found including barren branches')
    # Check that halonew contains each number exactly once, plus -1
    if myglobals.verbosity >= 3: print('halo_sorted=',halo_sorted)
    if set(halo_sorted[:-1]) == set([i for i in range(nhalos)]): 
        # Alter all pointers to use new positions
        for halo in range(nhalos):
            firstSubhalo[halo]=halo_sorted[firstSubhalo[halo]]
            neighbour[halo]=halo_sorted[neighbour[halo]]
            hostHalo[halo]=halo_sorted[hostHalo[halo]]
            firstProgenitor[halo]=halo_sorted[firstProgenitor[halo]]
            nextSibling[halo]=halo_sorted[nextSibling[halo]]
            descendant[halo]=halo_sorted[descendant[halo]]
        # Now sort the arrays: yes, all of them.
        # We need the inverse of halo_sorted for this
        halo_detros=numpy.full(nhalos+1, -1, numpy.int32)
        for i in range(nhalos): halo_detros[halo_sorted[i]]=i
        if myglobals.verbosity >=3: print('halo_detros=',halo_detros)
        haloids=[haloids[i] for i in halo_detros[:-1]]
        snapids=[snapids[i] for i in halo_detros[:-1]]
        density=[density[i] for i in halo_detros[:-1]]
        mass=[mass[i] for i in halo_detros[:-1]]
        firstSubhalo=[firstSubhalo[i] for i in halo_detros[:-1]]
        neighbour=[neighbour[i] for i in halo_detros[:-1]]
        hostHalo=[hostHalo[i] for i in halo_detros[:-1]]
        firstProgenitor=[firstProgenitor[i] for i in halo_detros[:-1]]
        nextSibling=[nextSibling[i] for i in halo_detros[:-1]]
        descendant=[descendant[i] for i in halo_detros[:-1]]
    else:
        print('set(halo_sorted[:-1]) != set([i for i in range(nhalos)]), so continuing without sorting')
        print('nhalos=',nhalos)
        print('len(set(halo_sorted[:-1])) = ',len(set(halo_sorted[:-1])))
        print('Halos not in sorted list = ',(set([i for i in range(nhalos)])).difference(set(halo_sorted[:-1])))
        #adjacent elements in sorted list
        temp=numpy.sort(halo_sorted)
        print('Duplicate halos:')
        for i in range(nhalos):
            if temp[i]==temp[i+1]: print('   ',halo_sorted[i])

    # Store datasets:
    # HaloID
    dset = mergerTree.create_dataset('HaloID', data=haloids)
    dset.attrs["Description"] = 'Original halo ID'
    dset.attrs["Units"] = 'None'
    # Snapshot
    dset = mergerTree.create_dataset('Snapshot', data=snapids)
    dset.attrs["Description"] = 'Snapshot of this halo'
    dset.attrs["Units"] = 'None'
    # Density
    dset=mergerTree.create_dataset('Density', data=density)
    dset.attrs["Description"] = 'Overdensity of the halo cf critical density'
    dset.attrs["Units"] = 'None'
    # Mass
    dset= mergerTree.create_dataset('Mass', data=mass)
    dset.attrs["Description"] = 'M200c'
    dset.attrs["Units"] = 'Msun/h?'
    # FirstSubhalo
    dset=mergerTree.create_dataset('FirstSubhaloIndex', 
                              data=numpy.array(firstSubhalo, dtype=numpy.int32))
    dset.attrs["Description"]='Index of FirstSubhalo in these arrays'
    dset.attrs["Units"]='None'
    #halodset['FirstSubhaloIndex'] = numpy.array(firstSubhalo)
    # Neighbour
    dset=mergerTree.create_dataset('NeighbourIndex',  
                              data=numpy.array(neighbour, dtype=numpy.int32))
    dset.attrs["Description"]='Index of Neighbour in these arrays'
    dset.attrs["Units"]='None'
    #halodset['NeighbourIndex'] = numpy.array(neighbour)
    # HostHalo
    dset=mergerTree.create_dataset('HostHaloIndex', 
                              data=numpy.array(hostHalo, dtype=numpy.int32))
    dset.attrs["Description"]='Index of HostHalo in these arrays'
    dset.attrs["Units"]='None'
    #halodset['HostHaloIndex'] = numpy.array(hostHalo)
    # FirstProgenitor 
    dset=mergerTree.create_dataset('FirstProgenitorIndex', 
                              data=numpy.array(firstProgenitor, dtype=numpy.int32))
    dset.attrs["Description"]='Index of FirstProgenitor in these arrays'
    dset.attrs["Units"]='None'
    #halodset['FirstProgenitorIndex'] = numpy.array(firstProgenitor, dtype=numpy.int32)
    # NextSibling
    dset=mergerTree.create_dataset('NextSiblingIndex', 
                              data=numpy.array(nextSibling, dtype=numpy.int32))
    dset.attrs["Description"]='Index of NextSibling in these arrays'
    dset.attrs["Units"]='None'
    #halodset['NextSiblingIndex'] = numpy.array(nextSibling, dtype=numpy.int32)
    # Descendant
    dset=mergerTree.create_dataset('DescendantIndex', 
                              data=numpy.array(descendant, dtype=numpy.int32))
    dset.attrs["Description"]='Index of Descendant in these arrays'
    dset.attrs["Units"]='None'
    #halodset['DescendantIndex'] = numpy.array(descendant, dtype=numpy.int32)

    # Write out table
    #halodset = mergerTree.create_dataset('Halo', data=halodset, dtype=mtreedtype)



