import glob
import os
import pickle
import sys

import numpy as np

import pynbody
import pynbody.analysis.profile as profile
import pynbody.filt as filt
import pynbody.plot as pp
import pynbody.plot.sph
import pynbody.units as units

simname = sys.argv[1]   # use first argument as simulation name to analyze
pp.plt.ion()            # plot in interactive mode to show all plots as
                        #     they are made

s = pynbody.load(simname) # load file into pynbody
h = s.halos()             # find and load halos (using AHF if in path)
diskf = filt.Disc('40 kpc','2 kpc')  # filter for particles in disk
notdiskf = filt.Not(filt.Disc('40 kpc','3 kpc')) # particles outside of disk
i=1
if (len(sys.argv) > 2):
    # if there's a second argument, it can be a photogenic file
    # and analysis will follow the halo that contains those particles
    photiords = np.genfromtxt(sys.argv[2],dtype='i8')
    frac = float(len(np.where(np.in1d(photiords,h[i]['iord']))[0]))/len(photiords)
    print('i: %d frac: %.2f'%(i,frac))
    while(((frac) < 0.5) & (i<100)):
        i=i+1
        frac = float(len(np.where(np.in1d(photiords,h[i]['iord']))[0]))/len(photiords)
        print('i: %d frac: %.2f'%(i,frac))
else:
    # otherwise follow largest halo with at least 2 stars
    while len(h[i].star) <2: i=i+1

if (i==100): sys.exit()
pynbody.analysis.angmom.faceon(h[i])  # align halo faceon
s.physical_units()   # change to physical units
