#!/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,importError
from pmx import jobscripts
import sqlite3, shutil,tempfile,re,datetime,collections
from pprint import pprint
import pandas as pd
from pmx.rosetta import muts_read_and_format

try:
    import sqlite3
except:
    raise(importError('sqlite3'))
    

class FlexDDG_analyze:
    '''Analysis of the results from the Flex_ddg protocol 
       for protein-protein binding affinity changes
       upon amino mutation [Barlow et al, JPCB, 2018].


    Parameters
    ----------
    mutFile : str, optional
        file name with mutations (chain resID mut)
    mutList : list, optional
        list of mutations
    path : str, optional
        path where the mutations are located.
    outfile : str, optional
        output file with all results summarized.
    overwrite: bool, optional
        if already present, should the results file
        be overwritten. Default: False
    nstruct : int,
        number of runs per mutation.
    
    '''
#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, mutList=[], path='.', outfile='', overwrite=False, nstruct=35 ):#, backrub_trajectory_stride=35000 ):

        # get parameters
        self.basepath = path
        self.outfile = outfile
        self.bOverwrite = overwrite
        self.nstruct = nstruct
        self.trajectory_stride = 35#backrub_trajectory_stride
        # other default rosetta parameters
        self.rosetta_output_file_name = 'rosetta.out'
        self.output_database_name = 'ddG.db3'
        self.script_output_folder = 'analysis_output'

        # get folders
        self.folders = []
        if (mutFile is None) and len(mutList)==0:
            self._mutations_from_folders( )
        elif mutFile is not None:
            self.sele = muts_read_and_format( mutFile, 'sss' )
            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])
                self.folders.append( folder )
        self.folderNum = len(self.folders)

        # output file
        self.outfile = outfile
        self.alreadyProcessedFolders = []
        self._process_outfile()

        # gam weighting params
        self.zemu_gam_params = {
            'fa_sol' :      (6.940, -6.722),
            'hbond_sc' :    (1.902, -1.999),
            'hbond_bb_sc' : (0.063,  0.452),
            'fa_rep' :      (1.659, -0.836),
            'fa_elec' :     (0.697, -0.122),
            'hbond_lr_bb' : (2.738, -1.179),
            'fa_atr' :      (2.313, -1.649),
        }

        # analyze every mutation
        counter = 1
        for folder in self.folders:
            if folder in self.alreadyProcessedFolders:
                print('Already processed: {0}/{1} {2} '.format(counter,self.folderNum,folder))
                counter+=1
                continue
            print('Analyzing: {0}/{1} {2} '.format(counter,self.folderNum,folder))
            # get the file paths
            finished_jobs = self._find_finished_jobs( self.basepath+'/'+folder )
            if len(finished_jobs.keys())==0:
                print('****Failed (no finished jobs found): {0}/{1} {2}****'.format(counter,self.folderNum,folder))
                counter+=1
                continue

            # get the values
            ddg_scores_dfs = []
            struct_scores_dfs = []
            for key in finished_jobs.keys():
                finished_job = finished_jobs[key]
                inner_scores_list = []
                inner_scores = self._process_finished_struct( finished_job, os.path.basename(finished_job) )
                inner_scores_list.append( inner_scores )
                scores = pd.concat( inner_scores_list )
                ddg_scores, struct_scores = self._calc_ddg( scores )
                struct_scores_dfs.append( struct_scores )
                ddg_scores_dfs.append( ddg_scores )
                ddg_scores_dfs.append( self._apply_zemu_gam(ddg_scores) )
                ddg_scores_dfs.extend( self._calc_dgs( scores ) )

#            pd.concat( struct_scores_dfs ).to_csv( 'tmp.csv' ) 
#            df = pd.concat( ddg_scores_dfs, sort=True )
#            df.to_csv( os.path.join('tmp2.csv') )

            # extract ddG and calculate mean, sderr
            df = pd.concat( ddg_scores_dfs,sort=True )
            ddgvalsGAM = df.loc[ (df['scored_state']=='ddG') & (df['score_function_name'].str.contains('-gam')) ]['total_score']
            ddgvals = df.loc[ (df['scored_state']=='ddG') & (df['score_function_name'].str.contains('-gam')==False) ]['total_score']
            # w/o -gam
            meanval = np.mean(ddgvals)
            sderrval = np.std(ddgvals,ddof=1)/np.sqrt(float(np.shape(ddgvals)[0]))
            # w/ -gam
            meanvalGAM = np.mean(ddgvalsGAM)
            sderrvalGAM = np.std(ddgvalsGAM,ddof=1)/np.sqrt(float(np.shape(ddgvalsGAM)[0]))
            # output
            outfp = open(self.outfile,'a')
            outfp.write( '%10s | %3.2f %3.2f | %3.2f %3.2f\n' % (folder,np.round(meanval,2),np.round(sderrval,2),np.round(meanvalGAM,2),np.round(sderrvalGAM,2)) )
            outfp.close()

            counter+=1


    def _process_outfile( self ):
        if self.outfile=='':
            self.outfile = self.basepath+'/results.dat'
        if self.bOverwrite==True:
            fp = open(self.outfile,'w')
            fp.close()
        # if file already exists, read its contents and append to them
        if os.path.isfile(self.outfile):
            fp = open(self.outfile,'r')
            for l in fp.readlines():
                l = l.rstrip()
                foo = l.split()
                if foo[0].startswith('mut_'):
                    self.alreadyProcessedFolders.append(foo[0])
            fp.close()

    def _gam_function(self, x, score_term = None ):
        return -1.0 * np.exp( self.zemu_gam_params[score_term][0] ) + 2.0 * np.exp( self.zemu_gam_params[score_term][0] ) / ( 1.0 + np.exp( -1.0 * x * np.exp( self.zemu_gam_params[score_term][1] ) ) )

    def _apply_zemu_gam(self, scores):
        new_columns = list(scores.columns)
        new_columns.remove('total_score')
        scores = scores.copy()[ new_columns ]
        for score_term in self.zemu_gam_params:
            assert( score_term in scores.columns )
            scores[score_term] = scores[score_term].apply( self._gam_function, score_term = score_term )
        scores[ 'total_score' ] = scores[ list(self.zemu_gam_params.keys()) ].sum( axis = 1 )
        scores[ 'score_function_name' ] = scores[ 'score_function_name' ] + '-gam'
#        self.gam_function_name = scores[ 'score_function_name' ] + '-gam'
#        self.score_function_name = scores['score_function_name']
        return(scores)

    def _mutations_from_folders( self ):
        folderContents = os.listdir(self.basepath)
        
        for folder in folderContents:
            if folder.startswith('mut_'):
                self.folders.append(folder)

    def _rosetta_output_succeeded( self,potential_struct_dir ):
        path_to_rosetta_output = os.path.join( potential_struct_dir, self.rosetta_output_file_name )
        if not os.path.isfile(path_to_rosetta_output):
            return False

        db3_file = os.path.join( potential_struct_dir, self.output_database_name )
        if not os.path.isfile( db3_file ):
            return False

        success_line_found = False
        no_more_batches_line_found = False
        with open( path_to_rosetta_output, 'r' ) as f:
            for line in f:
                if line.startswith( 'protocols.jd2.JobDistributor' ) \
                   and ( ('reported success in' in line) \
                   or ('did you forget to pass -overwrite' in line) ):
                    success_line_found = True
                if line.startswith( 'protocols.jd2.JobDistributor' ) and 'no more batches to process' in line:
                    no_more_batches_line_found = True

        return(no_more_batches_line_found and success_line_found)

    def _find_finished_jobs( self,folder ):
        return_dict = {}
        for i in range(1, self.nstruct+1):
            job_dir = folder+'/'+str(i)
            if i<10:
                job_dir = folder+'/0'+str(i)
            job_dir = os.path.abspath(job_dir)
            if os.path.isdir(job_dir):
                if self._rosetta_output_succeeded( job_dir ):
                    return_dict[i] = job_dir
        return(return_dict)

    def _process_finished_struct( self, output_path, case_name ):
        db3_file = os.path.join( output_path, self.output_database_name )
        assert( os.path.isfile( db3_file ) )
        struct_number = int( os.path.basename(output_path) )
        scores_df = self._get_scores_from_db3_file( db3_file, struct_number, case_name )
        return(scores_df)

    def _get_scores_from_db3_file(self, db3_file, struct_number, case_name):
        conn = sqlite3.connect(db3_file)
        conn.row_factory = sqlite3.Row
        c = conn.cursor()

        num_batches = c.execute('SELECT max(batch_id) from batches').fetchone()[0]
        scores = pd.read_sql_query('''
        SELECT batches.name, structure_scores.struct_id, score_types.score_type_name, structure_scores.score_value, score_function_method_options.score_function_name from structure_scores
        INNER JOIN batches ON batches.batch_id=structure_scores.batch_id
        INNER JOIN score_function_method_options ON score_function_method_options.batch_id=batches.batch_id
        INNER JOIN score_types ON score_types.batch_id=structure_scores.batch_id AND score_types.score_type_id=structure_scores.score_type_id
        ''', conn)

        def renumber_struct_id( struct_id ):
            return self.trajectory_stride * ( 1 + (int(struct_id-1) // num_batches) )

        scores['struct_id'] = scores['struct_id'].apply( renumber_struct_id )
        scores['name'] = scores['name'].apply( lambda x: x[:-9] if x.endswith('_dbreport') else x )
        scores = scores.pivot_table( index = ['name', 'struct_id', 'score_function_name'], columns = 'score_type_name', values = 'score_value' ).reset_index()
        scores.rename( columns = {
            'name' : 'state',
            'struct_id' : 'backrub_steps',
        }, inplace=True)
        scores['struct_num'] = struct_number
        scores['case_name'] = case_name
        conn.close()
        return(scores)

    def _calc_ddg( self, scores ):
        total_structs = np.max( scores['struct_num'] )

        nstructs_to_analyze = set([total_structs])
        for x in range(10, total_structs):
            if x % 10 == 0:
                nstructs_to_analyze.add(x)
        nstructs_to_analyze = sorted(nstructs_to_analyze)

        all_ddg_scores = []
        for nstructs in nstructs_to_analyze:
            ddg_scores = scores.loc[ ((scores['state'] == 'unbound_mut') | (scores['state'] == 'bound_wt')) & (scores['struct_num'] <= nstructs) ].copy()
            for column in ddg_scores.columns:
                if column not in ['state', 'case_name', 'backrub_steps', 'struct_num', 'score_function_name']:
                    ddg_scores.loc[:,column] *= -1.0
            ddg_scores = ddg_scores.append( scores.loc[ ((scores['state'] == 'unbound_wt') | (scores['state'] == 'bound_mut')) & (scores['struct_num'] <= nstructs) ].copy() )
            ddg_scores = ddg_scores.groupby( ['case_name', 'backrub_steps', 'struct_num', 'score_function_name'] ).sum().reset_index()

            if nstructs == total_structs:
                struct_scores = ddg_scores.copy()

            ddg_scores = ddg_scores.groupby( ['case_name', 'backrub_steps', 'score_function_name'] ).mean().round(decimals=5).reset_index()
            new_columns = list(ddg_scores.columns.values)
            new_columns.remove( 'struct_num' )
            ddg_scores = ddg_scores[new_columns]
            ddg_scores[ 'scored_state' ] = 'ddG'
            ddg_scores[ 'nstruct' ] = nstructs
            all_ddg_scores.append(ddg_scores)
        return(pd.concat(all_ddg_scores), struct_scores)

    def _calc_dgs( self,scores ):
        l = []

        total_structs = np.max( scores['struct_num'] )

        nstructs_to_analyze = set([total_structs])
        for x in range(10, total_structs):
            if x % 10 == 0:
                nstructs_to_analyze.add(x)
        nstructs_to_analyze = sorted(nstructs_to_analyze)

        for state in ['mut', 'wt']:
            for nstructs in nstructs_to_analyze:
                dg_scores = scores.loc[ (scores['state'].str.endswith(state)) & (scores['state'].str.startswith('unbound')) & (scores['struct_num'] <= nstructs) ].copy()
                for column in dg_scores.columns:
                    if column not in ['state', 'case_name', 'backrub_steps', 'struct_num', 'score_function_name']:
                        dg_scores.loc[:,column] *= -1.0
                dg_scores = dg_scores.append( scores.loc[ (scores['state'].str.endswith(state)) & (scores['state'].str.startswith('bound')) & (scores['struct_num'] <= nstructs) ].copy() )
                dg_scores = dg_scores.groupby( ['case_name', 'backrub_steps', 'struct_num', 'score_function_name'] ).sum().reset_index()
                dg_scores = dg_scores.groupby( ['case_name', 'backrub_steps', 'score_function_name'] ).mean().round(decimals=5).reset_index()
                new_columns = list(dg_scores.columns.values)
                new_columns.remove( 'struct_num' )
                dg_scores = dg_scores[new_columns]
                dg_scores[ 'scored_state' ] = state + '_dG'
                dg_scores[ 'nstruct' ] = nstructs
                l.append( dg_scores )
        return(l)
