#!/usr/bin/env python

import sys
import array
import struct
import copy

if len(sys.argv) < 2:
    print "\ncull_subhalos: Removes subhalos (not the particles, just the group)"
    print "Usage: cull_subhalos <bound halo file>"
    print "Output: <bound halo file>.culled"
    print "Exiting\n"
    sys.exit(1)

def cull_subgroups(groups):

    indexes = range(0,len(groups))
    indexes.reverse()

    for i in indexes:
        little_group = groups[i]
        for j in range(0,i):
            big_group = groups[j]
            bs = set(big_group)
            ls = set(little_group)
            if bs & ls == ls:
                groups.pop(i)
                break
            
def check_for_duplicates(groups):

    indexes = range(0,len(groups))
    indexes.reverse()

    dups = 0

    for i in indexes:
        little_group = groups[i]
        for j in range(0,i):
            big_group = groups[j]
            bs = set(big_group)
            ls = set(little_group)
            if bs & ls:
                dups += 1

    return dups

def remove_duplicates(groups):

    indexes = range(0,len(groups))
    indexes.reverse()

    for i in indexes:
        little_group = groups[i]
        for j in range(0,i):
            big_group = groups[j]
            bs = set(big_group)
            ls = set(little_group)
            matches = bs & ls
            if matches:
                for m in matches:
                    little_group.remove(m)


fnamevoboz = sys.argv[1]

voboz_fin = open(fnamevoboz,'rb')
ng_voboz = struct.unpack('=i',voboz_fin.read(4))[0]
voboz_groups = []

for i in range(0,ng_voboz):
    nph = struct.unpack('=i',voboz_fin.read(4))[0]
    if nph:
        p = array.array('i')
        p.read(voboz_fin,nph)        
        voboz_groups.append(copy.copy(p))

voboz_fin.close()

#sort from largest to smallest
voboz_groups.sort(key=len,reverse=True)


n_initial = len(voboz_groups)
print "Found %d initial groups." % n_initial
cull_subgroups(voboz_groups)
print "Removed %d subhalos." % (n_initial - len(voboz_groups))

n_dup = check_for_duplicates(voboz_groups)

if n_dup:
    print "Found %d groups sharing particles." % n_dup
    remove_duplicates(voboz_groups)
    n_dup2 = check_for_duplicates(voboz_groups)
    print "Removed shared particles from %d." % (n_dup - n_dup2)

culled_out = open(fnamevoboz+'.culled','wb')

culled_out.write(struct.pack('=i',len(voboz_groups)))

for group in voboz_groups:
    culled_out.write(struct.pack('=i',len(group)))
    group.write(culled_out)

culled_out.close()
