#!/usr/bin/env python

"""This module contains classes for calling rosetta scripts:
Flex_DDG, ...
"""

import os
import sys
from copy import deepcopy
import numpy as np
from . import library
from .atom import Atom
from .model import Model
from .resselection import InteractiveSelection
from pmx import resselection
from pmx.parser import parseList, kickOutComments
from .utils import create_folder,remove_netmount
from pmx import jobscripts

# =======
# Helper functions
# =======
def get_rosetta_path():
    rosettapath = ""

    # search in several environment variables
    if "ROSETTAPATH" in os.environ:
         rosettapath = os.environ.get('ROSETTAPATH')
         if os.path.isdir(rosettapath):
             return(os.path.abspath(rosettapath))
    if len(rosettapath)==0 and "ROSETTA" in os.environ:
         rosettapath = os.environ.get('ROSETTA')
         if os.path.isdir(rosettapath):
             return(os.path.abspath(rosettapath))
    if len(rosettapath)==0 and "PMXROSETTA" in os.environ:
         rosettapath = os.environ.get('PMXROSETTA')
         if os.path.isdir(rosettapath):
             return(os.path.abspath(rosettapath))
   
    print('Folder for Rosetta binaries not found.')
    print('Please, set one of the variables to rosetta/bin:')
    print('ROSETTAPATH, ROSETTA or PMXROSETTA')
    sys.exit(0) 

def muts_read_and_format(filename, format_string, comment='#',
                    ignore_missing=False):
    """
    formats the mutations
    entries can be separated by a comma, e.g. X 17 A, X 18 A
    format_string : str
        e.g. 'is' = [int, str]; 'iis' = [int, int, str] etc.
    """
    l = open(filename).readlines()
    if comment is not None:
        l = kickOutComments(l, comment)
    output = []
    for ll in l:
        foo = ll.split(',')
        n = parseList(format_string, foo, ignore_missing)
        output.append(n)
    return output


class Rosetta:
    '''Parent class for the other rosetta classes.
       Defines several parameters that are common for all rosetta
       classes.

    Parameters
    ----------
    protFile : str
        file name of the protein structure.
    scoringFunction : str, optional
        scoring function, e.g. fa_talaris2014
    weights : str, optional
        weights. If not defined, assigns according to scoringFunction
    max_cpus : int, optional
        number of cpu to use
    jobscript : str, optional
        queue'ing system (only sge is supported)
    clusterSimTime : int, optional
        simulation time in hours in jobscript
    clusterModulesToLoad : array of strings
        array of modules to load in jobscript
    clusterFilesToSource : array of strings
        array of files to source in jobscript
    
    '''

    def __init__(self, **kwargs):
        self.basepath = '.'
        self.protFiles = ''
        self.scoringFunction = 'fa_talaris2014'
        self.weights = 'talaris2014'
        self.max_cpus = 1
        self.jobscript = 'sge'
        self.clusterSimTime = 24
        self.clusterModulesToLoad = []
        self.clusterFilesToSource = []

        for key, val in kwargs.items():
            setattr(self, key, val)

        if self.protFile is None:
            raise ValueError('protein file is needed')  



# =======
# Classes for Rosetta protocols
# =======
class FlexDDG(Rosetta):
    '''Flex_ddg protocol for protein-protein binding affinity changes
       upon amino mutation [Barlow et al, JPCB, 2018].

       Brief description:...

    Parameters
    ----------
    protFile : str
        file name of the protein structure.
    mutFile : str, optional
        file name with mutations (chain resID mut)
    mutList : list
        list of mutations
    chainsToMove: str, optional
        which chain(s) to move (comma separated list).
        If None, the chain with the mutation will be moved.
    
    '''
#MIT License

#Copyright (c) 2018 Kyle Barlow

#Permission is hereby granted, free of charge, to any person obtaining a copy
#of this software and associated documentation files (the "Software"), to deal
#in the Software without restriction, including without limitation the rights
#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#copies of the Software, and to permit persons to whom the Software is
#furnished to do so, subject to the following conditions:

#The above copyright notice and this permission notice shall be included in all
#copies or substantial portions of the Software.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#SOFTWARE.

    def __init__(self, mutFile=None, chainsToMove=None, mutList=[],
                 nstruct=35, max_minimization_iter=5000, 
                 abs_score_convergence_thresh=1.0, number_backrub_trials=35000, 
                 backrub_trajectory_stride=35000, **kwargs):
#                 jobscript='sge',clusterSimTime=24,clusterModulesToLoad=[],clusterFilesToSource=[]):
#    def __init__(self, **kwargs):

        Rosetta.__init__(self,**kwargs)

        ### deal with protein ###
#        self.protFile = protFile
        self.m = Model(self.protFile,renumber_residues=False,bPDBTER=True,rename_atoms=False,scale_coords='A')

        ### deal with mutations ###
        if (mutFile is None) and len(mutList)==0:
            self.sele = mutate.InteractiveSelection(m=self.m, ff='', renumbered=False)
        elif mutFile is not None:
            self.sele = muts_read_and_format( mutFile, 'sss' )
        else:
            raise ValueError('mutation file or mutation list is needed')  

        ### chainsToMove
        self.chainsToMove = chainsToMove

        ### get parameters
#        self.scoringFunction = scoringFunction #'fa_talaris2014'
#        self.weights = weights #'talaris2014'
#        self.max_cpus = max_cpus
        self.nstruct = nstruct
        self.max_minimization_iter = max_minimization_iter
        self.abs_score_convergence_thresh = abs_score_convergence_thresh
        self.number_backrub_trials = number_backrub_trials
        self.backrub_trajectory_stride = backrub_trajectory_stride
        ### parameters for jobscript
#        self.jobscript = jobscript
#        self.clusterSimTime = clusterSimTime
#        self.clusterModulesToLoad = clusterModulesToLoad
#        self.clusterFilesToSource = clusterFilesToSource

        ### folder structure ###
#        self.basepath = path
        self.folders = []
        self.__gen_folder_structure( )

        ### rosetta input files ###
        self.__gen_flex_ddg_input_files( )

        ### rosetta flex_ddg xml protocol ###
        self.__write_flex_ddg_xml_protocol( )

        ### rosetta runner ###
        self.__write_flex_ddg_runner( )

        ### jobscript ###
        if self.jobscript!='':
            self.__gen_jobscript( )
            

    def __gen_jobscript( self ):
        jobscriptFiles = []

        for sel,folder in zip(self.sele,self.folders):
            simcpu = self.max_cpus
            jobname = 'mut_ch{0}_{1}{2}'.format(sel[0][0],sel[0][1],sel[0][2])
            for i in range(1,len(sel)): # several mutations simultaneously
                s = sel[i]
                jobname = jobname+'_ch{0}_{1}{2}'.format(s[0],s[1],s[2])
            jobscriptFile = folder+'/jobscript'
            jobscriptFiles.append(jobscriptFile)
            cmdline = 'python run.py'
            if 'sge' in self.jobscript.lower():
                jobscripts.SGE_jobscript( fname=jobscriptFile, jobname=jobname,
                                  simcpu=self.max_cpus, simtime=self.clusterSimTime,
                                  modules=self.clusterModulesToLoad,source=self.clusterFilesToSource,
                                  cmd=cmdline, path=folder )
        submitterFile = self.basepath+'/submit_jobscripts.py'
        jobscripts.submitter( jobscriptFiles, self.folders, self.basepath, submitterFile, self.jobscript )
            

    def __gen_flex_ddg_input_files( self ):
        for sel,folder in zip(self.sele,self.folders):
            ### nataa file
            fname = folder+'/nataa_mutations.resfile'
            fp = open(fname,'w')
            fp.write('NATAA\nstart\n')
            fp.write('{0} {1} PIKAA {2}\n'.format(sel[0][1],sel[0][0],sel[0][2]))
            for i in range(1,len(sel)): # several mutations simultaneously
                s = sel[i]
                fp.write('{0} {1} PIKAA {2}\n'.format(s[1],s[0],s[2]))
            fp.close()

            ### chains-to-move file
            fname = folder+'/chains_to_move.txt'
            chToMove = sel[0][0]
            chainsAlreadyFound = [chToMove]
            for i in range(1,len(sel)): # several mutations simultaneously
                if sel[i][0] not in chainsAlreadyFound:
                    chToMove = chToMove + ','+sel[i][0]
            if self.chainsToMove!=None:
                chToMove = self.chainsToMove
            fp = open(fname,'w')
            fp.write('{0}\n'.format(chToMove))
            fp.close()

    def __gen_folder_structure(self):
        for sel in self.sele:
            folder = '/mut_ch{0}_{1}{2}'.format(sel[0][0],sel[0][1],sel[0][2])
            for i in range(1,len(sel)): # several mutations simultaneously
                s = sel[i]
                folder = folder+'_ch{0}_{1}{2}'.format(s[0],s[1],s[2])
            dirname = os.path.abspath( self.basepath+folder )
            self.folders.append( dirname )
            create_folder(dirname)

    def __write_flex_ddg_runner( self ):
        rosettapath = get_rosetta_path()

        talaris = ''
        if 'talaris' in self.scoringFunction:
            talaris = '\'-restore_talaris_behavior\','
        elif 'score12' in self.scoringFunction:
            talaris = '\'-restore_pre_talaris_2013_behavior\','

        for sel,folder in zip(self.sele,self.folders):
            xmlprotocol = './ddG-backrub.xml'
            fname = folder+'/run.py'
            fp = open(fname,'w')
            runner = '''#!/usr/bin/python

import socket
import sys
import os
import subprocess

use_multiprocessing = True
if use_multiprocessing:
    import multiprocessing
    max_cpus = {max_cpus} # We might want to not run on the full number of cores, as Rosetta take about 2 Gb of memory per instance

rosetta_scripts_path = os.path.expanduser("{rpath}/rosetta_scripts.linuxgccrelease")
nstruct = {nstruct} # Normally 35
max_minimization_iter = {max_minimization_iter} #5 # Normally 5000
abs_score_convergence_thresh = {abs_score_convergence_thresh} #200.0 # Normally 1.0
number_backrub_trials = {number_backrub_trials} #10 # Normally 35000
backrub_trajectory_stride = {backrub_trajectory_stride} # Can be whatever you want, if you would like to see results from earlier time points in the backrub trajectory. 7000 is a reasonable number, to give you three checkpouints for a 35000 step run, but you could also set it to 35000 for quickest run time (as the final minimization and packing steps will only need to be run one time).
path_to_script = '{xmlfile}'

def run_flex_ddg( input_path, input_pdb_path, chains_to_move, nstruct_i ):
    output_directory = os.path.join( os.path.join( input_path, '%02d' % nstruct_i ) )
    if not os.path.isdir(output_directory):
        os.makedirs(output_directory)

    flex_ddg_args = [
        os.path.abspath(rosetta_scripts_path),
        "-s %s" % os.path.abspath(input_pdb_path),
        '-parser:protocol', os.path.abspath(path_to_script),
        '-parser:script_vars',
        'chainstomove=' + chains_to_move,
        'mutate_resfile_relpath=' + os.path.abspath( os.path.join( input_path, 'nataa_mutations.resfile' ) ),
        'number_backrub_trials=%d' % number_backrub_trials,
        'max_minimization_iter=%d' % max_minimization_iter,
        'abs_score_convergence_thresh=%.1f' % abs_score_convergence_thresh,
        'backrub_trajectory_stride=%d' % backrub_trajectory_stride ,
        {talaris}
        '-in:file:fullatom',
        '-ignore_unrecognized_res',
        '-ignore_zero_occupancy false',
        '-ex1',
        '-ex2',
    ]

    log_path = os.path.join(output_directory, 'rosetta.out')

    print('Running Rosetta with args:')
    print(' '.join(flex_ddg_args))
    print('Output logged to:', os.path.abspath(log_path))
#    print

    outfile = open(log_path, 'w')
    process = subprocess.Popen(flex_ddg_args, stdout=outfile, stderr=subprocess.STDOUT, close_fds = True, cwd = output_directory)
    returncode = process.wait()
    outfile.close()

if __name__ == '__main__':
    cases = []
    case_path = './'
    input_pdb_path = '{protFile}'
    with open( os.path.join( case_path, 'chains_to_move.txt' ), 'r' ) as f:
        chains_to_move = f.readlines()[0].strip()
    for nstruct_i in range(1, nstruct + 1 ):
        cases.append( (case_path, input_pdb_path, chains_to_move, nstruct_i) )

    if use_multiprocessing:
        pool = multiprocessing.Pool( processes = min(max_cpus, multiprocessing.cpu_count()) )

    for args in cases:
        if use_multiprocessing:
            pool.apply_async( run_flex_ddg, args = args )
        else:
            run_flex_ddg( *args )

    if use_multiprocessing:
        pool.close()
        pool.join()
'''.format(rpath=rosettapath,xmlfile=xmlprotocol,
           max_cpus=self.max_cpus,
           nstruct=self.nstruct,
           max_minimization_iter=self.max_minimization_iter,
           abs_score_convergence_thresh=self.abs_score_convergence_thresh,
           number_backrub_trials=self.number_backrub_trials,
           backrub_trajectory_stride=self.backrub_trajectory_stride,talaris=talaris,
           protFile=remove_netmount(self.protFile))
       
            fp.write( runner )
            fp.close()

    def __write_flex_ddg_xml_protocol( self ):
        xmlscript = '''<ROSETTASCRIPTS>
  <SCOREFXNS>
    <ScoreFunction name="fa_{scoreFunc}" weights="{weights}"/>
    <ScoreFunction name="fa_{scoreFunc}_cst" weights="{weights}">
      <Reweight scoretype="atom_pair_constraint" weight="1.0"/>
      <Set fa_max_dis="9.0"/>
    </ScoreFunction>
  </SCOREFXNS>

  <!-- ### Only required input file (other than PDB) - mutation resfile ### -->
  <!-- #### All residues must be set to be NATAA packable at top of resfile ### -->
  <TASKOPERATIONS>
    <ReadResfile name="res_mutate" filename="%%mutate_resfile_relpath%%"/>
  </TASKOPERATIONS>

  <RESIDUE_SELECTORS>
    <Task name="resselector" fixed="0" packable="0" designable="1" task_operations="res_mutate"/>
    <Neighborhood name="bubble" selector="resselector" distance="8.0"/>
    <PrimarySequenceNeighborhood name="bubble_adjacent" selector="bubble" lower="1" upper="1"/>
    <StoredResidueSubset name="restore_neighbor_shell" subset_name="neighbor_shell"/>
    <Not name="everythingelse" selector="restore_neighbor_shell"/>
  </RESIDUE_SELECTORS>
  <TASKOPERATIONS>
    <OperateOnResidueSubset name="repackonly" selector="restore_neighbor_shell">
      <RestrictToRepackingRLT/>
    </OperateOnResidueSubset>
    <OperateOnResidueSubset name="norepack" selector="everythingelse">
      <PreventRepackingRLT/>
    </OperateOnResidueSubset>
    <UseMultiCoolAnnealer name="multicool" states="6"/>
    <ExtraChiCutoff name="extrachizero" extrachi_cutoff="0"/>
    <InitializeFromCommandline name="commandline_init"/>
    <RestrictToRepacking name="restrict_to_repacking"/>
  </TASKOPERATIONS>

  <FILTERS>
  </FILTERS>

  <MOVERS>
    <StoreResidueSubset name="neighbor_shell_storer" subset_name="neighbor_shell" residue_selector="bubble_adjacent" />

    <AddConstraintsToCurrentConformationMover name="addcst" use_distance_cst="1" coord_dev="0.5" min_seq_sep="0" max_distance="9" CA_only="1" bound_width="0.0" cst_weight="0.0"/>
    <ClearConstraintsMover name="clearcst"/>
    <MinMover name="minimize" scorefxn="fa_{scoreFunc}_cst" chi="1" bb="1" type="lbfgs_armijo_nonmonotone" tolerance="0.000001" max_iter="%%max_minimization_iter%%" abs_score_convergence_threshold="%%abs_score_convergence_thresh%%"/>

    <PackRotamersMover name="repack" scorefxn="fa_{scoreFunc}" task_operations="commandline_init,repackonly,norepack,multicool"/>
    <PackRotamersMover name="mutate" scorefxn="fa_{scoreFunc}" task_operations="commandline_init,res_mutate,norepack,multicool"/>

    <ReportToDB name="dbreport" batch_description="interface_ddG" database_name="ddG.db3">
      <ScoreTypeFeatures/>
      <ScoreFunctionFeatures scorefxn="fa_{scoreFunc}"/>
      <StructureScoresFeatures scorefxn="fa_{scoreFunc}"/>
    </ReportToDB>

    <ReportToDB name="structreport" batch_description="interface_ddG_struct" database_name="struct.db3">
      <PoseConformationFeatures/>
      <PdbDataFeatures/>
      <JobDataFeatures/>
      <ResidueFeatures/>
      <PoseCommentsFeatures/>
      <ProteinResidueConformationFeatures/>
      <ResidueConformationFeatures/>
    </ReportToDB>

    <SavePoseMover name="save_wt_bound_pose" restore_pose="0" reference_name="wt_bound_pose"/>
    <SavePoseMover name="save_backrub_pose" restore_pose="0" reference_name="backrubpdb"/>
    <SavePoseMover name="restore_backrub_pose" restore_pose="1" reference_name="backrubpdb"/>

    <InterfaceDdGMover name="int_ddG_mover" wt_ref_savepose_mover="save_wt_bound_pose" chain_name="%%chainstomove%%" db_reporter="dbreport" scorefxn="fa_{scoreFunc}"/>

    <ScoreMover name="apply_score" scorefxn="fa_{scoreFunc}_cst" verbose="0"/>

    <!-- This ParsedProtocol allows the ddG calculation to take place multiple times along the backrub trajectory, if desired -->
    <ParsedProtocol name="finish_ddg_post_backrub">
      <Add mover_name="save_backrub_pose"/>
      <Add mover_name="structreport"/>

      <Add mover_name="repack"/>

      <Add mover_name="addcst"/>
      <Add mover_name="minimize"/>
      <Add mover_name="clearcst"/>

      <Add mover_name="save_wt_bound_pose"/>
      <Add mover_name="structreport"/>
      <Add mover_name="restore_backrub_pose"/>

      <Add mover_name="mutate"/>

      <Add mover_name="addcst"/>
      <Add mover_name="minimize"/>
      <Add mover_name="clearcst"/>
      <Add mover_name="structreport"/>

      <Add mover_name="int_ddG_mover"/>
    </ParsedProtocol>

    <BackrubProtocol name="backrub" mc_kt="1.2" ntrials="%%number_backrub_trials%%" pivot_residue_selector="restore_neighbor_shell" task_operations="restrict_to_repacking,commandline_init,extrachizero" recover_low="0" trajectory_stride="%%backrub_trajectory_stride%%" trajectory_apply_mover="finish_ddg_post_backrub"/>

  </MOVERS>
  <APPLY_TO_POSE>
  </APPLY_TO_POSE>
  <PROTOCOLS>
    <Add mover_name="addcst"/>
    <Add mover_name="apply_score"/> <!-- Necessary to initialize neighbor graph -->
    <Add mover_name="neighbor_shell_storer"/>

    <Add mover_name="minimize"/>
    <Add mover_name="clearcst"/>

    <Add mover_name="backrub"/>
  </PROTOCOLS>
  <OUTPUT />
</ROSETTASCRIPTS>
'''.format(scoreFunc=self.scoringFunction,weights=self.weights)
       
        for sel,folder in zip(self.sele,self.folders):
            fname = folder+'/ddG-backrub.xml'
            fp = open(fname,'w')
            fp.write( xmlscript )
            fp.close()


######################################
######################################
######################################
class DDGmonomer:
    '''DDGmonomer protocol for protein stability changes
       upon amino mutation [Kellogg et al, 2011, Proteins].

       Brief description:...

    Parameters
    ----------
    protFile : str
        file name of the protein structure.
    mutFile : str, optional
        file name with mutations (chain resID mut)
    mutList : list
        list of mutations
    path : str, optional
        path where to build the folder structure.
    
    '''

    def __init__(self, protFile=None, mutFile=None, basepath='.', mutList=[],
                 scoringFunction='fa_talaris2014', weights='talaris2014',         
                 max_cpus=1,jobscript='sge',clusterSimTime=24,clusterModulesToLoad=[],clusterFilesToSource=[]):

        ### deal with protein ###
        if protFile is None:
            raise ValueError('protein file is needed')  
        self.protFile = protFile
        self.m = Model(self.protFile,renumber_residues=False,bPDBTER=True,rename_atoms=False,scale_coords='A')

        ### deal with mutations ###
        if (mutFile is None) and len(mutList)==0:
            self.sele = mutate.InteractiveSelection(m=self.m, ff='', renumbered=False)
        elif mutFile is not None:
            self.sele = muts_read_and_format( mutFile, 'sss' )
        else:
            raise ValueError('mutation file or mutation list is needed')  

        ### get parameters
        self.scoringFunction = scoringFunction #'fa_talaris2014'
        self.weights = weights #'talaris2014'
        self.max_cpus = max_cpus
        self.nstruct = nstruct
        self.max_minimization_iter = max_minimization_iter
        self.abs_score_convergence_thresh = abs_score_convergence_thresh
        self.number_backrub_trials = number_backrub_trials
        self.backrub_trajectory_stride = backrub_trajectory_stride
        ### parameters for jobscript
        self.jobscript = jobscript
        self.clusterSimTime = clusterSimTime
        self.clusterModulesToLoad = clusterModulesToLoad
        self.clusterFilesToSource = clusterFilesToSource
 
        # ... Unfinished ...


