from ost import io
from promod3 import loop, modelling

# load example (has res. numbering starting at 1)
prot = io.LoadPDB('data/1CRN.pdb')
res_list = prot.residues
seqres_str = ''.join([r.one_letter_code for r in res_list])

# initialize AllAtom environment and sidechain reconstructor
env = loop.AllAtomEnv(seqres_str)
env.SetInitialEnvironment(prot)
sc_rec = modelling.SidechainReconstructor(keep_sidechains=False)
sc_rec.AttachEnvironment(env)

# reconstruct subset (res. num. 6..10)
res = sc_rec.Reconstruct(6, 5)
# reconstruct two loops (6..10 and 20..25)
res = sc_rec.Reconstruct(start_resnum_list=[6, 20],
                         num_residues_list=[5, 6],
                         chain_idx_list=[0, 0])
# update environment with solution
env.SetEnvironment(res.env_pos)
# store all positions of environment
io.SavePDB(env.GetAllAtomPositions().ToEntity(), 'sc_rec_test.pdb')