# 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.


'''
Automatically (re-)construct sidechains in a protein model.

Example usage:
  pm build-sidechains in.pdb 
    This reads a structure stored in in.pdb, strips all sidechains, 
    detects and models disulfid bonds and reconstructs all sidechains with the 
    flexible rotamer model. The result is stored as out.pdb
'''

import os
import ost
from ost import io, settings
from promod3 import modelling, sidechain
from promod3.core import pm3argparse, helper

# make sure we see output when passing '-h'
ost.PushVerbosityLevel(2)

# parse command line
parser = pm3argparse.PM3ArgumentParser(__doc__, action=True)

parser.AddStructure()
parser.AssembleParser()
parser.add_argument('-o', '--out-file', metavar='<FILENAME>', type=str,
                    default='out.pdb', help='File to store coordinates ' +
                                            '(default: %(default)s)')
parser.add_argument('-k', '--keep-sidechains', action='store_true',
                    help='Keep existing sidechains')
parser.add_argument('-n', '--no-disulfids', action='store_true',
                    help='Do not build disulfid bonds before sidechain ' + 
                         'optimization')
parser.add_argument('-r', '--rigid-rotamers', action='store_true',
                    help='Do not use rotamers with subrotamers')
parser.add_argument('-i', '--backbone-independent', action='store_true',
                    help='Use backbone independent rotamer library instead ' +
                         'of the default backbone dependent one')
parser.add_argument('-s', '--no-subrotamer-optimization', action='store_true',
                    help='Dont do subrotamer optimization if flexible ' +
                         'rotamer model is used.')
parser.add_argument('-f', '--energy_function', action = 'store', 
                    dest = "energy_function", default = "SCWRL4")

opts = parser.Parse()

ost.PushVerbosityLevel(3)

if len(opts.structures) != 1:
    helper.MsgErrorAndExit("Expect exactly one input structure", 4)

prot = opts.structures[0]
rotamer_model = "frm"
if opts.rigid_rotamers:
    rotamer_model = "rrm"

lib = None
if opts.backbone_independent:
    lib = sidechain.LoadLib()

opt_subrot = True
if opts.no_subrotamer_optimization:
    opt_subrot = False

# do sidechain modelling
modelling.ReconstructSidechains(prot, keep_sidechains=opts.keep_sidechains,
                                build_disulfids=opts.no_disulfids==False, 
                                rotamer_model=rotamer_model, 
                                rotamer_library=lib, 
                                optimize_subrotamers=opt_subrot,
                                energy_function = opts.energy_function)

# output
ost.PopVerbosityLevel()
io.SavePDB(prot, opts.out_file)
if not os.path.isfile(opts.out_file):
    helper.MsgErrorAndExit("Failed to write model file '%s'." % opts.model_file, 4)

