from ost import io, seq
from promod3 import loop

# load an example structure
prot = io.LoadPDB('data/1CRN.pdb')

# extract some additional information
seqres = ''.join([r.one_letter_code for r in prot.residues])
frag_pos = 35
frag_length = 9
frag_seq = seqres[frag_pos:frag_pos+frag_length]
frag_residues = prot.residues[frag_pos:frag_pos+frag_length]
ref_backbone = loop.BackboneList(frag_seq, frag_residues)

# let's load the StructureDB and a substitution matrix
db = loop.LoadStructureDB()
subst_matrix = seq.alg.BLOSUM62

# the fragger object needs to be initialized with its target sequence
fragger = loop.Fragger(frag_seq)

# we could now add an arbitrary number of parameters from different
# features for now we only go for the sequence similarity score
fragger.AddSeqSimParameters(1.0, subst_matrix)

# the Fragger can finally be filled by providing a StructureDB
fragger.Fill(db, 1.0, 100)

# let's see how good the fragments are...
below_three = 0
for i in range(len(fragger)):
    ca_rmsd = fragger[i].CARMSD(ref_backbone,True)
    print("Fragment %d has CA RMSD of %.3f" % (i, ca_rmsd))
    if ca_rmsd < 3.0:
        below_three += 1

fraction = float(below_three)/len(fragger)
print("Fraction of fragments below 3A: %.2f" % fraction)

# add into a cached map with ID based on frag_pos
fragger_map = loop.FraggerMap()
if not fragger_map.Contains(frag_pos):
	fragger_map[frag_pos] = fragger
# store it for future use
fragger_map.SaveBB("frag_map.dat")