from ost import io
from ost import conop
from promod3 import loop

sequence = "AAAAAAAA"

# the backbone gets initially constructed with backbone dihedrals
# typical for helices
bb_list = loop.BackboneList(sequence)

# let's have a look at the set dihedral angles
for i in range(len(bb_list)):
    print("Looking at position %d" % i)
    if i > 0:
        print("phi: %.4f" % bb_list.GetPhiTorsion(i))
    if i < len(bb_list)-1:
        print("psi: %.4f" % bb_list.GetPsiTorsion(i))


# we now use a TorsionSampler to set random dihedral angles
torsion_sampler = loop.LoadTorsionSampler()
idx = torsion_sampler.GetHistogramIndex(conop.ALA, 
                                        conop.ALA, 
                                        conop.ALA)
for i in range(len(bb_list)):
    dihedrals = torsion_sampler.Draw(idx)
    if i > 0:
        bb_list.SetPhiTorsion(i, dihedrals[0])
    if i < len(bb_list)-1:
        bb_list.SetPsiTorsion(i, dihedrals[1])

# let's save down the randomized fragment
io.SavePDB(bb_list.ToEntity(), "randomized_fragment.pdb")