import numpy as np
import utilities as util
import os.path
import visualization.image_maker as imaker
import gadget
import sys

## this is just here to call the movie maker 'master' routine
def movie_maker_call():
    snapshot_dir='/home/cchayward/scratch/VoT/sbcn_hr_f'
    output_dir='/home/cchayward/scratch/VoT/sbcn_hr_f/movie_frames'
    cosmo=0.
    size=20.
    type_to_make='star'
    movie_maker(xmax_of_box=size,snapdir=snapshot_dir,outputdir_master=output_dir,\
        show_gasstarxray=type_to_make,cosmological=cosmo,frame_min=0,frames_per_gyr=50.,\
        i_snap_min=57,i_snap_max=72,skip_bh=1,set_fixed_center=[0.,0.,0.],\
        theta_0=90.,phi_0=180.,add_extension='',\
        camera_opening_angle=80.,\
        pixels=500);
    # i_snap_min=53; i_snap_max=66;frames_per_gyr=400.

def test_im():
    smaster='/Users/phopkins/Documents/work/plots/zooms/'
    sdir='m11_hr'
    snum=271
    x00=50.

    out0,out1 = imaker.image_maker(sdir,snum,\
        snapdir_master=smaster,outdir_master='./',\
        cosmo=1,xrange=np.array([-1.,1.])*x00,\
        show_gasstarxray='gas');

##
## here is the main routine for the galaxy simulation movies:
##   it builds the snapshot list, does the centering, goes through frames and calls
##   the appropriate routines for interpolation between snapshots, and then calls the
##   image_maker routine to render the frame: many switches here (most don't need to be set),
##   but tune it as you need to (I recommend a low snapshot spacing to test)
##
def movie_maker(\
    xmax_of_box=50., #"""scale (in code units) of movie box size"""
    snapdir='/n/scratch2/hernquist_lab/phopkins/zoom_test/from_scinet/hires_bh', #"""location of snapshots"""
    outputdir_master='/n/scratch2/hernquist_lab/phopkins/movie_frames/', #"""parent folder for frames dump"""
    show_gasstarxray='stars', #"""determines if image is 'stars','gas','xr' etc"""
    frames_per_gyr=100., #"""sets spacing of frames (even in time)"""
    cosmological=1, #"""is this is a cosmological (comoving) simulation?"""
    time_min=0, #"""initial time of movie"""
    time_max=0, #"""final time (if<=0, just set to time of maximum snapshot)"""
    frame_min=0, #"""initial frame number to process"""
    frame_max=1.0e10, #"""final frame number to process"""
    i_snap_min=0, #"""minimum snapshot number in list to scan"""
    i_snap_max=0, #"""maximum snapshot number to scan (if =0, will just use largest in snapdir)"""
    pixels=720, #"""images will be pixels*pixels"""
    show_time_label=0, #"""do or don't place time label on each frame"""
    show_scale_label=0, #"""do or don't place physical scale label on each frame"""
    temp_max=1.0e6, #"""maximum gas temperature for color-weighted gas images (as 'temp_cuts' in 3-color)"""
    temp_min=3.0e2, #"""minimum gas temperature for color-weighted gas images (as 'temp_cuts' in 3-color)"""
    theta_0=90., #"""(initial) polar angle (if rotating)"""
    phi_0=90., #"""(initial) azimuthal angle (if rotating)"""
    center_on_bh=0, #"""if have a bh particle, uses it for the centering"""
    center_on_com=0, #"""simple centering on center of mass"""
    set_fixed_center=0, #"""option to set a fixed center for the whole movie"""
    sdss_colors=0, #"""use sdss color-scheme for three-color images"""
    nasa_colors=1, #"""use nasa color-scheme for three-color images"""
    use_h0=1, #"""correct snapshots to physical units"""
    four_char=0, #"""snapshots have 4-character filenames"""
    skip_bh=0, #"""skips reading bh information in snapshots"""
    use_old_extinction_routine=0, #"""uses older (faster but less stable) stellar extinction routine"""
    add_extension='', #"""extension for movie snapshot images"""
    do_with_colors=1, #"""make color images"""
    threecolor=1, #"""do or dont use the standard three-color projection for images"""
    camera_opening_angle=45., #"""camera opening angle (specifies camera distance)"""
    scattered_fraction=0.01, #"""fraction of light that is unattenuated"""
    z_to_add=0.0, #"""add this metallicity to all particles"""
    min_stellar_age=0.0 #"""force star particles to be at least this old to avoid interpolation artifacts"""
    ):

    ## first decide whether its a gas or stellar image
    SHOW_STARS=0; SHOW_GAS=1;
    if((show_gasstarxray=='star') or (show_gasstarxray=='stars') or (show_gasstarxray=='st')):
        SHOW_STARS=1; SHOW_GAS=0;

    ## parse the directory names (for file-naming conventions)
    s0=snapdir.split("/"); snapdir_specific=s0[len(s0)-1]; n_s0=1;
    if(len(snapdir_specific)<=1): snapdir_specific=s0[len(s0)-2]; n_s0=2;
    snapdir_master='';
    for j in s0[0:len(s0)+1-n_s0]: snapdir_master += str(j)+'/';
    outputdir_master+='/' ## just to be safe

    ## build the snapshot list in the directory of interest
    print '... building snapshot list in directory ...'
    snapshot_list = build_snapshot_list( snapdir, \
        use_h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )
    if(i_snap_max<=0): i_snap_max=snapshot_list[snapshot_list.size-2];

    ## to avoid 'camera jitter', the safest thing (albeit more expensive) is to
    ##   pre-compute the movie center at all snapshots, then use a smoothed tracking of it
    ##   -- call subroutine to check if its done for us, and if not, build it --
    print '... building/loading camera positions list ...'
    build_camera_centering( snapshot_list, snapdir, force_rebuild=0, \
        center_on_bh=center_on_bh,center_on_com=center_on_com,set_fixed_center=set_fixed_center,\
        use_h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )

    print '... building list of times for movie frames ...'
    ## now build the times for each frame of the movie
    time_frame_grid, a_scale_grid = build_time_frame_grid( snapdir, snapshot_list, \
          frames_per_gyr, time_min=time_min, time_max=time_max, \
          use_h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh)

    ##
    ## some (simple) defaults for the scaling of dynamic range, etc
    ##  (can do fancy things for movie panning/zooming etc by changing the variables here)
    ##
    t=0.0*time_frame_grid
    x_plot_scale_grid = t + xmax_of_box
    y_plot_scale_grid = x_plot_scale_grid
    #theta_frame_grid = t + 90. # should be in degrees
    theta_frame_grid = t + theta_0 # should be in degrees
    #phi_frame_grid = t + 90. ## x-y
    phi_frame_grid = t + phi_0 ## x-y
    dynrange_0 = 3.0e2
    maxden_0 = 3.8e7
    maxden_grid = t + maxden_0 * ((50./x_plot_scale_grid)**(0.3))
    if(SHOW_STARS==1):
        dynrange_0 *= 100.0 #XXX CCH: originally 3.0
        maxden_grid *= 100.0 #XXX CCH: originally 3.0
    maxden_grid /= 1.0e10 # recalling in gadget units, m=1.0e10 m_sun
    # layer below is for zoom_dw large-scale GAS sims ONLY
    #z_grid = 1./a_scale_grid - 1.
    #dynrange_0 *= 20. * (1.+np.exp(-(z_grid-2.)))
    #maxden_grid *= 3.
    dynrange_grid=dynrange_0 + 0.*x_plot_scale_grid

    print '... entering main snapshot loop ...'
    ## now enter main loop over snapshots in the simulation set
    snapshot_f_number_that_i_just_loaded = -1;
    snaps_to_do = (snapshot_list >= i_snap_min) & (snapshot_list <= i_snap_max)
    ii = np.arange(snapshot_list.size)
    snap_grid_to_loop = ii[snaps_to_do]
    for i_snap in snap_grid_to_loop:
        ## load header info for snapshot 'i'
        PPP_head = gadget.readsnap(snapdir,snapshot_list[i_snap],1,header_only=1,\
            h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh)
        if(cosmological==1):
            a_scale_i=PPP_head['time']; t_i=cosmological_time(a_scale_i);
        else:
            t_i=PPP_head['time']; a_scale_i=t_i
        a_scale_i=np.array(a_scale_i); t_i=np.array(t_i);

        ## load header info for snapshot 'f=i+1'
        PPP_head = gadget.readsnap(snapdir,snapshot_list[i_snap+1],1,header_only=1,\
            h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh)
        if(cosmological==1):
            a_scale_f=PPP_head['time']; t_f=cosmological_time(a_scale_f);
        else:
            t_f=PPP_head['time']; a_scale_f=t_f
        a_scale_f=np.array(a_scale_f); t_f=np.array(t_f);

        print '... snapshot headers loaded: getting timestep info: ...'
        ## now calculate whether there are any frames in the time-range between the snapshots
        delta_t = t_f - t_i
        print 'Timestep is ti/tf/delta_t ',t_i,t_f,delta_t
        i_frame = np.array(range(time_frame_grid.size));
        check_frames_to_do = (time_frame_grid >= t_i) & (time_frame_grid <= t_f) & \
            (i_frame <= frame_max) & (i_frame >= frame_min);
        frames_to_do = i_frame[check_frames_to_do];
        n_frames_to_do = frames_to_do.size;

        if(n_frames_to_do>0):
            print '... processing ',n_frames_to_do,' frames in this snapshot interval ...'

            center_i = get_precalc_zoom_center(snapdir,a_scale_i,cosmological=cosmological)
            center_f = get_precalc_zoom_center(snapdir,a_scale_f,cosmological=cosmological)
            ## correct to comoving coordinates for interpolation (so correctly capture hubble flow)
            if (cosmological==1):
                center_i /= a_scale_i
                center_f /= a_scale_f
            print '... ... centered (i/f) at ',center_i,center_f

            print '... ... loading snapshot bricks ... ...'
            ## tolerance for keeping particles outside the box, for this calculation
            xmax_tmp=xmax_of_box*1.5 + 10.
            ## load the actual --data-- for snapshot (i):
            id_i,m_i,x_i,y_i,z_i,vx_i,vy_i,vz_i,h_i,c_i,zm_i = \
                load_pos_vel_etc(snapdir, snapshot_list[i_snap],\
                    h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh, \
                    GAS=SHOW_GAS,BOX_MAX=xmax_tmp,CENTER_BOX=center_i, \
                    min_stellar_age=min_stellar_age);
            if (SHOW_STARS==1):
                id_i_g,m_i_g,x_i_g,y_i_g,z_i_g,vx_i_g,vy_i_g,vz_i_g,h_i_g,c_i_g,zm_i_g = \
                  load_pos_vel_etc(snapdir, snapshot_list[i_snap], \
                    h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh, \
                    GAS=1,SWAP_TEMP_RHO=1,BOX_MAX=xmax_tmp,CENTER_BOX=center_i, \
                    min_stellar_age=min_stellar_age);
            ## correct to comoving coordinates for interpolation (so correctly capture hubble flow)
            ## (put this here, splitting initial and final rescalings so dont do it twice on 'recycling' step)
            if (cosmological==1):
                for vec in [x_i,y_i,z_i,h_i,vx_i,vy_i,vz_i]: vec /= a_scale_i;
                if (SHOW_STARS==1):
                    for vec in [x_i_g,y_i_g,z_i_g,h_i_g,vx_i_g,vy_i_g,vz_i_g]: vec /= a_scale_i;

            ## now load snapshot (f) [should be new, have to load fresh]
            id_f,m_f,x_f,y_f,z_f,vx_f,vy_f,vz_f,h_f,c_f,zm_f = \
                load_pos_vel_etc(snapdir, snapshot_list[i_snap+1], \
                    h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh, \
                    GAS=SHOW_GAS,BOX_MAX=xmax_tmp,CENTER_BOX=center_f);
            if (SHOW_STARS==1):
                id_f_g,m_f_g,x_f_g,y_f_g,z_f_g,vx_f_g,vy_f_g,vz_f_g,h_f_g,c_f_g,zm_f_g = \
                  load_pos_vel_etc(snapdir, snapshot_list[i_snap+1], \
                    h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh, \
                    GAS=1,SWAP_TEMP_RHO=1,BOX_MAX=xmax_tmp,CENTER_BOX=center_f);
            ## correct to comoving coordinates for interpolation (so correctly capture hubble flow)
            if (cosmological==1):
                for vec in [x_f,y_f,z_f,h_f,vx_f,vy_f,vz_f]: vec /= a_scale_f;
                if (SHOW_STARS==1):
                    for vec in [x_f_g,y_f_g,z_f_g,h_f_g,vx_f_g,vy_f_g,vz_f_g]: vec /= a_scale_f;
            snapshot_f_number_that_i_just_loaded = snapshot_list[i_snap+1];

            ## before going further, check that we actually have particles to process,
            ##   and if not skip this frame
            if (m_i.size+m_f.size < 5):
                continue;

            print '... ... matching ids in the snapshots ... ...'
            ## predefine matches to save time in the loop below
            sub_id_i,sub_id_f,nomatch_from_i_in_f,nomatch_from_f_in_i = compile_matched_ids( id_i, id_f );
            if (SHOW_STARS==1):
                sub_id_i_g,sub_id_f_g,nomatch_from_i_in_f_g,nomatch_from_f_in_i_g = compile_matched_ids( id_i_g, id_f_g );

            print '... ... entering frame loop ... ...'
            ## loop over frames in between the two snapshots just pulled up
            for i_frame in range(n_frames_to_do):
                j_of_frame = frames_to_do[i_frame];
                dt = time_frame_grid[j_of_frame] - t_i
                time_of_frame = time_frame_grid[j_of_frame]

                print '... ... ... interpolating for frame ',i_frame+1,'/',n_frames_to_do,'... ... ...'
                ## this is the interpolation step between the two snapshots for each frame
                x_all,y_all,z_all,m_all,h_all,c_all,zm_all = interpolation_for_movie( dt, delta_t, \
                    m_i,h_i,c_i,zm_i, m_f,h_f,c_f,zm_f, \
                    x_i,y_i,z_i,vx_i,vy_i,vz_i, x_f,y_f,z_f,vx_f,vy_f,vz_f, \
                    sub_id_i,sub_id_f, nomatch_from_i_in_f,nomatch_from_f_in_i, \
                    center_i,center_f,use_polar_interpolation=0 )
                if (SHOW_STARS==1):
                    x_all_g,y_all_g,z_all_g,m_all_g,h_all_g,c_all_g,zm_all_g = \
                        interpolation_for_movie( dt, delta_t, \
                        m_i_g,h_i_g,c_i_g,zm_i_g,m_f_g,h_f_g,c_f_g,zm_f_g, \
                        x_i_g,y_i_g,z_i_g,vx_i_g,vy_i_g,vz_i_g, x_f_g,y_f_g,z_f_g,vx_f_g,vy_f_g,vz_f_g, \
                        sub_id_i_g,sub_id_f_g, nomatch_from_i_in_f_g,nomatch_from_f_in_i_g, \
                        center_i,center_f,use_polar_interpolation=0 )
                ## ok, now have the interpolated gas(+stellar) quantities needed for each image

                ## correct back from comoving coordinates for plotting (since units are physical)
                if (cosmological==1):
                    a = a_scale_grid[j_of_frame]
                    time_of_frame = a
                    for vec in [x_all,y_all,z_all,h_all]: vec *= a;
                    if (SHOW_STARS==1):
                        for vec in [x_all_g,y_all_g,z_all_g,h_all_g]: vec *= a;

                print '... ... ... centering and setting passed variables ... ... ...'
                ## re-center on the current frame center, interpolated between the snapshots
                cen = get_precalc_zoom_center(snapdir,a_scale_grid[j_of_frame],cosmological=cosmological)
                print 'Center for frame j=',j_of_frame,' at ',cen
                x_all-=cen[0]; y_all-=cen[1]; z_all-=cen[2];
                if (SHOW_STARS==1):
                    x_all_g-=cen[0]; y_all_g-=cen[1]; z_all_g-=cen[2];

                ## set dynamic ranges of image
                maxden=maxden_grid[j_of_frame]
                dynrange=dynrange_grid[j_of_frame]
                ## set image spatial scale
                xr_0=x_plot_scale_grid[j_of_frame]
                yr_0=y_plot_scale_grid[j_of_frame]
                xr=[-xr_0,xr_0]; yr=[-yr_0,yr_0]; zr_0=max([xr_0,yr_0]); zr=[-zr_0,zr_0]
                ## set viewing angle for image
                theta = theta_frame_grid[j_of_frame]
                phi   = phi_frame_grid[j_of_frame]

                ## some final variable cleaning/prep before imaging routine:
                if (SHOW_STARS==0):
                    gdum=np.array([0]);
                    x_all_g=y_all_g=z_all_g=m_all_g=h_all_g=c_all_g=zm_all_g=gdum;
                else:
                    gdum=np.zeros(m_all_g.size)+1.;

                print '... ... ... sending to main imaging routine ... ... ...'
                print snapdir_specific
                print j_of_frame
                print snapdir_master
                print outputdir_master
                print theta
                print phi
                print dynrange
                print maxden
                print show_gasstarxray
                print add_extension
                print show_time_label
                print show_scale_label
                print use_h0
                print cosmological
                print pixels
                print xr
                print yr
                print zr
                print nasa_colors
                print sdss_colors
                print use_old_extinction_routine
                print do_with_colors
                print threecolor
                print temp_min
                print temp_max
                print np.array([temp_min, temp_max])
                print SHOW_GAS
                print m_all
                print x_all
                print y_all
                print z_all
                print c_all
                print h_all
                print zm_all
                print gdum
                print c_all_g
                print h_all_g
                print zm_all_g
                print m_all_g
                print x_all_g
                print y_all_g
                print z_all_g
                print time_of_frame

                ## alright, now we can actually send this to the imaging routine
                ##   to make the frame image!
                image24, massmap = \
                  imaker.image_maker( snapdir_specific, j_of_frame, \
                    snapdir_master=snapdir_master, outdir_master=outputdir_master,
                    theta=theta, phi=phi, dynrange=dynrange, maxden=maxden,
                    show_gasstarxray=show_gasstarxray, add_extension=add_extension, \
                    show_time_label=show_time_label, show_scale_label=show_scale_label, \
                    use_h0=use_h0, cosmo=cosmological, coordinates_cylindrical=0, \
                    set_percent_maxden=0, set_percent_minden=0, pixels=pixels, \
                    center_on_com=0, center_on_bh=0, center=[0., 0., 0.], \
                    xrange=xr, yrange=yr, zrange=zr,
                    project_to_camera=1, camera_opening_angle=camera_opening_angle, \
                    center_is_camera_position=0, camera_direction=[0.,0.,-1.], \
                    nasa_colors=nasa_colors, sdss_colors=sdss_colors, \
                    use_old_extinction_routine=use_old_extinction_routine, \
                    do_with_colors=do_with_colors, threecolor=threecolor, \
                    log_temp_wt=1, set_temp_min=temp_min, set_temp_max=temp_max,
                    gas_map_temperature_cuts=np.array([temp_min, temp_max]),
                    input_data_is_sent_directly=1, include_lighting=SHOW_GAS, \
                    m_all=m_all,x_all=x_all,y_all=y_all,z_all=z_all,c_all=c_all,h_all=h_all,zm_all=zm_all+z_to_add,\
                    gas_u=gdum,gas_rho=c_all_g,gas_numh=gdum,gas_nume=gdum,gas_hsml=h_all_g,gas_metallicity=zm_all_g+z_to_add,\
                    gas_mass=m_all_g,gas_x=x_all_g,gas_y=y_all_g,gas_z=z_all_g,time=time_of_frame, \
                    scattered_fraction=scattered_fraction,min_stellar_age=min_stellar_age );

                ## don't need to do anything else here, the image & datafile are dumped!
                print '... ... ... frame ',j_of_frame,'complete! ... ... ...'

    return 1; ## success!

def ok_scan(input,xmax=1.0e10,pos=0):
    if (pos==1):
        return (np.isnan(input)==False) & (np.abs(input)<=xmax) & (input > 0.);
    else:
        return (np.isnan(input)==False) & (np.abs(input)<=xmax);

def idmatch(id_i, id_f):
    ## match items in set (i) to those in set (f)
    ##  --- this is particularly speedy when the ids are all unique (case here) --
    index_dict_i = dict((k,i) for i,k in enumerate(id_i));
    index_dict_f = dict((k,i) for i,k in enumerate(id_f));
    inter = set(id_i).intersection(set(id_f));
    indices_i = np.array([ index_dict_i[x] for x in inter ]);
    indices_f = np.array([ index_dict_f[x] for x in inter ]);
    return indices_i, indices_f;

def compile_matched_ids( id_i, id_f ):
    sub_id_i, sub_id_f = idmatch(id_i,id_f)

    nomatch_from_i_in_f = (id_i > -1.0e40) ## should return all true
    if (sub_id_i.size > 0):
        nomatch_from_i_in_f[sub_id_i] = False ## is matched

    nomatch_from_f_in_i = (id_f > -1.0e40)
    if (sub_id_f.size > 0):
        nomatch_from_f_in_i[sub_id_f] = False

    return sub_id_i, sub_id_f, nomatch_from_i_in_f, nomatch_from_f_in_i

def interpolation_for_movie_pos( dt, delta_t, u_i,u_f, vu_i,vu_f, \
    sub_id_i,sub_id_f, nomatch_from_i_in_f,nomatch_from_f_in_i , periodic=0, order=3 ):

    dt=np.array(dt); delta_t=np.array(delta_t); ## make sure casting is ok
    tau = dt/delta_t;
    Nmatched=u_i[sub_id_i].size;
    Nunmatched_i=u_i[nomatch_from_i_in_f].size;
    Nunmatched_f=u_f[nomatch_from_f_in_i].size;
    u=np.zeros(Nmatched+Nunmatched_i+Nunmatched_f);

    # first compute for objects in (i) with a match in (f)
    if (Nmatched>0):
        xi = np.copy(u_i[sub_id_i])  ; xf = np.copy(u_f[sub_id_f])  ;
        vi = np.copy(vu_i[sub_id_i]) ; vf = np.copy(vu_f[sub_id_f]) ;
        if (periodic==1):
            ## u is periodic in 2pi, so estimate its final position and wrap coordinate
            ## guess where its going to be to get 'best estimate' of the number of revolutions
            #x_exp = xi + 0.5*(vf+vi)*delta_t # actually can get some errors when vf/i switch signs...
            x_exp = xi + vi*delta_t
            ## now pin the final location to the appropriate value
            xf = get_closest_periodic_value( x_exp, xf , xi)

        if (order==3):
            x2 = 3.*(xf-xi) - (2.*vi+vf)*delta_t
            x3 = -2.*(xf-xi) + (vi+vf)*delta_t
            ## third-order interpolation: enables exact matching of x, v, but can 'overshoot'
            u[0:Nmatched] = xi + vi*dt + x2*tau*tau + x3*tau*tau*tau
        elif (order==2):
            ## second-order: minimizes absolute velocity difference (more stable, no exact velocity matching)
            x2 = (vf-vi)*delta_t/2.
            x1 = (xf-xi) - x2
            u[0:Nmatched] = xi + x1*tau + x2*tau*tau
        else:
            ## use linear interpolation below if for some reason this is unstable: least accurate, but most stable
            u[0:Nmatched] = u_i[sub_id_i] + (u_f[sub_id_f]-u_i[sub_id_i])*tau

    ## now unmatched: first those in (i) with no match in (f)
    if(Nunmatched_i>0):
        u[Nmatched:Nmatched+Nunmatched_i] = \
            u_i[nomatch_from_i_in_f]+vu_i[nomatch_from_i_in_f]*dt;
	## now unmatched: now those in (f) with no match in (i)
    if(Nunmatched_f>0):
        u[Nmatched+Nunmatched_i:Nmatched+Nunmatched_i+Nunmatched_f] = \
            u_f[nomatch_from_f_in_i]+vu_f[nomatch_from_f_in_i]*(dt-delta_t);

    return u;

def interpolation_for_movie_jhat( dt, delta_t, j_i,j_f, \
    sub_id_i,sub_id_f, nomatch_from_i_in_f,nomatch_from_f_in_i ):

    dt=np.array(dt); delta_t=np.array(delta_t); ## make sure casting is ok
    tau = dt/delta_t;
    u_i = np.zeros(j_i[:,0].size)
    u_f = np.zeros(j_f[:,0].size)
    Nmatched=u_i[sub_id_i].size;
    Nunmatched_i=u_i[nomatch_from_i_in_f].size;
    Nunmatched_f=u_f[nomatch_from_f_in_i].size;
    j_hat=np.zeros((Nmatched+Nunmatched_i+Nunmatched_f,3));

    # first compute for objects in (i) with a match in (f)
    if (Nmatched>0):
        ## initialize matrices
        pt_i = np.zeros(Nmatched); pt_f = 0.*pt_i;
        p_i = np.zeros((Nmatched,3)); p_f = 0.*p_i; v_f = 0.*p_f
        for j in [0,1,2]:
            p_i[:,j] = j_i[sub_id_i,j]
            p_f[:,j] = j_f[sub_id_f,j]
            pt_i += p_i[:,j]*p_i[:,j]
            pt_f += p_f[:,j]*p_f[:,j]
        ## normalize vectors to unity
        pt_i = np.sqrt(pt_i); pt_f = np.sqrt(pt_f);
        for j in [0,1,2]:
            p_i[:,j] /= pt_i
            p_f[:,j] /= pt_f

        ## build a rotation matrix between the j_hat values at two different timesteps:
        cos_rot_ang = p_i[:,0]*p_f[:,0] + p_i[:,1]*p_f[:,1] + p_i[:,2]*p_f[:,2]
        angle = np.arccos( cos_rot_ang )

        angle_to_rotate = angle * tau
        cos_ang = np.cos(angle_to_rotate)
        sin_ang = np.sin(angle_to_rotate)

        ## build perpendicular vectors and normalize
        k_rot = cross( p_i, p_f )
        kk=0.*pt_i;
        for j in [0,1,2]: kk += k_rot[:,j]*k_rot[:,j]
        kk = np.sqrt(kk)
        for j in [0,1,2]: k_rot[:,j] /= kk

        k_cross_p = cross( k_rot, p_i )
        kk=0.*pt_i;
        for j in [0,1,2]: kk += k_cross_p[:,j]*k_cross_p[:,j]
        kk = np.sqrt(kk)
        for j in [0,1,2]: k_cross_p[:,j] /= kk

        k_dot_p = k_rot[:,0]*p_i[:,0] + k_rot[:,1]*p_i[:,1] + k_rot[:,2]*p_i[:,2]
        for j in [0,1,2]:
            v_f[:,j] = p_i[:,j]*cos_ang + k_cross_p[:,j]*sin_ang + k_rot[:,j]*k_dot_p*(1.-cos_ang)

        for j in [0,1,2]:
            j_hat[0:Nmatched,j] = v_f[:,j]  ## now have the j_hat vector for this time

    ## now unmatched: first those in (i) with no match in (f)
    if(Nunmatched_i>0):
        for j in [0,1,2]:
            j_hat[Nmatched:Nmatched+Nunmatched_i,j] = j_i[nomatch_from_i_in_f,j]

	## now unmatched: now those in (f) with no match in (i)
    if(Nunmatched_f>0):
        for j in [0,1,2]:
            j_hat[Nmatched+Nunmatched_i:Nmatched+Nunmatched_i+Nunmatched_f] = j_f[nomatch_from_f_in_i]

    return j_hat;

def interpolation_for_movie_scalar(dt, delta_t, u_i,u_f, \
    sub_id_i,sub_id_f, nomatch_from_i_in_f,nomatch_from_f_in_i, \
    mass_rescale=0, hsml_rescale=0 ):

    dt=np.array(dt); delta_t=np.array(delta_t); ## make sure casting is ok
    tau = dt/delta_t;
    Nmatched=u_i[sub_id_i].size;
    Nunmatched_i=u_i[nomatch_from_i_in_f].size;
    Nunmatched_f=u_f[nomatch_from_f_in_i].size;
    u=np.zeros(Nmatched+Nunmatched_i+Nunmatched_f);

    # first compute for objects in (i) with a match in (f)
    if (Nmatched>0):
        u[0:Nmatched] = u_i[sub_id_i] + (u_f[sub_id_f]-u_i[sub_id_i])*tau
    ## now unmatched: first those in (i) with no match in (f)
    if(Nunmatched_i>0):
        if(mass_rescale==1):
            u_t = u_i[nomatch_from_i_in_f] * (1.-tau**(1./5.))**5.
        elif(hsml_rescale==1):
            u_t = u_i[nomatch_from_i_in_f] / ((1.-tau**(1./3.))**3.+1.0e-5)
        else:
            u_t = u_i[nomatch_from_i_in_f]
        u[Nmatched:Nmatched+Nunmatched_i] = u_t;
	## now unmatched: now those in (f) with no match in (i)
    if(Nunmatched_f>0):
        if(mass_rescale==1):
            u_t = u_f[nomatch_from_f_in_i] * (1.-(1.-tau)**(1./5.))**5.
        elif(hsml_rescale==1):
            u_t = u_f[nomatch_from_f_in_i] / ((1.-(1.-tau)**(1./3.))**3.+1.0e-5)
        else:
            u_t = u_f[nomatch_from_f_in_i]
        u[Nmatched+Nunmatched_i:Nmatched+Nunmatched_i+Nunmatched_f] = u_t;

    return u;

def cosmological_time(a,h=0.71,Omega_M=0.27):
    ## exact solution for a flat universe
    x=Omega_M/(1.-Omega_M) / (a*a*a);
    t=(2./(3.*np.sqrt(1.-Omega_M))) * np.log( np.sqrt(x) / (-1. + np.sqrt(1.+x)) );
    t *= 13.777 * (0.71/h); ## in Gyr
    return t;

def cross(x,y):
    #return np.cross(x,y,axis=1)
    c=0.*x
    c[:,0] = x[:,1]*y[:,2] - x[:,2]*y[:,1]
    c[:,1] =-x[:,0]*y[:,2] + x[:,2]*y[:,0]
    c[:,2] = x[:,0]*y[:,1] - x[:,1]*y[:,0]
    return c

def xyz_to_polar(x,y,z,vx,vy,vz):
    for q in [x,y,z,vx,vy,vz]: q=np.array(q,dtype='d')+1.0e-10
    x += 1.0e-5; vy += 1.0e-5;

    ## get the angular momentum vector (to use for interpolation)
    r = np.sqrt( x*x + y*y + z*z )
    pos = np.zeros([x.size,3],dtype='d');
    pos[:,0]=x/r; pos[:,1]=y/r; pos[:,2]=z/r;
    v = np.sqrt( vx*vx + vy*vy + vz*vz )
    vel = np.zeros([vx.size,3],dtype='d');
    vel[:,0]=vx/v; vel[:,1]=vy/v; vel[:,2]=vz/v;

    j_hat = cross( pos, vel )
    jj = np.sqrt( j_hat[:,0]*j_hat[:,0] + j_hat[:,1]*j_hat[:,1] + j_hat[:,2]*j_hat[:,2] )
    for j in [0,1,2]: j_hat[:,j] /= jj

    ## get spherical polar coordinates
    v_r = vx*pos[:,0] + vy*pos[:,1] + vz*pos[:,2]

    ## now get the vector perpendicular to both j and r
    phi_hat = cross( j_hat, pos )
    jj = np.sqrt( phi_hat[:,0]*phi_hat[:,0] + phi_hat[:,1]*phi_hat[:,1] + phi_hat[:,2]*phi_hat[:,2] )
    for j in [0,1,2]: phi_hat[:,j] /= jj

    v_phi = vx*phi_hat[:,0] + vy*phi_hat[:,1] + vz*phi_hat[:,2]
    v_phi /= r

    ## define the absolute 'phi' relative to an arbitrary (but fixed) 90-degree rotation
    ##   of the angular momentum vector j
    x_jhat = 0.*j_hat
    x_jhat[:,0]=0.*j_hat[:,0]; x_jhat[:,1]=-j_hat[:,2]; x_jhat[:,2]=j_hat[:,1]
    jj = np.sqrt( x_jhat[:,0]*x_jhat[:,0] + x_jhat[:,1]*x_jhat[:,1] + x_jhat[:,2]*x_jhat[:,2] )
    for j in [0,1,2]: x_jhat[:,j] /= jj

    ## generate y-vector by cross-product a x b = c (gaurantees right-hand rule)
    y_jhat = cross( j_hat, x_jhat )
    jj = np.sqrt( y_jhat[:,0]*y_jhat[:,0] + y_jhat[:,1]*y_jhat[:,1] + y_jhat[:,2]*y_jhat[:,2] )
    for j in [0,1,2]: y_jhat[:,j] /= jj

    ## now project r onto this, to obtain the components in this plane
    x_for_phi = x_jhat[:,0]*pos[:,0] + x_jhat[:,1]*pos[:,1] + x_jhat[:,2]*pos[:,2]
    y_for_phi = y_jhat[:,0]*pos[:,0] + y_jhat[:,1]*pos[:,1] + y_jhat[:,2]*pos[:,2]
    phi = np.arctan2( y_for_phi , x_for_phi )
    ## reset to a 0-2pi system instead of numpy's -pi,pi system
    lo = (phi < 0.); phi[lo] += 2.*np.pi

    return r, v_r, phi, v_phi, j_hat

## simple function to wrap x/y/z coordinates near the box edges, to give
##   the 'correct' x-x0 in periodic coordinates
def periodic_dx( x, x0, box ):
    b0 = box / 2.
    dx = x - x0
    too_large = (dx > b0)
    dx[too_large] -= box
    too_small = (dx < -b0)
    dx[too_small] += box
    return dx

def get_closest_periodic_value( x_expected, x_final, x_initial, limiter=1.0 ):
    x_remainder = np.mod(x_expected, 2.*np.pi)
    dx = periodic_dx( x_final, x_remainder, 2.*np.pi ) ## returns wrapped dx from x_remainder
    dx_full = x_expected + dx - x_initial

    return dx_full + x_initial

    # now put a limiter so we don't allow some crazy number of 'laps':
    sign_dx = 0.*dx_full + 1.
    sign_dx[dx_full<0.]=-1.
    dx_full_abs = (sign_dx*dx_full)
    dx_full_cycles = dx_full_abs / (2.*np.pi) # number of 'laps'
    dx_full_remainder = np.mod(dx_full_cycles, 1.0)

    dx_corrected = np.copy(dx_full)
    lminusone = limiter - 1.
    hi = (dx_full_cycles > limiter) & (dx_full_remainder > lminusone)
    dx_corrected[hi] = 2.*np.pi*sign_dx[hi]*dx_full_remainder[hi]
    hi = (dx_full_cycles > limiter) & (dx_full_remainder < lminusone)
    dx_corrected[hi] = 2.*np.pi*sign_dx[hi]*(1.+dx_full_remainder[hi])

    return dx_corrected + x_initial

def interpolation_for_movie(dt, delta_t, \
	    m_i,h_i,c_i,zm_i,\
	    m_f,h_f,c_f,zm_f, \
	    x_in,y_in,z_in,vx_in,vy_in,vz_in, \
	    x_fn,y_fn,z_fn,vx_fn,vy_fn,vz_fn, \
	    s_i,s_f,nm_i_f,nm_f_i, \
	    cen_i,cen_f , use_polar_interpolation=1 ):

    x_i = np.copy(x_in) - cen_i[0]
    y_i = np.copy(y_in) - cen_i[1]
    z_i = np.copy(z_in) - cen_i[2]
    vx_i = np.copy(vx_in) - (cen_f[0]-cen_i[0])/delta_t
    vy_i = np.copy(vy_in) - (cen_f[1]-cen_i[1])/delta_t
    vz_i = np.copy(vz_in) - (cen_f[2]-cen_i[2])/delta_t
    x_f = np.copy(x_fn) - cen_f[0]
    y_f = np.copy(y_fn) - cen_f[1]
    z_f = np.copy(z_fn) - cen_f[2]
    vx_f = np.copy(vx_fn) - (cen_f[0]-cen_i[0])/delta_t
    vy_f = np.copy(vy_fn) - (cen_f[1]-cen_i[1])/delta_t
    vz_f = np.copy(vz_fn) - (cen_f[2]-cen_i[2])/delta_t

    #for ui,uin,j in zip([x_i,y_i,z_i],[x_in,y_in,z_in],[0,1,2]): ui = uin - cen_i[j]
    #for uf,ufn,j in zip([x_f,y_f,z_f],[x_fn,y_fn,z_fn],[0,1,2]): uf = ufn - cen_f[j]
    #for vui,vuin,j in zip([vx_i,vy_i,vz_i],[vx_in,vy_in,vz_in],[0,1,2]): vui = vuin - (cen_f[j]-cen_i[j])/delta_t
    #for vuf,vufn,j in zip([vx_f,vy_f,vz_f],[vx_fn,vy_fn,vz_fn],[0,1,2]): vuf = vufn - (cen_f[j]-cen_i[j])/delta_t

    if (use_polar_interpolation==1):
        r_i,vr_i,phi_i,vphi_i,j_hat_i = xyz_to_polar(x_i,y_i,z_i,vx_i,vy_i,vz_i)
        r_f,vr_f,phi_f,vphi_f,j_hat_f = xyz_to_polar(x_f,y_f,z_f,vx_f,vy_f,vz_f)

        # r gets treated like a normal spatial coordinate (except cannot go <0; ignore for now?)
        r = interpolation_for_movie_pos(dt,delta_t, r_i,r_f,vr_i,vr_f, s_i,s_f,nm_i_f,nm_f_i,order=1)
        # predict the 'final' phi, theta, knowing that these are actually periodic:
        # note for phi, need third-order: inexact matching can leave 'wedges' missing from orbits
        phi = interpolation_for_movie_pos(dt,delta_t, phi_i,phi_f,vphi_i,vphi_f, s_i,s_f,nm_i_f,nm_f_i, periodic=1,order=3)
        x_t = r * np.cos(phi);  y_t = r * np.sin(phi);

        j_hat = interpolation_for_movie_jhat(dt,delta_t, j_hat_i,j_hat_f, s_i,s_f,nm_i_f,nm_f_i )
        ## define the absolute 'phi' relative to an arbitrary (but fixed) 90-degree rotation of the angular momentum vector j
        x_jhat = 0.*j_hat
        x_jhat[:,0]=0.*j_hat[:,0]; x_jhat[:,1]=-j_hat[:,2]; x_jhat[:,2]=j_hat[:,1]
        jj = np.sqrt( x_jhat[:,0]*x_jhat[:,0] + x_jhat[:,1]*x_jhat[:,1] + x_jhat[:,2]*x_jhat[:,2] )
        for j in [0,1,2]: x_jhat[:,j] /= jj
        ## generate y-vector by cross-product a x b = c (gaurantees right-hand rule)
        y_jhat = cross( j_hat, x_jhat )
        jj = np.sqrt( y_jhat[:,0]*y_jhat[:,0] + y_jhat[:,1]*y_jhat[:,1] + y_jhat[:,2]*y_jhat[:,2] )
        for j in [0,1,2]: y_jhat[:,j] /= jj

        x = x_t*x_jhat[:,0] + y_t*y_jhat[:,0]
        y = x_t*x_jhat[:,1] + y_t*y_jhat[:,1]
        z = x_t*x_jhat[:,2] + y_t*y_jhat[:,2]

        ## use cartesian interpolation for some points
        x_c=interpolation_for_movie_pos(dt,delta_t, x_i,x_f,vx_i,vx_f, s_i,s_f,nm_i_f,nm_f_i,order=1)
        y_c=interpolation_for_movie_pos(dt,delta_t, y_i,y_f,vy_i,vy_f, s_i,s_f,nm_i_f,nm_f_i,order=1)
        z_c=interpolation_for_movie_pos(dt,delta_t, z_i,z_f,vz_i,vz_f, s_i,s_f,nm_i_f,nm_f_i,order=1)

        ## combine velocities into initial and final matched vectors
        Nmatched=x_i[s_i].size;
        Nunmatched_i=x_i[nm_i_f].size;
        Nunmatched_f=x_f[nm_f_i].size;
        u=np.zeros(Nmatched+Nunmatched_i+Nunmatched_f);
        vr=0.*u; vphi=0.*u; rr=0.*u; rvphi_i=r_i*vphi_i; rvphi_f=r_f*vphi_f;
        if (Nmatched>0):
            vr[0:Nmatched] = np.sqrt((vr_i[s_i]**2.+vr_f[s_f]**2.)/2.)
            vphi[0:Nmatched] = np.sqrt((rvphi_i[s_i]**2.+rvphi_f[s_f]**2.)/2.)
            rr[0:Nmatched] = np.sqrt((r_i[s_i]**2.+r_f[s_f]**2.)/2.)
        if(Nunmatched_i>0):
            vr[Nmatched:Nmatched+Nunmatched_i] = vr_i[nm_i_f]
            vphi[Nmatched:Nmatched+Nunmatched_i] = rvphi_i[nm_i_f]
            rr[Nmatched:Nmatched+Nunmatched_i] = r_i[nm_i_f]
        if(Nunmatched_f>0):
            vr[Nmatched+Nunmatched_i:Nmatched+Nunmatched_i+Nunmatched_f] = vr_f[nm_f_i]
            vphi[Nmatched+Nunmatched_i:Nmatched+Nunmatched_i+Nunmatched_f] = rvphi_f[nm_f_i]
            rr[Nmatched+Nunmatched_i:Nmatched+Nunmatched_i+Nunmatched_f] = r_f[nm_f_i]
        ## attempt to guess what's 'rotation supported' and not, and use appropriate interpolation
        vr2 = vr*vr
        v2 = vr2 + vphi*vphi
        #no_rot_support = ( vr2 > v2/3. ) | ( rr > 5. ) # usually does ok, too much 'breathing' at late times
        no_rot_support = ( vr2 > v2/2. ) | ( rr > 30. )
        x[no_rot_support] = x_c[no_rot_support]
        y[no_rot_support] = y_c[no_rot_support]
        z[no_rot_support] = z_c[no_rot_support]

    else:
        x=interpolation_for_movie_pos(dt,delta_t, x_i,x_f,vx_i,vx_f, s_i,s_f,nm_i_f,nm_f_i)
        y=interpolation_for_movie_pos(dt,delta_t, y_i,y_f,vy_i,vy_f, s_i,s_f,nm_i_f,nm_f_i)
        z=interpolation_for_movie_pos(dt,delta_t, z_i,z_f,vz_i,vz_f, s_i,s_f,nm_i_f,nm_f_i)

    x = x + (cen_f[0]-cen_i[0]) * dt / delta_t + cen_i[0]
    y = y + (cen_f[1]-cen_i[1]) * dt / delta_t + cen_i[1]
    z = z + (cen_f[2]-cen_i[2]) * dt / delta_t + cen_i[2]

    m=interpolation_for_movie_scalar(dt,delta_t, m_i,m_f, s_i,s_f,nm_i_f,nm_f_i, mass_rescale=1)
    h=interpolation_for_movie_scalar(dt,delta_t, h_i,h_f, s_i,s_f,nm_i_f,nm_f_i, hsml_rescale=1)
    c=interpolation_for_movie_scalar(dt,delta_t, c_i,c_f, s_i,s_f,nm_i_f,nm_f_i)
    #c=np.log10(interpolation_for_movie_scalar(dt,delta_t, 10.**c_i,10.**c_f, s_i,s_f,nm_i_f,nm_f_i))
    zm=interpolation_for_movie_scalar(dt,delta_t, zm_i,zm_f, s_i,s_f,nm_i_f,nm_f_i)

    return x,y,z,m,h,c,zm;

## here's the grunt work of loading the data we'll need
def load_pos_vel_etc( snapdir, snapnum, \
        h0=1, four_char=0, cosmological=0, skip_bh=0, \
        use_rundir=1, #""" make sure to change this as appropriate for the system!"""
        GAS=0, SWAP_TEMP_RHO=0, BOX_MAX=0, CENTER_BOX=[0.,0.,0.], min_stellar_age=0. ):

    have=0; have_h_stars=0;
    ptypes=[0] ## just gas
    if(GAS==0):
        if (cosmological==1):
            ptypes=[4] ## just new stars
        else:
            ptypes=[2,3,4] ## bulge+disk+new stars

    for ptype in ptypes:
        ppp=gadget.readsnap(snapdir,snapnum,ptype,h0=h0,cosmological=cosmological,skip_bh=skip_bh);
        ppp_head=gadget.readsnap(snapdir,snapnum,ptype,header_only=1,h0=h0,cosmological=cosmological,skip_bh=skip_bh);
        if(ppp['k']==1):
            n=ppp['m'].size;
            if(n>1):
                m=ppp['m']; id=ppp['id'];
                p=ppp['p']; x=p[:,0]; y=p[:,1]; z=p[:,2];
                v=ppp['v']; vx=v[:,0]; vy=v[:,1]; vz=v[:,2];
                if (ptype==0):
                    h=ppp['h']; zm=ppp['z'];
                    if(len(zm.shape)>1): zm=zm[:,0]
                    if(SWAP_TEMP_RHO==1):
                        cc=ppp['rho']
                        cc_temp=gadget.gas_temperature(ppp['u'],ppp['ne']);
                        zm[cc_temp>1.0e6]=1.0e-6; ## don't allow hot gas to have dust
                    else:
                        cc=gadget.gas_temperature(ppp['u'],ppp['ne']);
                        #cc[cc<50.]=50.; ## purely to prevent very low values making the image bad
                if (ptype==4):
                    cc=gadget.get_stellar_ages(ppp,ppp_head,cosmological=cosmological);
                    cc=np.maximum(cc,min_stellar_age)
                    zm=ppp['z'];
                    if(len(zm.shape)>1): zm=zm[:,0]
                if (ptype==2): ## need to assign ages and metallicities
                    cc = np.maximum(min_stellar_age , np.zeros(n)+2.0+ppp_head['time'])
                    zm = np.zeros(n) + 0.01;
                if (ptype==3): ## need to assign ages and metallicities
                    cc = np.maximum(min_stellar_age , np.zeros(n)+6.0+ppp_head['time'])
                    zm = np.zeros(n) + 0.002;
                hstars_should_concat=0;
                if((ptype>0) & (have_h_stars==0)):
                    h=gadget.load_allstars_hsml(snapdir,snapnum,cosmo=cosmological, \
                        use_rundir=use_rundir,four_char=four_char,use_h0=h0);
                    have_h_stars=1;
                    hstars_should_concat=1;
                if (have==1):
                    m_all=np.concatenate((m_all,m)); c_all=np.concatenate((c_all,cc)); zm_all=np.concatenate((zm_all,zm));
                    x_all=np.concatenate((x_all,x)); y_all=np.concatenate((y_all,y)); z_all=np.concatenate((z_all,z));
                    vx_all=np.concatenate((vx_all,vx)); vy_all=np.concatenate((vy_all,vy)); vz_all=np.concatenate((vz_all,vz));
                    id_all=np.concatenate((id_all,id));
                    if(hstars_should_concat==1): h_all=np.concatenate((h_all,h));
                else:
                    m_all=m; x_all=x; y_all=y; z_all=z; zm_all=zm; c_all=cc; h_all=h;
                    vx_all=vx; vy_all=vy; vz_all=vz; id_all=id
                    have=1;

    dum = np.zeros(1)
    if (have != 1): return dum,dum,dum,dum,dum,dum,dum,dum,dum,dum,dum;
    if (m_all.size <= 1): return dum,dum,dum,dum,dum,dum,dum,dum,dum,dum,dum;

    if (BOX_MAX != 0):
        x=x_all-CENTER_BOX[0]; y=y_all-CENTER_BOX[1]; z=z_all-CENTER_BOX[2];
        ## make sure everything is cast correctly
        x=np.array(x,dtype='f'); y=np.array(y,dtype='f'); z=np.array(z,dtype='f')
        h_all=np.array(h_all,dtype='f'); c_all=np.array(c_all,dtype='f');
        m_all=np.array(m_all,dtype='f'); zm_all=np.array(zm_all,dtype='f')
        ok=ok_scan(c_all,pos=1) & ok_scan(m_all,pos=1) & ok_scan(h_all,pos=1) & ok_scan(zm_all,pos=1) & \
            ok_scan(x,xmax=BOX_MAX) & ok_scan(y,xmax=BOX_MAX) & ok_scan(z,xmax=BOX_MAX) & \
            (h_all < 100.) & (m_all > 1.0e-12);
        for vec in [id_all, m_all, x_all, y_all, z_all, vx_all, vy_all, vz_all, h_all, c_all, zm_all]:
            vec = vec[ok];

    h_all *= 1.25
    return id_all, m_all, x_all, y_all, z_all, vx_all, vy_all, vz_all, h_all, c_all, zm_all;

def build_snapshot_list( snapdir , use_h0=1,four_char=0,cosmological=0,skip_bh=0):
    i=0; imax=450; snums=[-1];
    while (i < imax):
        PPP_head = gadget.readsnap( snapdir,i,1,header_only=1,\
            h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )
        if(PPP_head['k']==0):
            snums=np.concatenate((snums,np.array([i])));
            if (i > imax-50): imax += 50;
        i += 1;
    snums=snums[snums>=0];
    return snums

def get_camera_centering_filename( snapdir ):
    ## (first parse the directory names (for file-naming conventions))
    s0=snapdir.split("/"); snapdir_specific=s0[len(s0)-1]; n_s0=1;
    if(len(snapdir_specific)<=1): snapdir_specific=s0[len(s0)-2]; n_s0=2;
    ext = '.mv_cen'
    fname = snapdir+'/'+snapdir_specific+ext
    return fname

def build_camera_centering( snums, snapdir ,
    center_on_bh = 0, # center on first BH particle (good for nuclear sims)
    center_on_com = 0, # center on center of mass (good for merger sims)
    set_fixed_center = 0, # option to set a fixed center
    force_rebuild = 0, # force the routine to build from scratch (don't check for file)
    use_h0=1,four_char=0,cosmological=0,skip_bh=0):

    ## first check if the tabulated centers file already exists:
    fname = get_camera_centering_filename( snapdir )
    if(force_rebuild==0):
        if(os.path.exists(fname)): return 1

    ## ok, if we're here, center-table doesn't exist, so let's build it:
    n_snums = np.array(snums).size
    time_all = np.zeros( (n_snums) )
    cen = np.zeros( (n_snums,3) )
    outfi_cen = open(fname,'a')
    set_fixed_center = np.array(set_fixed_center)
    center_on_zero = 0
    if (set_fixed_center.size==1):
        if (set_fixed_center==1): center_on_zero=1
    for i in range(n_snums):
        print 'centering for snap ',snums[i],' in ',snapdir
        PPP_head = gadget.readsnap( snapdir, snums[i], 0, header_only=1, \
            h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )
        time_all[i] = PPP_head['time']
        cen[i,:] = [0.,0.,0.]

        if (center_on_zero==1):
            continue
        if (set_fixed_center.size==3):
            cen[i,:]=set_fixed_center
            continue

        if (center_on_bh==1):
            PPP = gadget.readsnap( snapdir, snums[i], 5,\
                h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )
            if (PPP['k']==1):
                if (PPP['m'].size > 0):
                    cen[i,:] = PPP['p'][0,:]
            continue

        if (center_on_com==1):
            ptype = 1
            PPP = gadget.readsnap( snapdir, snums[i], ptype,\
                h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )
            if (PPP['k']==1):
                if (PPP['m'].size > 0):
                    wt = np.array(P['m']); wt /= np.sum(wt); p = np.array(P['p'])
                    for j in [0,1,2]: cen[i,j]=np.sum(p[:,j]*wt)
            continue

        ## if not using any of the methods above, use our fancier iterative
        ##   centering solution for the 'dominant' structure
        cen[i,:] = gadget.calculate_zoom_center( snapdir, snums[i], \
            cen=[0.,0.,0.], clip_size=2.e10, rho_cut=1.0e-5, \
            h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh)
        print '..zoom_cen at ',cen[i,:]
        sys.stdout.flush()

    ## all done, write out these results
    for i in range(len(cen[:,0])):
        v = np.array([ time_all[i], cen[i,0], cen[i,1], cen[i,2] ]).reshape(-1,4)
        np.savetxt(outfi_cen,v)
    outfi_cen.close()
    return 1

def get_precalc_zoom_center( snapdir, time_desired, cosmological=0 ):
    fname = get_camera_centering_filename( snapdir )
    cenfile = open(fname,'r'); brick=np.loadtxt(cenfile).reshape(-1,4)
    time=brick[:,0]; x0=brick[:,1]; y0=brick[:,2]; z0=brick[:,3]; ## define columns
    if (cosmological==1):
        for q in [x0,y0,z0]: q /= time ## correct to comoving for spline

    nsmoo = 13
    windowfun = 'flat' ## boxcar smoothing
    #nsmoo = 21
    #windowfun = 'hanning' ## gaussian-like smoothing (should use larger window function)
    cen = np.zeros(3)
    for j,p in zip([0,1,2],[x0,y0,z0]):
        pp = util.smooth(p, window_len=nsmoo, window=windowfun)
        #pp = p ## linear interpolation (in case we don't have enough points) #DEBUG!
        if (cosmological==1):
            pp_interp = time_desired * np.interp( np.log(time_desired), np.log(time), pp )
        else:
            pp_interp = np.interp( time_desired, time, pp )
        cen[j] = pp_interp

    print '... ... ... getting center for frame at t=',time_desired,' : ',cen
    return cen

# builds the list of times to dump frames
def build_time_frame_grid( snapdir, snapshot_list, frames_per_gyr, \
        time_min=0, time_max=0, use_h0=1, four_char=0, cosmological=0, skip_bh=1 ):

    ## set times for frames
    if(time_min==0):
        PPP_head = gadget.readsnap( snapdir, snapshot_list[0], 1, header_only=1, \
            h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )
        time_min = PPP_head['time'] ## set to first time in snapshot list
    if(time_max==0):
        PPP_head = gadget.readsnap( snapdir, snapshot_list[-1], 1, header_only=1, \
            h0=use_h0,four_char=four_char,cosmological=cosmological,skip_bh=skip_bh )
        time_max = PPP_head['time'] ## set to last time in snapshot list
    if(cosmological==1):
        time_min = cosmological_time(time_min) ## input in 'a_scale', not time (Gyr)
        time_max = cosmological_time(time_max)

    time_frame_grid = np.arange(time_min, time_max, 1./np.float(frames_per_gyr))

    if (cosmological==0):
        a_scale_grid = time_frame_grid ## to avoid error below
    else:
        a_tmp = 10.**np.arange(-3.,0.001,0.001)
        t_tmp = cosmological_time(a_tmp)
        a_scale_grid = np.exp(np.interp(np.log(time_frame_grid),np.log(t_tmp),np.log(a_tmp)))
        a_scale_grid[a_scale_grid < a_tmp[0]] = a_tmp[0]

    return time_frame_grid, a_scale_grid


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from torch.autograd import Variable
import numpy as np

from models.gcn import GCNLayer

from models.bilinear_scorer import BilinearScorer

from models.custom_allennlp.stacked_alternating_lstm import StackedAlternatingLstm
from models.custom_allennlp.elmo import Elmo

class SRL(nn.Module):
    def __init__(
        self,
        hidden_dim,
        tagset_size,
        num_layers,
        w_c_vocab_size,
        c_c_vocab_size,
        use_syntax,
        eln,
        num_layers_top,
        use_elmo,
        use_bert,
        params,
        gpu_id=-1,
    ):
        super(SRL, self).__init__()
        if gpu_id > -1:
            self.use_gpu = True
        else:
            self.use_gpu = False
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.num_layers_top = num_layers_top
        self.eln = eln
        self.use_elmo = use_elmo
        self.use_bert = use_bert
        self.params = params
        self.dropout = nn.Dropout(p=params.gcn_dropout)
        self.embedding_dropout = nn.Dropout(p=params.emb_dropout)

        if self.use_elmo:
            weight_file = "https://s3-us-west-2.amazonaws.com/allennlp/models/elmo/2x4096_512_2048cnn_2xhighway_5.5B/elmo_2x4096_512_2048cnn_2xhighway_5.5B_weights.hdf5"
            options_file = "https://s3-us-west-2.amazonaws.com/allennlp/models/elmo/2x4096_512_2048cnn_2xhighway_5.5B/elmo_2x4096_512_2048cnn_2xhighway_5.5B_options.json"
            fixed_dim = 1024

            self.elmo = Elmo(
                options_file, weight_file, 1, dropout=0, do_layer_norm=False
            )

            fixed_dim += 100

        elif self.use_bert:
            fixed_dim = 768

        else:
            fixed_dim = 100

        embedding_dim = self.params.emb_dim
        self.indicator_embeddings = nn.Embedding(2, embedding_dim)

        self.tagset_size = tagset_size
        self.use_syntax = use_syntax
        self.num_layers_top = num_layers_top

        gcn_type = GCNLayer

        if self.params.non_linearity == "relu":
            self.non_linearity = nn.ReLU()
        elif self.params.non_linearity == "tanh":
            self.non_linearity = nn.Tanh()
        elif self.params.non_linearity == "leakyrelu":
            self.non_linearity = nn.LeakyReLU()
        elif self.params.non_linearity == "celu":
            self.non_linearity = nn.CELU()
        elif self.params.non_linearity == "selu":
            self.non_linearity = nn.SELU()
        else:
            raise NotImplementedError

        self.lstm = StackedAlternatingLstm(
            fixed_dim + embedding_dim,
            hidden_dim,
            num_layers=num_layers,
            recurrent_dropout_probability=0.1,
        )

        if self.use_syntax:
            if num_layers_top > 0:
                self.lstm_top = StackedAlternatingLstm(
                    hidden_dim,
                    hidden_dim,
                    num_layers=num_layers_top,
                    recurrent_dropout_probability=0.1,
                )
                self.hidden2predicate = nn.Linear(hidden_dim, hidden_dim)
                self.hidden2argument = nn.Linear(hidden_dim, hidden_dim)
                self.bilinear_scorer = BilinearScorer(
                    hidden_dim, tagset_size, params.bilinear_dropout
                )

            else:
                self.hidden2predicate = nn.Linear(hidden_dim, hidden_dim)
                self.hidden2argument = nn.Linear(hidden_dim, hidden_dim)
                self.bilinear_scorer = BilinearScorer(
                    hidden_dim, tagset_size, params.bilinear_dropout
                )

            self.gcn_w_c = gcn_type(
                hidden_dim,
                hidden_dim,
                w_c_vocab_size,
                in_arcs=True,
                out_arcs=True,
                use_gates=True,
                batch_first=True,
                residual=True,
                no_loop=True,
                dropout=self.params.gcn_dropout,
                non_linearity=self.non_linearity,
                edge_dropout=self.params.edge_dropout,
            )

            self.gcn_c_w = gcn_type(
                hidden_dim,
                hidden_dim,
                w_c_vocab_size,
                in_arcs=True,
                out_arcs=True,
                use_gates=True,
                batch_first=True,
                residual=True,
                no_loop=True,
                dropout=self.params.gcn_dropout,
                non_linearity=self.non_linearity,
                edge_dropout=self.params.edge_dropout,
            )

            self.gcn_c_c = gcn_type(
                hidden_dim,
                hidden_dim,
                c_c_vocab_size,
                in_arcs=True,
                out_arcs=True,
                use_gates=True,
                batch_first=True,
                residual=True,
                no_loop=False,
                dropout=self.params.gcn_dropout,
                non_linearity=self.non_linearity,
                edge_dropout=self.params.edge_dropout,
            )

        else:
            self.hidden2predicate = nn.Linear(hidden_dim, hidden_dim)
            self.hidden2argument = nn.Linear(hidden_dim, hidden_dim)
            self.bilinear_scorer = BilinearScorer(
                hidden_dim, tagset_size, params.bilinear_dropout
            )

        if self.eln:
            self.layernorm = nn.LayerNorm(fixed_dim)

    def forward(
        self,
        sentence,
        predicate_flags,
        sent_mask,
        lengths,
        fixed_embs,
        constituents,
        GCN_w_c,
        GCN_c_w,
        GCN_c_c,
        mask_const_batch,
        predicate_index,
        elmo_character_ids,
        bert_embs,
    ):

        if self.use_elmo:
            embeds = self.elmo(elmo_character_ids)["elmo_representations"][0]
            if not self.params.elmo_proj:
                embeds = torch.cat([embeds, fixed_embs], dim=2)

        elif self.use_bert:
            embeds = bert_embs
        else:
            embeds = fixed_embs

        if self.eln:
            embeds = self.layernorm(embeds * sent_mask.unsqueeze(2))

        embeds = self.embedding_dropout(embeds)

        embeds = torch.cat(
            (embeds, self.indicator_embeddings(predicate_flags.long())), 2
        )

        b, t, e = embeds.data.shape

        sent_len = torch.sort(lengths, descending=True)[0]
        idx_sort = torch.argsort(-lengths)

        if self.use_gpu:
            embeds = embeds.index_select(
                0, Variable(torch.cuda.LongTensor(idx_sort.cuda()))
            )
        else:
            embeds = embeds.index_select(0, Variable(torch.LongTensor(idx_sort)))

        packed = pack_padded_sequence(embeds, sent_len, batch_first=True)
        lstm_out, _ = self.lstm(packed)

        lstm_out, _ = pad_packed_sequence(lstm_out, batch_first=True)  # [b, t, h]

        # Un-sort by length
        idx_unsort = torch.argsort(idx_sort)
        if self.use_gpu:
            lstm_out = lstm_out.index_select(
                0, Variable(torch.cuda.LongTensor(idx_unsort.cuda()))
            )
        else:
            lstm_out = lstm_out.index_select(0, Variable(torch.LongTensor(idx_unsort)))

        if self.use_syntax:

            # Here I must concatenate the constituents with the lstm_out
            gcn_in = torch.cat([lstm_out, constituents], dim=1)
            mask_all = torch.cat([sent_mask, mask_const_batch], dim=1)
            # Apply graph conv
            adj_arc_in_w_c, adj_arc_out_w_c, adj_lab_in_w_c, adj_lab_out_w_c, mask_in_w_c, mask_out_w_c, mask_loop_w_c = (
                GCN_w_c
            )

            adj_arc_in_c_w, adj_arc_out_c_w, adj_lab_in_c_w, adj_lab_out_c_w, mask_in_c_w, mask_out_c_w, mask_loop_c_w = (
                GCN_c_w
            )

            adj_arc_in_c_c, adj_arc_out_c_c, adj_lab_in_c_c, adj_lab_out_c_c, mask_in_c_c, mask_out_c_c, mask_loop_c_c = (
                GCN_c_c
            )

            gcn_out = self.gcn_w_c(
                gcn_in,
                adj_arc_in_w_c,
                adj_arc_out_w_c,
                adj_lab_in_w_c,
                adj_lab_out_w_c,
                mask_in_w_c,
                mask_out_w_c,
                mask_loop_w_c,
                mask_all,
            )

            gcn_out = self.gcn_c_c(
                gcn_out,
                adj_arc_in_c_c,
                adj_arc_out_c_c,
                adj_lab_in_c_c,
                adj_lab_out_c_c,
                mask_in_c_c,
                mask_out_c_c,
                mask_loop_c_c,
                mask_all,
            )

            gcn_out = self.gcn_c_w(
                gcn_out,
                adj_arc_in_c_w,
                adj_arc_out_c_w,
                adj_lab_in_c_w,
                adj_lab_out_c_w,
                mask_in_c_w,
                mask_out_c_w,
                mask_loop_c_w,
                mask_all,
            )

            # Take back the lstm out only
            lstm_out = gcn_out.narrow(1, 0, t)
            if self.num_layers_top > 0:
                if self.use_gpu:
                    lstm_out = lstm_out.index_select(
                        0, Variable(torch.cuda.LongTensor(idx_sort.cuda()))
                    )
                else:
                    lstm_out = lstm_out.index_select(
                        0, Variable(torch.LongTensor(idx_sort))
                    )

                packed = pack_padded_sequence(lstm_out, sent_len, batch_first=True)
                lstm_out_, _ = self.lstm_top(packed)

                lstm_out_, _ = pad_packed_sequence(
                    lstm_out_, batch_first=True
                )  # [b, t, h]

                # Un-sort by length
                if self.use_gpu:
                    lstm_out_ = lstm_out_.index_select(
                        0, Variable(torch.cuda.LongTensor(idx_unsort.cuda()))
                    )
                else:
                    lstm_out_ = lstm_out_.index_select(
                        0, Variable(torch.LongTensor(idx_unsort))
                    )
                lstm_out = lstm_out_
        lstm_out_view = lstm_out.contiguous().view(b * t, -1)
        predicate_index = predicate_index.view(b * t)

        predicates_repr = lstm_out_view.index_select(0, predicate_index).view(b, t, -1)

        pred_repr = self.non_linearity(
            self.hidden2predicate(self.dropout(predicates_repr))
        )
        arg_repr = self.non_linearity(self.hidden2argument(self.dropout(lstm_out)))
        tag_scores = self.bilinear_scorer(pred_repr, arg_repr)  # [b*t, label_size]

        return tag_scores.view(b, t, self.tagset_size)

class SRL_Framenet(nn.Module):
    def __init__(
        self,
        hidden_dim,
        tagset_size,
        num_layers,
        w_c_vocab_size,
        c_c_vocab_size,
        use_syntax,
        eln,
        num_layers_top,
        params,
        gpu_id=-1,
    ):
        super(SRL_Framenet, self).__init__()

        if gpu_id > -1:
            self.use_gpu = True
        else:
            self.use_gpu = False
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.num_layers_top = num_layers_top
        self.eln = eln
        self.params = params
        self.dropout = nn.Dropout(p=params.gcn_dropout)

        fixed_dim = 100

        embedding_dim = self.params.emb_dim
        self.indicator_embeddings = nn.Embedding(2, embedding_dim)

        self.tagset_size = tagset_size
        self.use_syntax = use_syntax
        self.num_layers_top = num_layers_top

        gcn_type = GCNLayer

        self.lstm = StackedAlternatingLstm(
            fixed_dim + embedding_dim,
            hidden_dim,
            num_layers=num_layers,
            recurrent_dropout_probability=0.1,
        )

        if self.use_syntax:
            if num_layers_top > 0:
                if params.alter_top:
                    self.lstm_top = StackedAlternatingLstm(
                        hidden_dim,
                        hidden_dim,
                        num_layers=num_layers_top,
                        recurrent_dropout_probability=0.1,
                    )
                    self.hidden2predicate = nn.Linear(hidden_dim, hidden_dim)
                    self.hidden2argument = nn.Linear(hidden_dim, hidden_dim)
                    self.bilinear_scorer = BilinearScorer(
                        hidden_dim, tagset_size, params.bilinear_dropout
                    )
                else:
                    self.lstm_top = nn.LSTM(
                        hidden_dim,
                        hidden_dim,
                        num_layers=num_layers_top,
                        batch_first=True,
                        bidirectional=True,
                        dropout=self.params.gcn_dropout,
                    )

                    self.hidden2predicate = nn.Linear(2 * hidden_dim, hidden_dim)
                    self.hidden2argument = nn.Linear(2 * hidden_dim, hidden_dim)
                    self.bilinear_scorer = BilinearScorer(
                        hidden_dim, tagset_size, params.bilinear_dropout
                    )
            else:
                self.hidden2predicate = nn.Linear(2 * hidden_dim, hidden_dim)
                self.hidden2argument = nn.Linear(2 * hidden_dim, hidden_dim)
                self.bilinear_scorer = BilinearScorer(
                    hidden_dim, tagset_size, params.bilinear_dropout
                )

            self.gcn_w_c = gcn_type(
                hidden_dim,
                hidden_dim,
                w_c_vocab_size,
                in_arcs=True,
                out_arcs=True,
                use_gates=True,
                batch_first=True,
                residual=True,
                no_loop=True,
                dropout=self.params.gcn_dropout,
            )

            self.gcn_c_w = gcn_type(
                hidden_dim,
                hidden_dim,
                w_c_vocab_size,
                in_arcs=True,
                out_arcs=True,
                use_gates=True,
                batch_first=True,
                residual=True,
                no_loop=True,
                dropout=self.params.gcn_dropout,
            )

            self.gcn_c_c = gcn_type(
                hidden_dim,
                hidden_dim,
                c_c_vocab_size,
                in_arcs=True,
                out_arcs=True,
                use_gates=True,
                batch_first=True,
                residual=True,
                no_loop=False,
                dropout=self.params.gcn_dropout,
            )

        else:
            self.hidden2predicate = nn.Linear(hidden_dim, hidden_dim)
            self.hidden2argument = nn.Linear(hidden_dim, hidden_dim)
            self.bilinear_scorer = BilinearScorer(
                hidden_dim, tagset_size, params.bilinear_dropout
            )

        if self.eln:
            self.layernorm = nn.LayerNorm(fixed_dim)

    def forward(
        self,
        sentence,
        predicate_flags,
        sent_mask,
        lengths,
        fixed_embs,
        constituents,
        GCN_w_c,
        GCN_c_w,
        GCN_c_c,
        mask_const_batch,
        predicate_index,
        softmax_constraints,
        frame_emb_batch,
    ):

        embeds = fixed_embs

        if self.eln:
            embeds = self.layernorm(embeds * sent_mask.unsqueeze(2))

        if self.params.emb_dropout:
            embeds = self.dropout(embeds)

        embeds = torch.cat(
            (embeds, self.indicator_embeddings(predicate_flags.long())), 2
        )

        b, t, e = embeds.data.shape

        # Sort by length (keep idx)
        sent_len = torch.sort(lengths, descending=True)[0]
        idx_sort = torch.argsort(-lengths)

        if self.use_gpu:
            embeds = embeds.index_select(
                0, Variable(torch.cuda.LongTensor(idx_sort.cuda()))
            )
        else:
            embeds = embeds.index_select(0, Variable(torch.LongTensor(idx_sort)))

        packed = pack_padded_sequence(embeds, sent_len, batch_first=True)

        lstm_out, _ = self.lstm(packed)

        lstm_out, _ = pad_packed_sequence(lstm_out, batch_first=True)  # [b, t, h]

        # Un-sort by length
        idx_unsort = np.argsort(idx_sort)
        if self.use_gpu:
            lstm_out = lstm_out.index_select(
                0, Variable(torch.cuda.LongTensor(idx_unsort.cuda()))
            )
        else:
            lstm_out = lstm_out.index_select(0, Variable(torch.LongTensor(idx_unsort)))

        if self.use_syntax:
            # Here I must concatenate the constituents with the lstm_out
            gcn_in = torch.cat([lstm_out, constituents], dim=1)
            mask_all = torch.cat([sent_mask, mask_const_batch], dim=1)
            # Apply graph conv
            adj_arc_in_w_c, adj_arc_out_w_c, adj_lab_in_w_c, adj_lab_out_w_c, mask_in_w_c, mask_out_w_c, mask_loop_w_c = (
                GCN_w_c
            )

            adj_arc_in_c_w, adj_arc_out_c_w, adj_lab_in_c_w, adj_lab_out_c_w, mask_in_c_w, mask_out_c_w, mask_loop_c_w = (
                GCN_c_w
            )

            adj_arc_in_c_c, adj_arc_out_c_c, adj_lab_in_c_c, adj_lab_out_c_c, mask_in_c_c, mask_out_c_c, mask_loop_c_c = (
                GCN_c_c
            )

            gcn_out = self.gcn_w_c(
                gcn_in,
                adj_arc_in_w_c,
                adj_arc_out_w_c,
                adj_lab_in_w_c,
                adj_lab_out_w_c,
                mask_in_w_c,
                mask_out_w_c,
                mask_loop_w_c,
                mask_all,
            )

            gcn_out = self.gcn_c_c(
                gcn_out,
                adj_arc_in_c_c,
                adj_arc_out_c_c,
                adj_lab_in_c_c,
                adj_lab_out_c_c,
                mask_in_c_c,
                mask_out_c_c,
                mask_loop_c_c,
                mask_all,
            )

            gcn_out = self.gcn_c_w(
                gcn_out,
                adj_arc_in_c_w,
                adj_arc_out_c_w,
                adj_lab_in_c_w,
                adj_lab_out_c_w,
                mask_in_c_w,
                mask_out_c_w,
                mask_loop_c_w,
                mask_all,
            )

            # Take back the lstm out only
            lstm_out = gcn_out.narrow(1, 0, t)
            if self.num_layers_top > 0:
                if self.use_gpu:
                    lstm_out = lstm_out.index_select(
                        0, Variable(torch.cuda.LongTensor(idx_sort.cuda()))
                    )
                else:

                    lstm_out = lstm_out.index_select(
                        0, Variable(torch.LongTensor(idx_sort))
                    )

                packed = pack_padded_sequence(lstm_out, sent_len, batch_first=True)

                lstm_out, _ = self.lstm_top(packed)

                lstm_out, _ = pad_packed_sequence(
                    lstm_out, batch_first=True
                )  # [b, t, h]

                # Un-sort by length
                if self.use_gpu:
                    lstm_out = lstm_out.index_select(
                        0, Variable(torch.cuda.LongTensor(idx_unsort.cuda()))
                    )
                else:
                    lstm_out = lstm_out.index_select(
                        0, Variable(torch.LongTensor(idx_unsort))
                    )

        lstm_out_view = lstm_out.contiguous().view(b * t, -1)
        predicate_index = predicate_index.view(b * t)

        predicates_repr = lstm_out_view.index_select(0, predicate_index).view(b, t, -1)

        pred_repr = F.relu(self.hidden2predicate(self.dropout(predicates_repr)))
        arg_repr = F.relu(self.hidden2argument(self.dropout(lstm_out)))
        tag_scores = self.bilinear_scorer(pred_repr, arg_repr)  # [b*t, label_size]

        tag_scores = tag_scores.view(b, t, -1)
        tag_scores = tag_scores.masked_fill(
            (1 - softmax_constraints.view(b, 1, -1)).byte(), float("-1e13")
        )  # 1e-13)

        tag_scores = tag_scores.view(b * t, -1)

        return tag_scores.view(b, t, self.tagset_size)


from __future__ import division
import csv
import numpy
from copy import deepcopy
import argparse
from deep_disfluency.utils.tools import \
    convert_from_eval_tags_to_inc_disfluency_tags

def wer(r, h, macro=False):
    """
        Calculation of WER with Levenshtein distance.
        O(nm) time ans space complexity.

        >>> wer("who is there".split(), "is there".split())
        33.3333333
        >>> wer("who is there".split(), "".split())
        100.0
        >>> wer("".split(), "who is there".split())
        100.0

        macro :: Return the overall cost, else the WER
    """
    # initialisation
    d = numpy.zeros((len(r) + 1) * (len(h) + 1), dtype=numpy.int)
    d = d.reshape((len(r) + 1, len(h) + 1))

    for i in range(len(r) + 1):
        for j in range(len(h) + 1):
            if i == 0:
                d[0][j] = j
            elif j == 0:
                d[i][0] = i

    # computation
    for i in range(1, len(r) + 1):
        for j in range(1, len(h) + 1):
            if r[i - 1] == h[j - 1]:
                d[i][j] = d[i - 1][j - 1]
            else:
                substitution = d[i - 1][j - 1] + 1
                insertion = d[i][j - 1] + 1
                deletion = d[i - 1][j] + 1
                d[i][j] = min(substitution, insertion, deletion)
    if macro:
        return d[len(r)][len(h)]
    wer = 1 if len(r) == 0 and 0 < len(
        h) else d[len(r)][len(h)] / float(len(r))
    return 100 * float(wer)

def get_tags(s, open_delim='<',
             close_delim='/>'):
    """Iterator to spit out the xml style disfluency tags in a given string.

    Keyword arguments:
    s -- input string
    """
    while True:
        # Search for the next two delimiters in the source text
        start = s.find(open_delim)
        end = s.find(close_delim)
        # We found a non-empty match
        if -1 < start < end:
            # Skip the length of the open delimiter
            start += len(open_delim)
            # Spit out the tag
            yield open_delim + s[start:end].strip() + close_delim
            # Truncate string to start from last match
            s = s[end + len(close_delim):]
        else:
            return

def get_turn_number(triple_string):
    triple = triple_string[1:].strip(")").split(":")
    return int(triple[0])

def interactive_dialogue(pair):
    """Return the data from both speakers in the pair as
    a single dialogue with interactive speaker tags
    for switching speakers or sticking with the speakers
    """
    speaker_1, speaker_2 = pair
    frames_data, lex_data, pos_data, indices_data, labels_data = speaker_1
    # normalize for the sort so that the turn can't change when <c is in
    # the uttseg
    for i, label in enumerate(labels_data):
        if "<t" not in label:
            indices_data[i] = indices_data[i - 1]  # dynamically updates
    # put a search criteria on the first speaker
    # then sort the rest in place
    # turn into two list of tuples
    s1 = deepcopy([(f, l, p, i, y, "a") for f, l, p, i, y in zip(
        frames_data,
        lex_data,
        pos_data,
        indices_data,
        labels_data)])
    frames_data, lex_data, pos_data, indices_data, labels_data = speaker_2
    for i, label in enumerate(labels_data):
        if "<t" not in label:
            indices_data[i] = indices_data[i - 1]  # dynamically updates
    s2 = deepcopy([(f, l, p, i, y, "b") for f, l, p, i, y in zip(
        frames_data,
        lex_data,
        pos_data,
        indices_data,
        labels_data)])
    all_data = sorted(s1 + s2, key=lambda x: get_turn_number(x[3]))
    previous_speaker = ""
    previous_tag = ""
    for l in range(len(all_data)):
        line = list(all_data[l])
        # print line
        speaker = 'keep'
        if not line[5] == previous_speaker:
            speaker = 'take'
            previous_speaker = line[5]
        line[4] = line[4] + '<speaker floor="{}"/>'.format(speaker)
        all_data[l] = tuple(line)
        if "t/>" in previous_tag and "<c" in line[4]:
            print "warning interactive, continuation after end"
            print line
        previous_tag = line[4]
    final_data = []
    for i in range(5):
        final_data.append([x[i] for x in all_data])
    return final_data

def concat_all_data_all_speakers(dialogues, interactive_sort=False,
                                 divide_into_utts=False,
                                 convert_to_dnn_format=False,
                                 representation="disf_1",
                                 limit=8):
    """Concatenates all the data together as lists of lists"""
    frames = []
    words = []
    pos = []
    indices = []
    labels = []
    current_pair = []
    for d in sorted(dialogues, key=lambda x: x[0]):
        dialogueID, data = d
        if interactive_sort:
            current_pair.append(deepcopy(d))
            if len(current_pair) == 2:
                assert(current_pair[0][0][:-1] == current_pair[1][0][:-1])
                print[x[0] for x in current_pair]
                data = deepcopy(interactive_dialogue(deepcopy([x[1] for x in
                                                               current_pair])))
                current_pair = []
            else:
                continue
        frames_data, lex_data, pos_data, indices_data, labels_data = data
        lex_data = [x[0][1] for x in lex_data]
        pos_data = [x[1] for x in pos_data]
        if divide_into_utts:
            frames.append(deepcopy(frames_data))
            current_lex = []
            current_pos = []
            current_idx = []
            current_label = []
            started = False
            last_idx = ""
            for i, w_i, p_i, i_i, l_i in zip(range(0, len(lex_data)),
                                             lex_data,
                                             pos_data,
                                             indices_data,
                                             labels_data):
                print w_i, p_i, i_i, l_i
                if (started and i_i.split(":")[0] != last_idx) or \
                        i == len(lex_data) - 1:
                    if convert_to_dnn_format:
                        current_label = \
                            convert_from_eval_tags_to_inc_disfluency_tags(
                                current_label,
                                current_lex,
                                representation,
                                limit)
                    words.append(deepcopy(current_lex))
                    pos.append(deepcopy(current_pos))
                    indices.append(deepcopy(current_idx))
                    labels.append(deepcopy(current_label))
                    current_lex = []
                    current_pos = []
                    current_idx = []
                    current_label = []
                current_lex.append(w_i)
                current_pos.append(p_i)
                current_idx.append(i_i)
                last_idx = i_i.split(":")[0]
                current_label.append(l_i)
                started = True
        else:
            frames.append(deepcopy(frames_data))
            words.append(deepcopy(lex_data))
            pos.append(deepcopy(pos_data))
            indices.append(deepcopy(indices_data))
            labels.append(deepcopy(labels_data))

    print "concatenated all data"
    # print [len(x) for x in [frames, words, pos, indices, labels]]
    return frames, words, pos, indices, labels

def add_word_continuation_tags(tags):
    """Returns list with continuation tags for each word:
    <cc/> continues current dialogue act and the next word will also continue
    <ct/> continues current dialogue act and is the last word of it
    <tc/> starts this dialogue act tag and the next word continues it
    <tt/> starts and ends dialogue act (single word dialogue act)
    """
    tags = list(tags)
    for i in range(0, len(tags)):
        if i == 0:
            tags[i] = tags[i] + "<t"
        else:
            tags[i] = tags[i] + "<c"
        if i == len(tags) - 1:
            tags[i] = tags[i] + "t/>"
        else:
            tags[i] = tags[i] + "c/>"
    return tags

def sort_into_dialogue_speakers(IDs, mappings, utts, pos_tags=None,
                                labels=None,
                                add_uttseg=True,
                                add_dialogue_acts=True,
                                convert_to_dnn_tags=False):
    # for each utterance, given its ID get its conversation number
    # and dialogue participant in the format
    # needed for MSALign files
    # return a list of tuples of (dialogueID+speaker [eg. sw4004A],
    # dialogue_act
    # ((ID1,utterance)),(ID2,utterance))
    dialogue_speakers = []
    currentA = ""
    currentB = ""
    A_utts = []
    B_utts = []
    A_mappings = []
    B_mappings = []
    A_pos = []
    B_pos = []
    A_labels = []
    B_labels = []

    for ID, map, utt, pos, tags in zip(IDs, mappings, utts, pos_tags, labels):
        # if "3756" in ID:
        # print ID, mapping, utt
        split = ID.split(":")
        dialogue = split[0]
        speaker = split[1]
        uttID = split[2]
        # mapping = [uttID] * len(utt)
        dialogue_act = split[3]
        current_speaker = "".join([dialogue, speaker])
        if "A" in current_speaker:
            if current_speaker != currentA and not currentA == "":
                dialogue_speakers.append((currentA, A_mappings, A_utts, A_pos,
                                          A_labels))
                A_utts = []
                A_mappings = []
                A_pos = []
                A_labels = []
            currentA = current_speaker
            A_utts.extend(list(utt))
            A_mappings.extend(list(map))
            A_pos.extend(list(pos))
            if convert_to_dnn_tags:
                tags = convert_from_eval_tags_to_inc_disfluency_tags(tags, utt)
            if add_uttseg:
                tags = add_word_continuation_tags(tags)
            if add_dialogue_acts:
                tags = [this_tag + '<diact type="{0}"/>'
                        .format(dialogue_act) for this_tag in tags]
            A_labels.extend(tags)
        elif "B" in current_speaker:
            if current_speaker != currentB and not currentB == "":
                dialogue_speakers.append((currentB, B_mappings, B_utts, B_pos,
                                          B_labels))
                B_utts = []
                B_mappings = []
                B_pos = []
                B_labels = []
            currentB = current_speaker
            B_utts.extend(list(utt))
            B_mappings.extend(list(map))
            B_pos.extend(list(pos))
            if convert_to_dnn_tags:
                tags = convert_from_eval_tags_to_inc_disfluency_tags(tags, utt)
            if add_uttseg:
                tags = add_word_continuation_tags(tags)
            if add_dialogue_acts:
                tags = [this_tag + '<diact type="{0}"/>'
                        .format(dialogue_act) for this_tag in tags]
            B_labels.extend(tags)

    if not (currentA, A_mappings, A_utts, A_pos, A_labels) in\
            dialogue_speakers[-2:]:
        # if current_speaker != currentA and not currentA == "":
        dialogue_speakers.append((currentA, A_mappings, A_utts, A_pos,
                                  A_labels))  # concatenate them all together
    if not (currentB, B_mappings, B_utts, B_pos, B_labels) in \
            dialogue_speakers[-2:]:
        # if current_speaker != currentB and not currentB == "":
        dialogue_speakers.append((currentB, B_mappings, B_utts, B_pos,
                                  B_labels))  # concatenate them all together
    return dialogue_speakers

def load_data_from_disfluency_corpus_file(fp, representation="disf1", limit=8,
                                          convert_to_dnn_format=False):
    """Loads from file into five lists of lists of strings of equal length:
    one for utterance iDs (IDs))
    one for word timings of the targets (start,stop)
    one for words (seq),
    one for pos (pos_seq)
    one for tags (targets).

    NB this does not convert them into one-hot arrays, just outputs lists of string tags."""
    print "loading from", fp
    f = open(fp)
    print "loading data", f.name
    count_seq = 0
    IDs = []
    seq = []
    pos_seq = []
    targets = []
    timings = []
    currentTimings = []
    # marks the current fake time for the dialogue (i.e. end of word)
    current_fake_time = 0
    current_dialogue = ""

    reader = csv.reader(f, delimiter='\t')
    counter = 0
    utt_reference = ""
    currentWords = []
    currentPOS = []
    currentTags = []
    current_fake_time = 0

    # corpus = "" # can write to file
    for ref, timing, word, postag, disftag in reader:  # mixture of POS and Words
        # print ref, timing, word, postag, disftag
        counter += 1
        if not ref == "":
            if count_seq > 0:  # do not reset the first time
                # convert to the inc tags
                if convert_to_dnn_format:
                    currentTags = \
                        convert_from_eval_tags_to_inc_disfluency_tags(
                            currentTags,
                            currentWords,
                            representation,
                            limit)
                # corpus+=utt_reference #write data to a file for checking
                # convert to vectors
                seq.append(tuple(currentWords))
                pos_seq.append(tuple(currentPOS))
                targets.append(tuple(currentTags))
                IDs.append(utt_reference)
                timings.append(tuple(currentTimings))
                # reset the words
                currentWords = []
                currentPOS = []
                currentTags = []
                currentTimings = []
            # set the utterance reference
            count_seq += 1
            utt_reference = ref
            if not utt_reference.split(":")[0] == current_dialogue:
                current_dialogue = utt_reference.split(":")[0]
                current_fake_time = 0  # TODO fake for now- reset the current beginning of word time
        currentWords.append(word)
        currentPOS.append(postag)
        currentTags.append(disftag)
        currentTimings.append(timing)
        current_fake_time += 1
    # flush
    if not currentWords == []:
        if convert_to_dnn_format:
            currentTags = \
                convert_from_eval_tags_to_inc_disfluency_tags(
                    currentTags,
                    currentWords,
                    representation,
                    limit)
        seq.append(tuple(currentWords))
        pos_seq.append(tuple(currentPOS))
        targets.append(tuple(currentTags))
        IDs.append(utt_reference)
        timings.append(tuple(currentTimings))

    assert len(seq) == len(targets) == len(pos_seq)
    print "loaded " + str(len(seq)) + " sequences"
    f.close()
    return (IDs, timings, seq, pos_seq, targets)

def load_data_from_corpus_file(filename, limit=8,
                               representation="disf1",
                               convert_to_dnn_format=False):
    """Loads from disfluency detection with timings file.
    """
    all_speakers = []
    lex_data = []
    pos_data = []
    frames = []
    labels = []
    indices = []

    latest_increco = []
    latest_pos = []
    latest_labels = []
    latest_indices = []

    a_file = open(filename)
    started = False
    conv_no = ""
    prev_word = -1
    prev_pos = -1
    prev_index = ""
    print "loading in timings file", filename, "..."
    for line in a_file:
        if "Speaker:" in line:
            if not started:
                started = True
            else:
                # flush
                # print line
                if not latest_increco == []:
                    shift = -1
                    for i in range(0, len(latest_increco)):
                        triple = latest_increco[i]
                        if triple[1] == triple[2]:
                            # print "same timing!", triple
                            shift = i
                            break
                    if shift > -1:
                        latest_increco = fill_in_time_approximations(
                            latest_increco,
                            shift)
                    lex_data.extend(deepcopy(latest_increco))
                    pos_data.extend(deepcopy(latest_pos))
                    # convert to the disfluency tags for this
                    # latest_labels = convertFromEvalTagsToIncDisfluencyTags()
                    if convert_to_dnn_format:
                        latest_labels = \
                            convert_from_eval_tags_to_inc_disfluency_tags(
                                latest_labels,
                                latest_increco,
                                representation,
                                limit)
                    labels.extend(deepcopy(latest_labels))
                    indices.extend(deepcopy(latest_indices))
                # fake
                # print lex_data
                frames = [x[-1] for x in lex_data]  # last word end time
                # acoustic_data = [0] * len(lex_data)  # fakes
                #indices = [0] * len(lex_data)

                all_speakers.append((conv_no, (frames, lex_data, pos_data,
                                               indices, labels)))
                # reset
                lex_data = []
                pos_data = []
                labels = []
                indices = []
                latest_increco = []
                latest_pos = []
                latest_labels = []
                latest_indices = []
                prev_word = -1
                prev_pos = -1
                prev_index = ""

            conv_no = line.strip("\n").replace("Speaker: ", "")
            continue
        if line.strip("\n") == "":
            continue
        spl = line.strip("\n").split("\t")
        # print "@@@@" + line + "@@@@"
        # raw_input()
        idx = spl[0]
        start = float(spl[1])
        end = float(spl[2])
        word = spl[3]
        pos = spl[4]
        # need to convert to the right rep here
        tag = spl[5]

        latest_increco.append(([prev_word, word], start, end))
        latest_pos.append(deepcopy([prev_pos, pos]))
        latest_labels.append(tag)
        latest_indices.append(idx)
        prev_word = word
        prev_pos = pos

    # flush
    if not latest_increco == []:
        shift = -1
        for i in range(0, len(latest_increco)):
            triple = latest_increco[i]
            if triple[1] == triple[2]:
                shift = i
                break
        if shift > -1:
            latest_increco = fill_in_time_approximations(latest_increco, shift)
        lex_data.extend(latest_increco)
        pos_data.extend(latest_pos)
        if convert_to_dnn_format:
            latest_labels = \
                convert_from_eval_tags_to_inc_disfluency_tags(
                    latest_labels,
                    latest_increco,
                    representation,
                    limit)
        labels.extend(latest_labels)
        indices.extend(latest_indices)
    frames = [x[-1] for x in lex_data]  # last word end time
    # acoustic_data = [0,] * len(lex_data)  # fakes..
    # indices = [0] * len(lex_data)
    all_speakers.append((conv_no, (frames, lex_data, pos_data,
                                   indices, labels)))
    print len(all_speakers), "speakers with timings input"
    # print "first few lengths"
    # limit = 10
    # for s in all_speakers:
    #     id, data = s
    #     print [len(x) for x in data]
    #     limit-=1
    #     if limit<0: break
    return all_speakers

def get_diff_and_new_prefix(current, newprefix, verbose=False):
    """Only get the different right frontier according to the timings
    and change the current hypotheses"""
    if verbose:
        print "current", current
        print "newprefix", newprefix
    rollback = 0
    original_length = len(current)
    original_current = deepcopy(current)
    for i in range(len(current) - 1, -2, -1):
        if verbose:
            print "oooo", newprefix[0]
            if not current == []:
                print current[i]
        if i == -1 or (float(newprefix[0][1]) >= float(current[i][2])):
            if i == len(current) - 1:
                current = current + newprefix
                break
            k = 0
            marker = i + 1
            for j in range(i + 1, len(current)):
                if k == len(newprefix):
                    break
                if verbose:
                    print "...", j, k, current[j], newprefix[k], len(newprefix)
                if not current[j] == newprefix[k]:
                    break
                else:
                    if verbose:
                        print "repeat"
                    k += 1
                    marker = j + 1
            rollback = original_length - marker
            current = current[:marker] + newprefix[k:]
            newprefix = newprefix[k:]
            break
    if newprefix == []:
        rollback = 0  # just no rollback if no prefix
        current = original_current  # reset the current
    if verbose:
        print "current after call", current
        print "newprefix after call", newprefix
        print "rollback after call", rollback
    return (current, newprefix, rollback)

def simulate_increco_data(frame, acoustic_data, lexical_data, pos_data):
    """For transcripts + timings, create tuples of single hypotheses
    to simulate perfect ASR at the end of each word."""
    new_lexical_data = []
    new_pos_data = []
    new_acoustic_data = []
    current_time = 0
    for my_frame, acoust, word, pos in zip(frame, acoustic_data,
                                           lexical_data, pos_data):
        # print type(my_frame), type(acoust), type(word), type(pos)
        new_lexical_data.append([(word, current_time / 100, my_frame / 100)])
        current_time = my_frame
        new_pos_data.append([pos])
        new_acoustic_data.append([acoust])
    return new_acoustic_data, new_lexical_data, new_pos_data

def fill_in_time_approximations(word_timing_tuples, idx):
    """Called when the start and end time of a word are the same
    (i.e. it has no length, due to technical error), it searches forward
    in the list until
    it meets a word with a different end time, then spreads the time of
    the affected timing tuples stand and end time accordingly.

    word_timing_tuples :: list, of (word, start_time, end_time) triples
    idx :: int, the index of the first offending timing in the list
    """
    affected_indices = []
    start_time = word_timing_tuples[idx][1]
    end_time = word_timing_tuples[idx][2]
    assert start_time == end_time,\
        "Function not needed- different start and end times!"
    assert len(word_timing_tuples) > 1
    if end_time == word_timing_tuples[-1][2]:
        # this end time is the same as the latest one
        # need backward search for time
        # print "backward search"
        idx = len(word_timing_tuples) - 1
        for i in range(len(word_timing_tuples) - 1, -1, -1):
            affected_indices.append(i)
            idx = i
            if not word_timing_tuples[i][1] == start_time:
                start_time = word_timing_tuples[i][1]
                break
    else:
        for i in range(idx, len(word_timing_tuples)):
            affected_indices.append(i)
            if not word_timing_tuples[i][2] == start_time:
                end_time = word_timing_tuples[i][2]
                break
    total_time = end_time - start_time
    assert total_time > 0.0, str(word_timing_tuples[affected_indices[0]:]) +\
        str(idx)
    mean_len = total_time / len(affected_indices)
    for i in range(idx, idx + len(affected_indices)):
        end_time = start_time + mean_len
        word_timing_tuples[i] = (word_timing_tuples[i][0], start_time,
                                 end_time)
        start_time = end_time
    return word_timing_tuples

def process_arguments(config=None,
                      exp_id=None,
                      heldout_file="../data/disfluency_detection/switchboard/swbd_heldout_partial_data.csv",
                      test_file="../data/disfluency_detection/switchboard/swbd_test_partial_data.csv",
                      use_saved=None,
                      hmm=False,
                      verbose=True):
    """Loads arguments for an experiment from a config file

    Keyword arguments:
    config -- the config file location, default None
    exp_id -- the experiment ID name, default None
    """
    parser = argparse.ArgumentParser(
        description='This script trains a RNN for disfluency detection and saves the best models and results to disk.')
    parser.add_argument('-c', '--config', type=str,
                        help='The location of the config file.',
                        default=config)
    parser.add_argument('-e', '--exp_id', type=str,
                        help='The experiment number from which to load arguments from the config file.',
                        default=exp_id)
    parser.add_argument('-v', '--heldout_file', type=str,
                        help='The path to the validation file.',
                        default=heldout_file)
    parser.add_argument('-t', '--test_file', type=str,
                        help='The path to the test file.',
                        default=test_file)
    parser.add_argument('-m', '--use_saved_model', type=int,
                        help='Epoch number of the pre-trained model to load.',
                        default=use_saved)
    parser.add_argument('-hmm', '--hmm', type=bool,
                        help='Whether to use hmm disfluency decoder.',
                        default=hmm)
    parser.add_argument('-verb', '--verbose', type=bool,
                        help='Whether to output training progress.',
                        default=verbose)
    args = parser.parse_args()

    header = [
        'exp_id',  # experiment id
        'model',  # can be elman/lstm/mt_elman/mt_lstm
        'lr',  # learning rate
        'decay',  # decay on the learning rate if improvement stops
        'seed',  # random seed
        'window',
        # number of words in the context window (backwards only for disfluency)
        'bs',  # number of backprop through time steps
        'emb_dimension',  # dimension of word embedding
        'nhidden',  # number of hidden units
        'nepochs',  # maximum number of epochs
        'train_data',  # which training data
        'loss_function',  # default will be nll, unlikely to change
        'reg',  # regularization type
        'pos',  # whether pos tags or not
        'acoustic',  # whether using aoustic features or not
        'embeddings',  # embedding files, if any
        'update_embeddings',  # whether the embeddings should be updated at runtime or not
        'batch_size',  # batch size, 'word' or 'utterance'
        'tags',  # the output tag representations used
        'end_utterance'  # whether we do combined end of utterance detection too
    ]
    print header
    if args.config:
        for line in open(args.config):
            features = line.strip("\n").split("\t")
            # print features
            if features[0] != args.exp_id:
                continue
            for i in range(1, len(header)):
                feat_value = features[i].strip()  # if string
                # print feat_value, "feat"
                if header[i] in ['lr']:
                    feat_value = float(feat_value)
                elif header[i] in ['decay', 'pos', 'acoustic', 'update_embeddings', 'end_utterance']:
                    if feat_value == 'True':
                        feat_value = True
                    else:
                        feat_value = False
                elif header[i] in ['seed', 'window', 'bs', 'emb_dimension', 'nhidden', 'nepochs']:
                    feat_value = int(feat_value)
                elif feat_value == 'None':
                    feat_value = None
                setattr(args, header[i], feat_value)
    # print args
    return args

if __name__ == "__main__":
    print wer("who is there".split(), "is there".split(), macro=True)
    print wer("who is there".split(), "".split(), macro=True)
    print wer("".split(), "who is there".split(), macro=True)