def main():
    cutoff = {}
    cutoff[(40,40)] = 3.6
    cutoff[(13,29)] = 3.6
    cutoff[(29,13)] = 3.6
    cutoff[(40,13)] = 3.6
    cutoff[(13,40)] = 3.6
    cutoff[(29,40)] = 3.6
    cutoff[(40,29)] = 3.6
    cutoff[(13,13)] = 3.6
    cutoff[(29,29)] = 3.6

    cutoff[(41,41)] = 3.7
    cutoff[(28,28)] = 3.7
    cutoff[(41,28)] = 3.7
    cutoff[(28,41)] = 3.7

    cutoff[(46,46)] = 3.45
    cutoff[(14,14)] = 3.45
    cutoff[(46,14)] = 3.45
    cutoff[(14,46)] = 3.45

    paramfile = sys.argv[1]
    vp_dict = load_param_file(paramfile)

    modelfiles = sys.argv[2:]
    count = defaultdict(int) # Stores how many VPs have been found of each index type
    count = 0
    direc = 'ZrCuAl/md_80k/'
    for modelfile in modelfiles:
        print(modelfile)
        m = Model(modelfile)
        m.generate_neighbors(cutoff)
        #voronoi_3d(m,cutoff)
        #set_atom_vp_types(m,vp_dict)
        #vor_stats(m)
        #cats = index_stats(m)
        #for atom in m.atoms:
        #    #new = Model('0,0,12,0; number of atoms is {0};'.format(count), m.lx, m.ly, m.lz, atom.neighs + [atom])
        #    new = Model('{0}; number of atoms is {1};'.format(atom.vp.index, count[atom.vp.index]), m.lx, m.ly, m.lz, atom.neighs + [atom])
        #    fix_cluster_pbcs(new)
        #    val = normalize_bond_distances(new)
        #    new.comment = '{0}; number of atoms is {1}; bond length scaling factor is {2}'.format(atom.vp.index, count,val)
        #    center = find_center_atom(new)
        #    new.remove(center)
        #    new.add(center)
        #    vp_str = ''.join([str(x) for x in atom.vp.index])
        #    if not os.path.exists(direc+vp_str):
        #        os.makedirs(direc+vp_str)
        #    new.write(direc+'{0}/{0}.{1}.xyz'.format(vp_str, count[atom.vp.index]))
        #    count[atom.vp.index] += 1
        #print(count)
        cn = 0.0
        for atom in m.atoms:
            cn += atom.cn
        cn = float(cn)/m.natoms
        print(cn)

        for atom in m.atoms:
            new_cut = copy.copy(cutoff)
            old_cn = atom.cn
            inc = 0.0
            while atom.cn < 12:
                for key,val in new_cut.items(): new_cut[key] = val + 0.1
                inc += 0.1
                atom.neighs = m.get_atoms_in_cutoff(atom, new_cut)
                if(atom in atom.neighs): atom.neighs.remove(atom)
                atom.cn = len(atom.neighs)
            new = Model('CN changed from {0} to {1};'.format(old_cn, atom.cn), m.lx, m.ly, m.lz, atom.neighs + [atom])
            new.write('temp/{0}.xyz'.format(count))
            if inc > 0.0: print("Increased shell by {0} Ang. for atom {1}".format(inc, count))
            count += 1
        cn = 0.0
        for atom in m.atoms:
            cn += atom.cn
        cn = float(cn)/m.natoms
        print(cn)

    return 0

    modelfile = sys.argv[2]
    m = Model(modelfile)
    xtal_atoms = sys.argv[3]
    xtal_atoms = Model(xtal_atoms).atoms

    #x,y,z = (round(x,6) for x in self.coord)
    #a,b,c = (round(x,6) for x in other.coord)
    #print("HERE")
    #print([round(x,7) for x in xtal_atoms[13].coord])
    #print([round(x,7) for x in m.atoms[153].coord])
    #print(type(xtal_atoms[13]))
    #print(type(m.atoms[153]))
    #print(xtal_atoms[13] == m.atoms[153])
    #return 0

    glassy_atoms = []
    for atom in m.atoms:
        if atom not in xtal_atoms:
            glassy_atoms.append(atom)
    print(len(glassy_atoms))
    print(len(xtal_atoms))
    assert len(glassy_atoms) + len(xtal_atoms) == m.natoms
    voronoi_3d(m, cutoff)
    set_atom_vp_types(m, vp_dict)
    m.generate_neighbors(3.45)
    head, tail = os.path.split(modelfile)
    head = head + '/'
    if not os.path.exists(head+'glassy/'):
        os.makedirs(head+'glassy/')
    if not os.path.exists(head+'xtal/'):
        os.makedirs(head+'xtal/')
    for count, atom in enumerate(xtal_atoms):
        i = m.atoms.index(atom)
        atom = m.atoms[i]
        new = Model('{0}'.format(atom.vp.index), m.lx, m.ly, m.lz, atom.neighs + [atom])
        fix_cluster_pbcs(new)
        val = normalize_bond_distances(new)
        new.comment = '{0}; bond length scaling factor is {1}'.format(atom.vp.index, val)
        center = find_center_atom(new)
        new.remove(center)
        new.add(center)
        vp_str = ''.join([str(x) for x in atom.vp.index])
        new.write(head+'xtal/{0}.xyz'.format(count))
    for count, atom in enumerate(glassy_atoms):
        i = m.atoms.index(atom)
        atom = m.atoms[i]
        new = Model('{0}'.format(atom.vp.index), m.lx, m.ly, m.lz, atom.neighs + [atom])
        fix_cluster_pbcs(new)
        val = normalize_bond_distances(new)
        new.comment = '{0}; bond length scaling factor is {1}'.format(atom.vp.index, val)
        center = find_center_atom(new)
        new.remove(center)
        new.add(center)
        vp_str = ''.join([str(x) for x in atom.vp.index])
        new.write(head+'glassy/{0}.xyz'.format(count))
    return 0
        




    for atom in volume_atoms.atoms:
        for i,atom2 in enumerate(m.atoms):
            if(atom.z == atom2.z and [round(x, 5) for x in atom.coord] == [round(x, 5) for x in atom2.coord]): good[i] = True
    count = defaultdict(int) # Stores how many VPs have been found of each index type
    for modelfile in modelfiles:
        print(modelfile)
        m = Model(modelfile)
        voronoi_3d(m,cutoff)
        set_atom_vp_types(m,vp_dict)
        #vor_stats(m)
        #cats = index_stats(m)
        for i,atom in enumerate(m.atoms):
            if not good[i]: continue
            #new = Model('0,0,12,0; number of atoms is {0};'.format(count), m.lx, m.ly, m.lz, atom.neighs + [atom])
            new = Model('{0}; number of atoms is {1};'.format(atom.vp.index, count), m.lx, m.ly, m.lz, atom.neighs + [atom])
            fix_cluster_pbcs(new)
            val = normalize_bond_distances(new)
            new.comment = '{0}; number of atoms is {1}; bond length scaling factor is {2}'.format(atom.vp.index, count,val)
            center = find_center_atom(new)
            new.remove(center)
            new.add(center)
            vp_str = ''.join([str(x) for x in atom.vp.index])
            if not os.path.exists(vp_str):
                os.makedirs(vp_str)
            new.write('{0}/{0}.{1}.xyz'.format(vp_str, count[atom.vp.index]))
            count[atom.vp.index] += 1
        print(count)
        for c,v in count.items():
            print("{0}: {1}".format(c,v))
        print(sum(count.values()))
        cn = 0.0
        for atom in m.atoms:
            cn += atom.cn
        cn = float(cn)/m.natoms
        print(cn)
def main():
    # Load the MD model
    modelfile = sys.argv[1]
    md = Model(modelfile)

    # Make a model to hold all the atoms we pull out to make sure the original model was sampled uniformly.
    # Only the center atoms are added to this model.
    holding_model = Model(comment='holding box', xsize=md.xsize, ysize=md.ysize, zsize=md.zsize, atoms=[])

    # Load the cutoff dictionary so that we can generate neighbors for every atom
    from cutoff import cutoff

    # Directory name to save the files to
    dir = 'all_91200_clusters'
    if not os.path.exists(dir):
        os.makedirs(dir)

    # Set the number of clusters to randomly select
    num_clusters = 'all'
    if num_clusters == 'all' or num_clusters == md.natoms:
        random_selection = False
        num_clusters = md.natoms
    else:
        num_clusters = min(num_clusters, md.natoms)
        random_selection = True

    if random_selection:
        unpicked = range(md.natoms)

    start = 0
    for n in range(start, num_clusters+start):
        if random_selection:
            rand = random.choice(unpicked)
            unpicked.remove(rand)
            atom = md.atoms[rand]
        else:
            atom = md.atoms[n-start]
            rand = n-start

        # Generate neighbors for this atom
        atom.neighs = md.get_atoms_in_cutoff(atom, cutoff)
        if(atom in atom.neighs): atom.neighs.remove(atom)
        atom.cn = len(atom.neighs)

        holding_model.add(atom)

        # Create the cluster, normalize the bond distances, and write to disk
        atoms = atom.neighs + [atom]
        # c = create_cluster(md, atom, start, n, rand, center_included=True)
        c = Cluster(
            comment='cluster #{0} from atom {1}'.format(n-start, rand),
            xsize=md.xsize,
            ysize=md.ysize,
            zsize=md.zsize,
            atoms=atoms,
            center_included = True
        )
        ratio = c.normalize_bond_distances()
        c.comment='cluster #{0} from atom {1}; normalized bond distances by {2}'.format(n-start, rand, ratio)
        c.write(os.path.join(dir, '{0}.xyz'.format(n)))
        print(n-start, atom.id)
        #print(c)

    holding_model.write(os.path.join(dir, 'holding_model.xyz'))