# Copyright (c) 2013-2020, SIB - Swiss Institute of Bioinformatics and
#                          Biozentrum - University of Basel
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#   http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from ost import geom, mol, conop
from promod3 import core, sidechain
import traceback

###############################################################################
# helper functions
def _GetRotamerIDs(res_list):
    '''Return list (length = len(res_list)) of rotamer IDs for all residues.'''
    rotamer_ids = [sidechain.ALA] * len(res_list)
    for i,r in enumerate(res_list):
        rot_id = sidechain.TLCToRotID(r.GetName())
        if rot_id == sidechain.XXX:
            continue # no idea what it is, so we stick with ALA 
        rotamer_ids[i] = rot_id
    return rotamer_ids

def _GetPhiAngle(r):
    '''Extract phi angle for residue r.'''
    # def. fallback = helix
    phi = -1.0472
    # try to get phi from torsion angles
    tor = r.GetPhiTorsion()
    if tor.IsValid():
        phi = tor.GetAngle()
    else:
        r_prev = r.handle.prev
        if r_prev.IsValid() and mol.InSequence(r_prev, r.handle):
            c_prev = r_prev.FindAtom("C")
            n = r.FindAtom("N")
            ca = r.FindAtom("CA")
            c = r.FindAtom("C")
            if c_prev.IsValid() and n.IsValid() and ca.IsValid() and c.IsValid():
                phi = geom.DihedralAngle(c_prev.GetPos(),n.GetPos(),
                                         ca.GetPos(),c.GetPos())
    return phi

def _GetPsiAngle(r):
    '''Extract psi angle for residue r.'''
    # def. fallback = helix 
    psi = -0.7854
    # try to get psi from torsion angles
    tor = r.GetPsiTorsion()
    if tor.IsValid():
        psi = tor.GetAngle()
    else:
        r_next = r.handle.next
        if r_next.IsValid() and mol.InSequence(r.handle, r_next):
            n = r.FindAtom("N")
            ca = r.FindAtom("CA")
            c = r.FindAtom("C")
            n_next = r_next.FindAtom("N")
            if n.IsValid() and ca.IsValid() and c.IsValid() and n_next.IsValid():
                psi = geom.DihedralAngle(n.GetPos(), ca.GetPos(),
                                         c.GetPos(), n_next.GetPos())
    return psi

def _GetDihedrals(res_list):
    '''Extract dihedral angles for all residues.
    Returns phi and psi angles as 2 lists with length = len(res_list).
    '''
    prof_name = 'sidechain::_GetDihedrals'
    prof = core.StaticRuntimeProfiler.StartScoped(prof_name)
    phi_angles = [0.0] * len(res_list)
    psi_angles = [0.0] * len(res_list)
    for i,r in enumerate(res_list):
        phi_angles[i] = _GetPhiAngle(r)
        psi_angles[i] = _GetPsiAngle(r)
    return phi_angles, psi_angles

def _AddBackboneFrameResidues(frame_residues, res_list, rotamer_ids,
                              rotamer_constructor, phi_angles, psi_angles,
                              n_ter_residues, c_ter_residues):
    '''Update frame_residues (list) with BackboneFrameResidues for res_list.'''
    for i,r in enumerate(res_list):
        try:
            is_n_ter = r.handle.GetHashCode() in n_ter_residues
            is_c_ter = r.handle.GetHashCode() in c_ter_residues
            frame_residue = rotamer_constructor.ConstructBackboneFrameResidue(\
                                r.handle, rotamer_ids[i], i,
                                phi_angles[i], psi_angles[i], 
                                is_n_ter, is_c_ter)
            frame_residues.append(frame_residue)
        except:
            continue

def _AddLigandFrameResidues(frame_residues, ent_lig, rotamer_constructor, offset):
    '''Update frame_residues (list) with FrameResidues for res. in ent_lig.
    Set offset >= number of non-ligand residues (used for residue_index).
    '''
    # parse ligand residues
    for i, res in enumerate(ent_lig.residues):
        res_idx = offset + i
        is_done = False
        # special treatment for peptides
        if res.IsPeptideLinking():
            rot_id = sidechain.TLCToRotID(res.GetName())
            if rot_id != sidechain.XXX:
                # get more info
                phi = _GetPhiAngle(res)
                psi = _GetPsiAngle(res)
                r_prev = res.handle.prev
                n_ter = not r_prev.IsValid() \
                        or not mol.InSequence(r_prev, res.handle)
                r_next = res.handle.next
                c_ter = not r_next.IsValid() \
                        or not mol.InSequence(res.handle, r_next)
                # try to add frame residues (ignore exceptions)
                try:
                    fr1 = rotamer_constructor.ConstructBackboneFrameResidue(\
                              res.handle, rot_id, res_idx,
                              phi, psi, n_ter, c_ter)

                    fr2 = rotamer_constructor.ConstructSidechainFrameResidue(\
                              res.handle, rot_id, res_idx, phi, psi, n_ter, c_ter)

                    frame_residues.extend([fr1,fr2])
                except:
                    pass   # ignore peptide treatment and treat below
                else:
                    is_done = True
        # if it failed, treat it as an unknown entity
        if not is_done:
            # try to add frame residues (skip exceptions)
            try:
                # NOTES:
                # - ConstructFrameResidueHeuristic has fall back if res unknown
                # - it only deals with few possible ligand cases and has not
                #   been tested extensively! 
                comp_lib = conop.GetDefaultLib()
                fr = rotamer_constructor.ConstructFrameResidueHeuristic(\
                         res.handle, res_idx, comp_lib)
                frame_residues.append(fr)
            except:
                continue

def _AddSidechainFrameResidues(frame_residues, incomplete_sidechains,
                               keep_sidechains, res_list, rotamer_ids,
                               phi_angles, psi_angles,
                               rotamer_constructor, cystein_indices,
                               n_ter_residues, c_ter_residues):
    '''Update frame_residues (list) with SidechainFrameResidues for res_list,
    incomplete_sidechains (list of indices) with sidechains to be constructed,
    and (if given) cystein_indices (list of indices) with all CYS (appended).
    Each residue can only end up in one of the 3 lists.
    '''
    if keep_sidechains:
        # try to generate frame residues for all existing side chains
        # skip non-existing sidechains and CYS (if cystein_indices) and update
        # incomplete_sidechains and cystein_indices
        for i,r in enumerate(res_list):

            if cystein_indices is not None and rotamer_ids[i] == sidechain.CYS:
                cystein_indices.append(i)
                continue

            try:
                is_n_ter = r.handle.GetHashCode() in n_ter_residues
                is_c_ter = r.handle.GetHashCode() in c_ter_residues
                frame_residue = rotamer_constructor.ConstructSidechainFrameResidue(\
                                    r.handle, rotamer_ids[i], i, phi_angles[i],
                                    psi_angles[i], is_n_ter, is_c_ter)
                frame_residues.append(frame_residue)
            except Exception as e:
                incomplete_sidechains.append(i)
    else:
        # no frame residues to create, just update incomplete_sidechains
        # and cystein_indices if needed
        for i,r in enumerate(res_list):

            if cystein_indices is not None and rotamer_ids[i] == sidechain.CYS:
                cystein_indices.append(i)
                continue

            incomplete_sidechains.append(i)

def _AddCysteinFrameResidues(frame_residues, incomplete_sidechains,
                             keep_sidechains, res_list, rotamer_ids,
                             phi_angles, psi_angles,
                             rot_constructor, cystein_indices,
                             disulfid_indices, disulfid_rotamers,
                             n_ter_residues, c_ter_residues):
    '''Update frame_residues (list) with cysteins.
    Parameters as in _AddSidechainFrameResidues.
    Some cysteins (in disulfid_indices) get special treatment as disulfid
    bridges (disulfid_indices, disulfid_rotamers from _GetDisulfidBridges).
    '''
    # handle cysteins participating in a disulfid bond
    for cys_idx, cys_rot in zip(disulfid_indices, disulfid_rotamers):
        # add FrameResidue
        frame_residues.append(cys_rot.ToFrameResidue(cys_idx))
        # set the position in the proteins residues
        cys_rot.ApplyOnResidue(res_list[cys_idx].handle,
                               consider_hydrogens=False)
        sidechain.ConnectSidechain(res_list[cys_idx].handle, sidechain.CYS)

    # add remaining ones according the given flags
    for idx in cystein_indices:
        if idx in disulfid_indices:
            continue # already handled
        if keep_sidechains:
            try:
                is_n_ter = res_list[idx].handle.GetHashCode() in n_ter_residues
                is_c_ter = res_list[idx].handle.GetHashCode() in c_ter_residues
                frame_residue = rot_constructor.ConstructSidechainFrameResidue(\
                                    res_list[idx].handle, rotamer_ids[idx],
                                    idx, phi_angles[idx], psi_angles[idx],
                                    is_n_ter, is_c_ter)
                frame_residues.append(frame_residue)
            except:
                incomplete_sidechains.append(idx) 
        else:
            incomplete_sidechains.append(idx)

def _GetRotamerGroup(res_handle, rot_id, res_idx, rot_lib, rot_constructor,
                     phi, psi, use_frm, bbdep, probability_cutoff,
                     n_ter_residues, c_ter_residues):

    is_n_ter = res_handle.handle.GetHashCode() in n_ter_residues 
    is_c_ter = res_handle.handle.GetHashCode() in c_ter_residues 

    if use_frm:
        return rot_constructor.ConstructFRMRotamerGroup(res_handle, rot_id,
                                                        res_idx, rot_lib,
                                                        phi, psi, 
                                                        is_n_ter, is_c_ter,
                                                        probability_cutoff)
    else:        
        return rot_constructor.ConstructRRMRotamerGroup(res_handle, rot_id,
                                                        res_idx, rot_lib,
                                                        phi, psi,
                                                        is_n_ter, is_c_ter,
                                                        probability_cutoff)


def _GetRotamerGroups(res_list, rot_ids, indices, rot_lib, rot_constructor,
                      phi_angles, psi_angles, use_frm, bbdep, frame_residues,
                      n_ter_residues, c_ter_residues):
    '''Get list of rotamer groups from subset of res_list.
    Residues are chosen as res_list[i] for i in indices and only if a rotamer
    group can be created.
    Rotamer groups are filtered to keep only best ones (given frame).
    Returns list of rotamer groups and list of res. indices they belong to.
    '''
    prof_name = 'sidechain::_GetRotamerGroups'
    prof = core.StaticRuntimeProfiler.StartScoped(prof_name)

    # res.index (res_list[i]) for each modelled sc
    residues_with_rotamer_group = list()
    #  linked to residue in residues_with_rotamer_group
    rotamer_groups = list()
    # get frame for score evaluation
    frame = sidechain.Frame(frame_residues)
    # build rotamers for chosen sidechains
    for i in indices:
        # get rotamer ID
        r = res_list[i]
        rot_id = rot_ids[i]

        if rot_id == sidechain.CYS:
            rot_id = sidechain.CYH

        if rot_id == sidechain.PRO:
            tor = r.GetOmegaTorsion()
            omega = None
            if tor.IsValid():
                omega = tor.GetAngle()
            elif i > 0:
                # fallback computation of omega as in OST-code
                prev = res_list[i-1]
                if prev.IsValid() and prev.IsPeptideLinking():
                    ca_prev = prev.FindAtom("CA")
                    c_prev = prev.FindAtom("C")
                    n = r.FindAtom("N")
                    ca = r.FindAtom("CA")
                    valid = ca_prev.IsValid() and c_prev.IsValid() \
                            and n.IsValid() and ca.IsValid()
                    if valid and mol.BondExists(c_prev.handle, n.handle):
                        omega = geom.DihedralAngle(ca_prev.GetPos(),
                                                   c_prev.GetPos(),
                                                   n.GetPos(), ca.GetPos())
            # omega not set if prev. res. missing
            if omega is not None:
                if abs(omega) < 1.57:
                    rot_id = sidechain.CPR
                else:
                    rot_id = sidechain.TPR

        # get RotamerGroup
        try:
            rot_group = _GetRotamerGroup(r.handle, rot_id, i, rot_lib,
                                         rot_constructor, phi_angles[i],
                                         psi_angles[i], use_frm, bbdep, 0.98,
                                         n_ter_residues, c_ter_residues)
        except:
            continue
        # keep best ones
        rot_group.SetFrameEnergy(frame)
        rot_group.ApplySelfEnergyThresh()
        rotamer_groups.append(rot_group)
        residues_with_rotamer_group.append(i)

    return rotamer_groups, residues_with_rotamer_group

def _GetDisulfidBridges(frame_residues, keep_sidechains, cystein_indices, 
                        res_list, rotamer_library, use_frm, bbdep, 
                        rotamer_constructor, phi_angles, psi_angles,
                        n_ter_residues, c_ter_residues):
    '''Get disulfid bridges for CYS and according rotamers.
    CYS are identified by by items in cystein_indices (into res_list).
    Returns: disulfid_indices: list of res. index in bridge,
             disulfid_rotamers: list of rotamers (best one for bridge).
    '''
    # this is required for the disulfid score evaluation
    frame = sidechain.Frame(frame_residues)

    # some info we have to keep track of when evaluating disulfid bonds
    cystein_rot_groups = list()
    cys_ca_positions = list()
    cys_cb_positions = list()
    indices_with_rot_groups = list()

    for i in cystein_indices:
        # check ca, cb
        r = res_list[i]
        ca = r.FindAtom("CA")
        cb = r.FindAtom("CB")
        if not (ca.IsValid() and cb.IsValid()):
            continue
        cys_ca_positions.append(ca.GetPos())
        cys_cb_positions.append(cb.GetPos())

        # get RotamerGroup in case of cysteins, we do FRM in any case...
        # If we're suposed to keep the existing sidechains and the processed
        # cystein contains all required atoms, we manually construct an FRM
        # rotamer to still allow the rotamer to enter the disulfid bridge 
        # resolving algorithm.  
        rot_group = None
        if keep_sidechains:
            sg = r.FindAtom("SG")
            if sg.IsValid():
                # we're constructing the required particle through a frame 
                # residue in the rotamer constructor
                phi = _GetPhiAngle(r)
                psi = _GetPsiAngle(r)
                is_n_ter = r.handle.GetHashCode() in n_ter_residues
                is_c_ter = r.handle.GetHashCode() in c_ter_residues
                cys_frame_res = \
                rotamer_constructor.ConstructSidechainFrameResidue(r.handle, 
                                                                   sidechain.CYD, 
                                                                   0, phi, psi,
                                                                   is_n_ter, is_c_ter)
                for j in range(len(cys_frame_res)):
                    if cys_frame_res[j].GetName() == "SG":
                        particle_list = [cys_frame_res[j]]
                        # The temperature and self internal_e_prefactor 
                        # parameter have been copied from the 
                        # SCWRL4RotamerConstructor. Hacky...
                        frm_rotamer = sidechain.FRMRotamer(particle_list, 
                                                           1.69, 1.0, 4.07)
                        frm_rotamer.AddSubrotamerDefinition([0])               
                        frm_rotamer.SetInternalEnergy(0.0)
                        rot_group = sidechain.FRMRotamerGroup([frm_rotamer], i)
                        break
                if rot_group == None:
                    raise RuntimeError("Could not find SG atom in CYS frame res")

        if rot_group == None:
            rot_group = _GetRotamerGroup(r.handle, sidechain.CYD, i, 
                                         rotamer_library,
                                         rotamer_constructor, 
                                         phi_angles[i],
                                         psi_angles[i], True, bbdep, 0.98,
                                         n_ter_residues, c_ter_residues)

        rot_group.AddFrameEnergy(frame)
        cystein_rot_groups.append(rot_group)
        indices_with_rot_groups.append(i)

    bond_result, \
    rotamer_result = sidechain.ResolveCysteins(cystein_rot_groups, 
                                               cys_ca_positions,
                                               cys_cb_positions, 45.0,
                                               True)

    # get CYS with disulfid bonds and the chosen rotamers
    disulfid_indices = list()
    disulfid_rotamers = list()
  
    for a, b in zip(bond_result, rotamer_result):
        disulfid_indices.append(indices_with_rot_groups[a[0]])
        disulfid_indices.append(indices_with_rot_groups[a[1]])
        disulfid_rotamers.append(cystein_rot_groups[a[0]][b[0]])
        disulfid_rotamers.append(cystein_rot_groups[a[1]][b[1]])

    return disulfid_indices, disulfid_rotamers


def _AddCB(prot):
    '''Checks if every residue has a CB atom. Constructs them if necessary.'''
    edi = prot.handle.EditXCS(mol.BUFFERED_EDIT)
    for r in prot.residues:
        olc = r.one_letter_code
        if olc not in ['G', 'g']:
            cb = r.handle.FindAtom('CB')
            if not cb.IsValid():
                n = r.handle.FindAtom('N')
                ca = r.handle.FindAtom('CA')
                c = r.handle.FindAtom('C')
                if n.IsValid() and ca.IsValid() and c.IsValid():
                    cb_pos = core.ConstructCBetaPos(n.GetPos(), ca.GetPos(), 
                                                    c.GetPos())
                    cb = edi.InsertAtom(r.handle, "CB", cb_pos, "C");
                    # prot is a view on the underlying handle, so we also have
                    # to add it to the view
                    r.AddAtom(cb)


###############################################################################

def ReconstructSidechains(ent, keep_sidechains=False, build_disulfids=True,
                          rotamer_model="frm", consider_ligands=True, 
                          rotamer_library=None, optimize_subrotamers=True,
                          graph_max_complexity=100000000, 
                          graph_initial_epsilon=0.02, 
                          energy_function = "SCWRL4"):
    '''Reconstruct sidechains for the given structure.

    :param ent: Structure for sidechain reconstruction. Note, that the sidechain
                reconstruction gets directly applied on the structure itself.
    :type ent:  :class:`ost.mol.EntityHandle`

    :param keep_sidechains: Flag, whether complete sidechains in *ent* (i.e. 
                            containing all required atoms) should be kept rigid
                            and directly be added to the frame.
    :type keep_sidechains:  :class:`bool`

    :param build_disulfids: Flag, whether possible disulfid bonds should be 
                            searched. If a disulfid bond is found, the two
                            participating cysteins are fixed and added to
                            the frame.
    :type build_disulfids:  :class:`bool`

    :param rotamer_model: Rotamer model to be used, can either be "frm" or "rrm"
    :type rotamer_model:  :class:`str`

    :param consider_ligands: Flag, whether to add ligands (anything in chain
                             '_') as static objects.
    :type consider_ligands:  :class:`bool`

    :param rotamer_library: A rotamer library to extract the rotamers from. The
                            default is to call :meth:`<LoadBBDepLib>`.
    :type rotamer_library:  :class:`BBDepRotamerLib` / :class:`RotamerLib`

    :param optimize_subrotamers: Only considered when *rotamer_model*
                                 is "frm".
                                 If set to True, the FRM solution undergoes 
                                 some postprocessing by calling 
                                 :func:`SubrotamerOptimizer` with default 
                                 parametrization.
    :type optimize_subrotamers: :class:`bool`

    :param graph_max_complexity: Max. complexity for
                                 :meth:`RotamerGraph.TreeSolve`.
    :type graph_max_complexity:  :class:`int`
    :param graph_intial_epsilon: Initial epsilon for
                                 :meth:`RotamerGraph.TreeSolve`.
    :type graph_intial_epsilon:  :class:`float`
    :param energy_function:      What energy function to use can be any in 
                                 ["SCWRL4", "SCWRL3", "VINA"]
    :type energy_function:       :class:`str`
    '''
    prof_name = 'modelling::ReconstructSidechains'
    prof = core.StaticRuntimeProfiler.StartScoped(prof_name)

    # setup settings
    if rotamer_model.lower() == "frm":
        use_frm = True
    elif rotamer_model.lower() == "rrm":
        use_frm = False
    else:
        raise RuntimeError("Only \"rrm\" and \"frm\" allowed for rotamer_model!")

    if rotamer_library == None: 
        rotamer_library = sidechain.LoadBBDepLib()
    bbdep = False
    if type(rotamer_library) is sidechain.BBDepRotamerLib:
        bbdep = True

    rotamer_constructor = None

    if energy_function == "SCWRL4":
        rotamer_constructor = sidechain.SCWRL4RotamerConstructor(False)
    elif energy_function == "SCWRL3":
        rotamer_constructor = sidechain.SCWRL3RotamerConstructor(False)
    elif energy_function == "VINA":
        rotamer_constructor = sidechain.VINARotamerConstructor(False)
    else:
        raise RuntimeError("Only \"SCWRL4\", \"SCWRL3\" and \"VINA\" allowed "\
                           "for energy_function")

    if rotamer_constructor == None:
        raise RuntimeError("Invalid Energy function to reconstruct sidechains!")

    # take out ligand chain and any non-peptides
    prot = ent.Select("peptide=true and cname!='_'")

    # make sure that we have all CB atoms
    _AddCB(prot)
    
    # parse residues (all lists of length len(prot.residues))
    rotamer_ids = _GetRotamerIDs(prot.residues)
    phi_angles, psi_angles = _GetDihedrals(prot.residues)

    # set nter and cter (needed in _AddBackboneFrameResidues)
    n_ter_residues = set()
    c_ter_residues = set()
    for c in prot.chains:
        n_ter_residues.add(c.residues[0].handle.GetHashCode())
        c_ter_residues.add(c.residues[-1].handle.GetHashCode())

    # build up frame
    frame_residues = list()         # list of frame residues connected to frame
    incomplete_sidechains = list()  # residue indices
    _AddBackboneFrameResidues(frame_residues, prot.residues, rotamer_ids,
                              rotamer_constructor, phi_angles, psi_angles,
                              n_ter_residues, c_ter_residues)
    
    # add ligands?
    if consider_ligands:
        ligs = ent.Select("cname='_'")
        offset = len(prot.residues)
        _AddLigandFrameResidues(frame_residues, ligs, rotamer_constructor, offset)

    # check special handling of cysteins
    if build_disulfids:
        # residue indices of cysteins
        cystein_indices = list()
        # update frame_residues, incomplete_sidechains, cystein_indices
        _AddSidechainFrameResidues(frame_residues, incomplete_sidechains,
                                   keep_sidechains, prot.residues, rotamer_ids,
                                   phi_angles, psi_angles,
                                   rotamer_constructor, cystein_indices,
                                   n_ter_residues, c_ter_residues)
        # update frame_residues, incomplete_sidechains with cysteins (if needed)
        if len(cystein_indices) > 0:
            # get disulfid bridges and according rotamers
            disulfid_indices, disulfid_rotamers = \
                _GetDisulfidBridges(frame_residues, keep_sidechains, 
                                    cystein_indices, prot.residues,
                                    rotamer_library, use_frm, bbdep,
                                    rotamer_constructor, phi_angles, psi_angles,
                                    n_ter_residues, c_ter_residues)
            # update frame_residues, incomplete_sidechains
            _AddCysteinFrameResidues(frame_residues, incomplete_sidechains,
                                     keep_sidechains, prot.residues, rotamer_ids,
                                     phi_angles, psi_angles,
                                     rotamer_constructor, cystein_indices,
                                     disulfid_indices, disulfid_rotamers,
                                     n_ter_residues, c_ter_residues)
    else:
        # update frame_residues, incomplete_sidechains
        _AddSidechainFrameResidues(frame_residues, incomplete_sidechains,
                                   keep_sidechains, prot.residues, rotamer_ids,
                                   phi_angles, psi_angles,
                                   rotamer_constructor, None,
                                   n_ter_residues, c_ter_residues)
    
    # get rotamer groups and residues they're linked to
    rotamer_groups, residues_with_rotamer_group = \
        _GetRotamerGroups(prot.residues, rotamer_ids, incomplete_sidechains,
                          rotamer_library, rotamer_constructor, phi_angles,
                          psi_angles, use_frm, bbdep, frame_residues,
                          n_ter_residues, c_ter_residues)

    # set up graph and solve to get best rotamers
    if use_frm:
        graph = sidechain.RotamerGraph.CreateFromFRMList(rotamer_groups)
    else:
        graph = sidechain.RotamerGraph.CreateFromRRMList(rotamer_groups)

    solution = graph.TreeSolve(max_complexity=graph_max_complexity,
                               initial_epsilon=graph_initial_epsilon)[0]

    if use_frm and optimize_subrotamers:
        rotamers_to_optimize = list()

        for rot_group, sol in zip(rotamer_groups, solution): 
            rotamers_to_optimize.append(rot_group[sol])
        sidechain.SubrotamerOptimizer(rotamers_to_optimize)

    # update structure
    for i, rot_group, sol in zip(residues_with_rotamer_group, rotamer_groups,
                                 solution):
        try:
            res_handle = prot.residues[i].handle
            rot_group[sol].ApplyOnResidue(res_handle, consider_hydrogens=False)
            sidechain.ConnectSidechain(res_handle, rotamer_ids[i])
        except:
            print("there is a backbone atom missing... ", \
                  res_handle.GetQualifiedName())

# these methods will be exported into module
__all__ = ('ReconstructSidechains',)