def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('assignments', default='Macro4/MacroAssignments.h5', help='Path to an assignments file. (default=Macro4/MacroAssignments.h5)')
    parser.add_argument('--project', default='ProjectInfo.yaml', help='Path to ProjectInfo.yaml file. (default=ProjectInfo.yaml)')
    args = parser.parse_args()

    project = Project.load_from(args.project)
    t = reduce(operator.add, (project.load_traj(i) for i in range(project.n_trajs)))

    phi_angles = md.compute_dihedrals(t, [PHI_INDICES]) * 180.0 / np.pi
    psi_angles = md.compute_dihedrals(t, [PSI_INDICES]) * 180.0 / np.pi
    state_index = np.hstack(io.loadh(args.assignments)['arr_0'])

    for i in np.unique(state_index):
        pp.plot(phi_angles[np.where(state_index == i)],
                psi_angles[np.where(state_index == i)],
                'x', label='State %d' % i)


    pp.title("Alanine Dipeptide Macrostates")
    pp.xlabel(r"$\phi$")
    pp.ylabel(r"$\psi$")
    annotate()

    pp.legend(loc=1, labelspacing=0.075, prop={'size': 8.0}, scatterpoints=1,
              markerscale=0.5, numpoints=1)
    pp.xlim([-180, 180])
    pp.ylim([-180, 180])
    pp.show()

def test_generator():
    N_FRAMES = 2
    N_ATOMS = 5
    xyz = np.asarray(np.random.randn(N_FRAMES, N_ATOMS, 3), dtype=np.float32)
    ptraj = md.Trajectory(xyz=xyz, topology=None)

    quartets = np.array(list(itertools.combinations(range(N_ATOMS), 4)), dtype=np.int32)
    quartets2 = itertools.combinations(range(N_ATOMS), 4)
    a = md.compute_dihedrals(ptraj, quartets)
    b = md.compute_dihedrals(ptraj, quartets2)
    eq(a, b)

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space via calculation
        of dihedral (torsion) angles of alpha carbon backbone

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        """

        ca = [a.index for a in traj.top.atoms if a.name == 'CA']
        if len(ca) < 4:
            return np.zeros((len(traj), 0), dtype=np.float32)

        alpha_indices = np.array(
            [(ca[i - 1], ca[i], ca[i+1], ca[i + 2]) for i in range(1, len(ca) - 2)])
        result = md.compute_dihedrals(traj, alpha_indices)

        x = []
        if self.atom_indices is None:
            self.atom_indices = np.vstack(alpha_indices)
        if self.sincos:
            x.extend([np.cos(result), np.sin(result)])
        else:
            x.append(result)
        return np.hstack(x)

def test_no_indices():
    for fn in ['2EQQ.pdb', '1bpi.pdb']:
        for opt in [True, False]:
            t = md.load(get_fn(fn))
            assert md.compute_distances(t, np.zeros((0,2), dtype=int), opt=opt).shape == (t.n_frames, 0)
            assert md.compute_angles(t, np.zeros((0,3), dtype=int), opt=opt).shape == (t.n_frames, 0)
            assert md.compute_dihedrals(t, np.zeros((0,4), dtype=int), opt=opt).shape == (t.n_frames, 0)

def test_dihedral_1():
    pymol = find_executable('pymol')
    if pymol is None:
        raise SkipTest("pymol executable not found")
    
    xyz = '''MODEL        0
ATOM      1    A ACE     1       4.300  13.100   8.600  1.00  0.00
ATOM      2    B ACE     1       5.200  13.600   8.800  1.00  0.00
ATOM      3    C ACE     1       4.900  14.300   9.600  1.00  0.00
ATOM      4    D ACE     1       5.600  14.200   7.900  1.00  0.00
    '''
    script = '''
with open('output.txt', 'w') as f:
    f.write('%f' % cmd.get_dihedral('1/A', '1/B', '1/C', '1/D'))
'''
    
    with enter_temp_directory():
        with open('xyz.pdb', 'w') as f:
            f.write(xyz)
        with open('pymolscript.py', 'w') as f:
            f.write(script)

        os.system('%s %s -cr %s' % (pymol, 'xyz.pdb', 'pymolscript.py'))
        with open('output.txt') as f:
            pymol_value = np.deg2rad(float(f.read()))
        t = md.load('xyz.pdb')

    mdtraj_value = md.compute_dihedrals(t, [[0,1,2,3]])[0,0]
    
    np.testing.assert_array_almost_equal(pymol_value, mdtraj_value)

 def get_phipsi(self, trajs, phi, psi):
     #phi = [6, 8, 14, 16]
     #psi = [4, 6, 8, 14]
     PHI_INDICES = []
     PSI_INDICES = []
     for i in xrange(len(phi)):
         PHI_INDICES.append(self.atom_indices.index(phi[i]))
         PSI_INDICES.append(self.atom_indices.index(psi[i]))
     len_trajs = len(trajs)
     print "PSI:", PSI_INDICES
     print "PHI:", PHI_INDICES
     phi_angles = md.compute_dihedrals(trajs, [PHI_INDICES]) * 180.0 / np.pi
     psi_angles = md.compute_dihedrals(trajs, [PSI_INDICES]) * 180.0 / np.pi
     #phi_psi=np.column_stack((phi_angles, psi_angles))
     #return phi_psi
     return phi_angles, psi_angles

 def chi1_feat(traj, res):
     chi1 = traj.topology.select('resid %i and (name C or name CA or name CB or name CG or name SG or name CG1 or name OG or name OG1)' %res)
     if chi1.shape[0] != 4:
         return None
     chi1 = chi1.reshape([1,4])
     traj_chi1 = md.compute_dihedrals(traj, chi1)
     return traj_chi1

def sidechain_example(yaml_file):
    # Parse a YAML configuration, return as Dict
    cfg = Settings(yaml_file).asDict()
    structure = cfg['Structure']

    #Select move type
    sidechain = SideChainMove(structure, [1])
    #Iniitialize object that selects movestep
    sidechain_mover = MoveEngine(sidechain)

    #Generate the openmm.Systems outside SimulationFactory to allow modifications
    systems = SystemFactory(structure, sidechain.atom_indices, cfg['system'])

    #Generate the OpenMM Simulations
    simulations = SimulationFactory(systems, sidechain_mover, cfg['simulation'], cfg['md_reporters'],
                                    cfg['ncmc_reporters'])

    # Run BLUES Simulation
    blues = BLUESSimulation(simulations, cfg['simulation'])
    blues.run()

    #Analysis
    import mdtraj as md
    import numpy as np

    traj = md.load_netcdf('vacDivaline-test/vacDivaline.nc', top='tests/data/vacDivaline.prmtop')
    indicies = np.array([[0, 4, 6, 8]])
    dihedraldata = md.compute_dihedrals(traj, indicies)
    with open("vacDivaline-test/dihedrals.txt", 'w') as output:
        for value in dihedraldata:
            output.write("%s\n" % str(value)[1:-1])

def test_dihedral_1(pymol, tmpdir):
    xyz = '''MODEL        0
ATOM      1    A ACE     1       4.300  13.100   8.600  1.00  0.00
ATOM      2    B ACE     1       5.200  13.600   8.800  1.00  0.00
ATOM      3    C ACE     1       4.900  14.300   9.600  1.00  0.00
ATOM      4    D ACE     1       5.600  14.200   7.900  1.00  0.00
    '''
    script = '''
from pymol import cmd
with open('output.txt', 'w') as f:
    f.write('%f' % cmd.get_dihedral('1/A', '1/B', '1/C', '1/D'))
'''
    prevdir = os.path.abspath('.')
    try:
        os.chdir(tmpdir)
        with open('xyz.pdb', 'w') as f:
            f.write(xyz)
        with open('pymolscript.py', 'w') as f:
            f.write(script)

        os.system('%s %s -cr %s' % (pymol, 'xyz.pdb', 'pymolscript.py'))
        with open('output.txt') as f:
            pymol_value = np.deg2rad(float(f.read()))
        t = md.load('xyz.pdb')
    finally:
        os.chdir(prevdir)

    mdtraj_value = md.compute_dihedrals(t, [[0, 1, 2, 3]])[0, 0]

    np.testing.assert_array_almost_equal(pymol_value, mdtraj_value)

    def calc_dihedral_energy(self, traj, improper=False, sum=True):
        """Energy for dihedral interactions

        Parameters
        ----------
        traj : mdtraj.Trajectory

        sum : bool (opt.)
            If sum=True return the total energy.
        """
        phi = md.compute_dihedrals(traj, self._dihedral_idxs)
        #if improper:
        #    phi = np.pi + md.compute_dihedrals(traj, self._dihedral_idxs) # ?
        #else:
        #    phi = -temp_phi.copy()
        #    phi[temp_phi > 0] = 2.*np.pi - temp_phi[temp_phi > 0]

        if sum:
            Edihedral = np.zeros(traj.n_frames, float)
        else:
            Edihedral = np.zeros((traj.n_frames, self.n_dihedrals), float)

        for i in range(self.n_dihedrals):
            if sum:
                Edihedral += self._dihedrals[i].V(phi[:,i])
            else:
                Edihedral[:,i] = self._dihedrals[i].V(phi[:,i])
        return Edihedral

def DFG_dihedral_states(trajectories,def_DFG):

    dihedral = []

    for traj in trajectories:
        dihedral.append(md.compute_dihedrals(traj,[def_DFG]))

    return [dihedral]

 def load_data(self):
     load_time_start = time.time()
     data = []
     for tfn in self.filenames:
         kwargs = {} if tfn.endswith('h5') else {'top': self.top}
         for t in md.iterload(tfn, chunk=self.args.split, **kwargs):
             item = np.asarray(md.compute_dihedrals(t, self.indices), np.double)
             data.append(item)
     return data

def calculate_dihedrals():
    print("Calculating dihedrals...")
    traj_files = sorted(glob.glob("traj*xtc"))
    traj = [ md.load(filename, top='structure.gro') for filename in traj_files ]
    indices = list(traj[0].topology.select('backbone'))
    dihedral_quartets = np.array([indices[i:i+4] for i in range(len(indices)-4)])
    for i in range(len(traj)):
        thetas = md.compute_dihedrals(traj[i], dihedral_quartets)
        cos_sin_dihedrals = np.hstack([np.cos(thetas), np.sin(thetas)])
        np.save('out_' + str(i) + '.npy', thetas)

    def transform(self, traj):

        # compute dihedral energy
        phi = md.compute_dihedrals(traj, self.dihedrals)
        Edih = np.array(map(lambda x,y: x(y), self.Vdih, phi.T)).T

        # compute pair energy
        r = md.compute_distances(traj, self.pairs)
        Epair = np.array(map(lambda x,y: x(y), self.Vpair, r.T)).T

        return np.hstack((Edih, Epair))

def DFG_dihedral(trajectories, def_DFG):

    dihedral = []

    for traj in trajectories:

        dihedral.append(md.compute_dihedrals(traj, [def_DFG]))

    flattened_dihedral = np.asarray([val for sublist in dihedral for val in sublist])

    return [flattened_dihedral]

    def calculate_rama_energy(self, traj, total=True):
        """Calculate the one-body burial potential

        Parameters
        ----------
        traj : mdtraj.Trajectory
            Trajectory to calculate energy over.
        sum : opt, bool
            If true (default) return the sum of the burial potentials. If
            false, return the burial energy of each individual residue.
        """
        rama = self.potential_forms["RAMA"]
        pro_rama = self.potential_forms["RAMA_PROLINE"]

        bb_traj = self.backbone_mapping.map_traj(traj)

        # where does AWSEM define 0.
        phi = md.compute_dihedrals(bb_traj, self._phi_idxs)
        psi = md.compute_dihedrals(bb_traj, self._psi_idxs)

        if total:
            Vrama = np.zeros(bb_traj.n_frames, float)
        else:
            Vrama = np.zeros((bb_traj.n_frames, self.n_phi + self.n_pro_phi), float)

        for i in range(self.n_phi):
            if total:
                Vrama += rama.V(phi[:,i], psi[:,i])
            else:
                Vrama[:, i] = rama.V(phi[:,i], psi[:,i])

        if self.n_pro_phi > 0:
            pro_phi = md.compute_dihedrals(bb_traj, self._pro_phi_idxs)
            pro_psi = md.compute_dihedrals(bb_traj, self._pro_psi_idxs)
            for i in range(self.n_pro_phi):
                if total:
                    Vrama += pro_rama.V(pro_phi[:,i], pro_psi[:,i])
                else:
                    Vrama[:,self.n_phi + i] = pro_rama.V(pro_phi[:,i], pro_psi[:,i])
        return Vrama

    def test_dihedral_op(self):
        """ Create a dihedral order parameter """
        psi_atoms = [6,8,14,16]
        dihedral_op = op.CV_MD_Function("psi", md.compute_dihedrals,
                                    indices=[psi_atoms])

        mdtraj_version = self.storage.trajectories.load(0).md()
        md_dihed = md.compute_dihedrals(mdtraj_version, indices=[psi_atoms])
        traj = self.storage.trajectories.load(0)

        my_dihed =  dihedral_op( traj )

        np.testing.assert_allclose(md_dihed, my_dihed)

    def load_data(self):
        load_time_start = time.time()
        data = []
        for tfn in self.filenames:
            kwargs = {} if tfn.endswith('h5') else {'top': self.top}
            for t in md.iterload(tfn, chunk=self.args.split, **kwargs):
                item = np.asarray(md.compute_dihedrals(t, self.indices), np.double)
                data.append(item)

        print('Loading data into memory + vectorization: %f s' %
              (time.time() - load_time_start))
        print('''Fitting with %s timeseries from %d trajectories with %d
                total observations''' % (len(data), len(self.filenames),
                                         sum(len(e) for e in data)))
        return data

def label_alanine(traj):
    ''' use dihedral angles to cluster and label alanine dipeptide
    '''
    # calculating psi and phi angles
    psi_atoms = [6,8,14,16]
    phi_atoms = [4,6,8,14]
    indices = np.asarray([phi_atoms,psi_atoms])
    dihedrals = md.compute_dihedrals(traj,indices)
#    print(dihedrals)
    trans_di = np.transpose(dihedrals)
#    print(trans_di)
    plt.scatter(trans_di[0],trans_di[1])
    plt.xlabel('phi')
    plt.ylabel('psi')
    plt.savefig("/output/tempplot")

    # deal with the periodical condition: manually put same cluster together
    def shift_phi(x):
        if x>2.2:
            return -2*np.pi+x
        else: 
            return x
    def shift_psi(x):
        if x<-2.2:
            return 2*np.pi+x 
        else:
            return x
    dihedrals_shift = np.asarray(
                [np.vectorize(shift_phi)(trans_di[0]),
                 np.vectorize(shift_psi)(trans_di[1])]) 
#    print(trans_di)
    plt.figure()
    plt.scatter(dihedrals_shift[0],dihedrals_shift[1])
    plt.xlabel('phi')
    plt.ylabel('psi')
    plt.savefig("/output/tempplot1")
    print(dihedrals.shape)
    print(dihedrals_shift.shape)
   
    # do clustering with given initial centers
    centers = np.array([[55,48],[-77,138],[-77, -39],[60, -72]])*np.pi/180.0
    clu = kmeans(n_clusters=4,init=centers)
#    labels = clu.fit_predict(dihedrals)
    labels = clu.fit_predict(np.transpose(dihedrals_shift))
    labels = 
    print("centers:")
    print(clu.cluster_centers_*180.0/np.pi)
    return labels

def test_DihedralFeaturizer_describe_features_nosincos():
    feat = DihedralFeaturizer(sincos=False)
    rnd_traj = np.random.randint(len(trajectories))
    features = feat.transform([trajectories[rnd_traj]])
    df = pd.DataFrame(feat.describe_features(trajectories[rnd_traj]))

    for f in range(25):
        f_index = np.random.choice(len(df))

        atom_inds = df.iloc[f_index].atominds
        feature_value = md.compute_dihedrals(trajectories[rnd_traj],
                                             [atom_inds])
        if feat.sincos:
            func = getattr(np, '%s' % df.iloc[f_index].otherinfo)
            feature_value = func(feature_value)

        assert (features[0][:, f_index] == feature_value.flatten()).all()