def test_pdb_from_url():
    # load pdb from URL
    t1 = load_pdb('http://www.rcsb.org/pdb/files/4K6Q.pdb.gz')
    t2 = load_pdb('http://www.rcsb.org/pdb/files/4K6Q.pdb')
    eq(t1.n_frames, 1)
    eq(t2.n_frames, 1)
    eq(t1.n_atoms, 2208)
    eq(t2.n_atoms, 2208)

def loader(file):

    try:
        md.load_pdb(file)
    except:
        return None
            
    print(file)
    
    return None

    def _construct_traj(self):
        logger.debug('Loading Trajectory object for model {0} ({1}/{2})'.format(self.df.templateid.iloc[0], 0, len(self.df.model_filepath)))
        traj = mdtraj.load_pdb(self.df.model_filepath[0])
        remove_disulfide_bonds_from_topology(traj.topology)
        self.traj = traj

        for m, model_filepath in enumerate(self.df.model_filepath[1:]):
            logger.debug('Loading Trajectory object for model {0} ({1}/{2})'.format(self.df.templateid.iloc[m+1], m+1, len(self.df.model_filepath)))
            traj = mdtraj.load_pdb(model_filepath)
            remove_disulfide_bonds_from_topology(traj.topology)
            self.traj += traj

def test_1vii_url_and_gz():
    t1 = load_pdb('http://www.rcsb.org/pdb/files/1vii.pdb.gz')
    t2 = load_pdb('http://www.rcsb.org/pdb/files/1vii.pdb')
    t3 = load_pdb(get_fn('1vii.pdb.gz'))
    t4 = load_pdb(get_fn('1vii.pdb'))
    eq(t1.n_frames, 1)
    eq(t1.n_frames, t2.n_frames)
    eq(t1.n_frames, t3.n_frames)
    eq(t1.n_frames, t4.n_frames)
    
    eq(t1.n_atoms, t2.n_atoms)
    eq(t1.n_atoms, t3.n_atoms)
    eq(t1.n_atoms, t4.n_atoms)

def test_segment_id(get_fn):
    pdb = load_pdb(get_fn('ala_ala_ala.pdb'))
    pdb.save_pdb(temp)
    pdb2 = load_pdb(temp)

    correct_segment_id = 'AAL'
    # check that all segment ids are set correctly
    for ridx,r in enumerate(pdb.top.residues):
        assert r.segment_id == correct_segment_id, "residue %i (0-indexed) does not have segment_id set correctly from ala_ala_ala.pdb"%(ridx)

    # check that all segment ids are set correctly after a new pdb file is written
    for ridx,(r1,r2) in enumerate(zip(pdb.top.residues,pdb2.top.residues)):
        assert r1.segment_id == r2.segment_id, "segment_id of residue %i (0-indexed) in ala_ala_ala.pdb does not agree with value in after being written out to a new pdb file"%(ridx)

    def _get_models(self):
        self.model = {}
        root, dirnames, filenames = next(os.walk(self.models_target_dir))
        for dirname in dirnames:
            if 'implicit' in self.model and 'explicit' in self.model:
                break
            if 'implicit' not in self.model:
                implicit_model_filename = os.path.join(self.models_target_dir, dirname, 'implicit-refined.pdb.gz')
                if os.path.exists(implicit_model_filename):
                    self.model['implicit'] = mdtraj.load_pdb(implicit_model_filename)

            if 'explicit' not in self.model:
                explicit_model_filename = os.path.join(self.models_target_dir, dirname, 'explicit-refined.pdb.gz')
                if os.path.exists(explicit_model_filename):
                    self.model['explicit'] = mdtraj.load_pdb(explicit_model_filename)

def test_3(tmpdir):
    # 1COY gives a small error, due to a broken chain.
    pdbids = ['1GAI', '6gsv', '2AAC']
    for pdbid in pdbids:
        t = md.load_pdb('http://www.rcsb.org/pdb/files/%s.pdb' % pdbid)
        t = t.atom_slice(t.top.select_atom_indices('minimal'))
        assert_(call_dssp(tmpdir, t), md.compute_dssp(t, simplified=False)[0])

def metal_searcher(file):
    
    metal_searcher_results = {}
    metal_searcher_results['one_atom_residue_names'] = []
    metal_searcher_results['ok_file_count'] = 0
    metal_searcher_results['error_file_count'] = 0
    
    print(file)
    
    try:
        traj = md.load_pdb(file)
        metal_searcher_results['ok_file_count'] = 1
    except:
        metal_searcher_results['error_file_count'] = 1
        return metal_searcher_results
        
    topo = traj.topology
    
    one_atom_residue_atoms = [atom for atom in topo.atoms if atom.name == atom.residue.name]
    
    if one_atom_residue_atoms:
        for atom in one_atom_residue_atoms:
            metal_searcher_results['one_atom_residue_names'].append(str(atom.name))
            
    print(metal_searcher_results)        
    return metal_searcher_results

def loader(file):

    print(file)
    try:
        traj = md.load_pdb(file)
    except:    
        return file

def metal_searcher(file):
    
    metal_searcher_results = {}
    
    metal_searcher_results['ok_file_count'] = 0
    metal_searcher_results['error_file_count'] = 0
    metal_searcher_results['metals_in_file'] = {}
    metal_searcher_results['any_metal_in_file'] = 0
    
    for i in metals_list:
        metal_searcher_results['metals_in_file'][i] = 0
    
    print(file)
    
    try:
        traj = md.load_pdb(file)
        metal_searcher_results['ok_file_count'] = 1
    except:
        metal_searcher_results['error_file_count'] = 1
        return metal_searcher_results
        
    topo = traj.topology
    
    one_atom_residue_atoms = [atom.name for atom in topo.atoms if atom.name == atom.residue.name and atom.name in metals_list]
    
    if one_atom_residue_atoms:
        metal_searcher_results['any_metal_in_file'] = 1
        for atom_name in one_atom_residue_atoms:
            metal_searcher_results['metals_in_file'][atom_name] += 1
            
    print(metal_searcher_results)        
    return metal_searcher_results

    def get(self, pdbid):
        pdbid = pdbid.lower()

        if len(pdbid) != 4:
            raise Exception("pdb id must be four characters long")

        filename = os.path.join(self.temp_dir, '%s.pdb.gz' % pdbid)
        with open(filename, 'wb') as filehandler:
            self.conn.retrbinary('RETR pub/pdb/data/structures/divided/pdb/%s/pdb%s.ent.gz' % (pdbid[1:3], pdbid),
                filehandler.write)

        if self.just_files:
            return filename

        else:
            # then we actually want to load it as a pdb file object and return an mdtraj trajectory
            with gzip.open(filename) as filehandler:
                text = filehandler.read()

            with open(filename[:-3], 'w') as filehandler:
                filehandler.write(text)

            pdb = mdtraj.load_pdb(filename[:-3])

            return pdb

def test_refine_explicit_md_short():
    with integrationtest_context(set_up_project_stage='solvated'):
        targetid = 'EGFR_HUMAN_D0'
        templateid = 'KC1D_HUMAN_D0_4KB8_D'
        refine_explicit_md(
            process_only_these_targets=[targetid],
            process_only_these_templates=[templateid],
            sim_length=2.0*unit.femtosecond,
            nsteps_per_iteration=1,
            verbose=True
        )
        explicit_metadata_filepath = os.path.join(
            default_project_dirnames.models, targetid, 'refine_explicit_md-meta0.yaml'
        )
        explicit_model_filepath = os.path.join(
            default_project_dirnames.models, targetid, templateid, 'explicit-refined.pdb.gz'
        )
        explicit_energies_filepath = os.path.join(
            default_project_dirnames.models, targetid, templateid, 'explicit-energies.txt'
        )
        explicit_log_filepath = os.path.join(
            default_project_dirnames.models, targetid, templateid, 'explicit-log.yaml'
        )

        assert all(map(
            os.path.exists,
            [explicit_model_filepath, explicit_energies_filepath, explicit_log_filepath]
        ))
        with open(explicit_log_filepath) as explicit_log_file:
            explicit_log = yaml.load(explicit_log_file)
        assert explicit_log.get('finished') is True
        assert explicit_log.get('successful') is True
        explicit_model_traj = mdtraj.load_pdb(explicit_model_filepath)

def test_1():
    for fn in ['1bpi.pdb', '1vii.pdb', '4K6Q.pdb', '1am7_protein.pdb']:
        t = md.load_pdb(get_fn(fn))
        t = t.atom_slice(t.top.select_atom_indices('minimal'))
        f = lambda : assert_(call_dssp(t), md.compute_dssp(t, simplified=False)[0])
        f.description = 'test_1: %s' % fn
        yield f

def test_4():
    t = md.load_pdb(get_fn('1am7_protein.pdb'))
    a = md.compute_dssp(t, simplified=True)
    b = md.compute_dssp(t, simplified=False)
    assert len(a) == len(b)
    assert len(a[0]) == len(b[0])
    assert list(np.unique(a[0])) == ['C', 'E', 'H']

    def __init__(self, PDB_filename, free_energy, expdata_filename=None, use_log_normal_distances=False,
                       dloggamma=np.log(1.01), gamma_min=0.2, gamma_max=10.0):
        """Initialize the class.
        INPUTS
	conf		A molecular structure as an msmbuilder Conformation() object.
                        NOTE: For cases where the structure is an ensemble (say, from clustering)
                        and the modeled NOE distances and coupling constants are averaged, 
                        the structure itself can just be a placeholder with the right atom names
                        and numbering
              
        free_energy     The (reduced) free energy f = beta*F of this conformation
        """

        self.PDB_filename = PDB_filename
        self.expdata_filename = expdata_filename
        self.conf = mdtraj.load_pdb(PDB_filename)
        # Convert the coordinates from nm to Angstrom units 
        self.conf.xyz = self.conf.xyz*10.0 

        # The (reduced) free energy f = beta*F of this structure, as predicted by modeling
        self.free_energy = free_energy

        # Flag to use log-normal distance errors log(d/d0)
        self.use_log_normal_distances = use_log_normal_distances

        # Store info about gamma^(-1/6) scaling  parameter array
        self.dloggamma = dloggamma
        self.gamma_min = gamma_min
        self.gamma_max = gamma_max
        self.allowed_gamma = np.exp(np.arange(np.log(self.gamma_min), np.log(self.gamma_max), self.dloggamma))

        # Store distance restraint info
        self.distance_restraints = []
        self.distance_equivalency_groups = {}
        self.ambiguous_groups = []  # list of pairs of group indices, e.g.:   [ [[1,2,3],[4,5,6]],   [[7],[8]], ...]
        self.ndistances = 0

        # Store dihedral restraint info
        self.dihedral_restraints = [] 
        self.dihedral_equivalency_groups = {}
        self.dihedral_ambiguity_groups = {}
        self.ndihedrals = 0

        # Create a KarplusRelation object
        self.karplus = KarplusRelation()

        # variables to store pre-computed SSE and effective degrees of freedom (d.o.f.)
        self.sse_distances = np.array([0.0 for gamma in self.allowed_gamma])
        self.Ndof_distances = None
        self.sse_dihedrals = None
        self.Ndof_dihedrals = None
        self.betas = None   # if reference is used, an array of N_j betas for each distance
        self.neglog_reference_priors = None
        self.sum_neglog_reference_priors = 0.0

        # If an experimental data file is given, load in the information
        self.expdata_filename = expdata_filename
        if expdata_filename != None:
        	self.load_expdata(expdata_filename)

def test_topology(get_fn):
    top = md.load_pdb(get_fn('native.pdb')).topology

    with HDF5TrajectoryFile(temp, 'w') as f:
        f.topology = top

    with HDF5TrajectoryFile(temp) as f:
        assert f.topology == top

def test_load_trajectory(get_fn):
    # Compare a TNG file to the PDB file it was created from.

    pdbtraj = md.load_pdb(get_fn('frame0.pdb'))
    tngtraj = md.load_tng(get_fn('frame0.tng'), top=pdbtraj.topology)
    eq(pdbtraj.n_frames, tngtraj.n_frames)
    eq(pdbtraj.unitcell_vectors, tngtraj.unitcell_vectors)
    eq(pdbtraj.xyz, tngtraj.xyz)

 def _store_highest_seqid_model(self):
     models_sorted = self.df.sort('seqid').templateid
     for modelid in models_sorted:
         model_filepath = os.path.join(self.models_target_dir, modelid, self.model_filename)
         if os.path.exists(model_filepath):
             self.ref_modelid = modelid
             self.ref_model_filepath = model_filepath
             break
     self.ref_model_traj = mdtraj.load_pdb(self.ref_model_filepath)

 def _align_structures(self):
     for irow in range(self.nrows):
         row_first_col = irow * self.ncols
         self._ref_traj = mdtraj.load_pdb(self.structure_filepaths[row_first_col])
         self._heavy_atoms = self._ref_traj.topology.select('not name H')
         for icol in range(self.ncols):
             alignment_index = row_first_col + icol
             structure_filepath = self.structure_filepaths[alignment_index]
             self._align_structure(structure_filepath, alignment_index)

def test_3nch_conect():
    # This has conect entries that use all available digits, good failure case.
    t1 = load_pdb(get_fn('3nch.pdb.gz'))
    top, bonds = t1.top.to_dataframe()
    bonds = dict(((a, b), 1) for (a, b) in bonds)
    eq(bonds[19782, 19783], 1)  # Check that last SO4 molecule has right bonds
    eq(bonds[19782, 19784], 1)  # Check that last SO4 molecule has right bonds
    eq(bonds[19782, 19785], 1)  # Check that last SO4 molecule has right bonds
    eq(bonds[19782, 19786], 1)  # Check that last SO4 molecule has right bonds