#!/usr/bin/env python
import os
import sys
import shutil
import math
import argparse

from glob import glob
import pandas as pd
import numpy as np

from mdkit.utility import mol2
from dockbox.dbxtools import *

# command-line arguments and options
parser = argparse.ArgumentParser(description="Extract best docking poses after rundbx finished.")

parser.add_argument('-all-targets',
    dest='combine_targets',
    action='store_true',
    default=False,
    help='Select best poses over all the targets. If not specified, extract best pose separately for each target. A "%s/%s/%s" architecture \
of the folders is assumed'%(ligdir_prefix,tardir_prefix,isodir_prefix))

parser.add_argument('-all-isomers',
    dest='combine_isomers',
    action='store_true',
    default=False,
    help='Select best poses over all the isomers. If not specified, extract best pose separately for every isomer. A "%s/%s/%s" architecture \
of the folders is assumed'%(ligdir_prefix,tardir_prefix,isodir_prefix))

parser.add_argument('-copy',
    default=False,
    action='store_true',
    help='Copy best poses in folder.')

parser.add_argument('-csvl',
    type=str,
    dest='csvfile_l',
    metavar='FILE',
    help='Filename containing info about compounds. Used to add names of compounds. Default: none')

parser.add_argument('-csvr',
    dest='csvfile_r',
    default=None,
    metavar='FILENAME',
    help='Filename containing info about targets. If none, will look for a receptor file in the "poses" folders.  Default: none')

parser.add_argument('-cutoff',
    dest='cutoff',
    type=float,
    metavar='RMSD_VALUE',
    default=2.0,
    help='RMSD cutoff used for consensus docking or score-based consensus docking. Default: 2.0 A')

parser.add_argument('-d',
    dest='docking_programs',
    nargs='+',
    metavar=('PRGM1', 'PRGM2'),
    help='Docking programs (instances) to be considered when extracting best poses')

parser.add_argument('-dirs',
    dest='dirs',
    nargs='+',
    default=['.'],
    metavar=('DIR1', 'DIR2'),
    help='Directories considered for analysis. Should contain a folder called "poses". Default: curr. dir')

parser.add_argument('-label',
    dest='label',
    default='results',
    metavar='LABEL NAME',
    help='Name of results label. Default: results')

group = parser.add_mutually_exclusive_group(required=False)

group.add_argument('-sf',
    nargs='+',
    dest='sf',
    metavar='FUNC',
    help='Scoring functions used to extract the best pose (combination of scores)')

group.add_argument('-cd',
    dest='cd',
    nargs='+',
    metavar='PRGM',
    help='Docking programs used for standard consensus docking')

group.add_argument('-sbcd',
    dest='sbcd',
    nargs='+',
    metavar='FUNC',
    help='Scoring functions used for score-based consensus docking')

# update parsers with arguments
args = parser.parse_args()

def add_names(csvfile, df):
    df_ligands = pd.read_csv(csvfile)
    if 'isomer' in df_ligands:
        df_ligands = df_ligands[df_ligands['isomer']==1]

    df = df.merge(df_ligands[['ligID', 'name']], on='ligID')
    return df

dirs = []
for dir in args.dirs:
    if os.path.isdir(dir+'/poses'):
        dirs.append(os.path.relpath(dir))
    else:
        raise ValueError('directory '+dir+'/poses does not exist!')
iscwd, isligID, istargetID, isisomerID = check_directories(dirs)

if not istargetID and args.combine_targets:
    args.combine_targets = False

if not isisomerID and args.combine_isomers:
    args.combine_isomers = False

# check if info related to targets is there!
is_rec_pdb = True
if args.csvfile_r and istargetID: 
    df_targets = pd.read_csv(args.csvfile_r)
    csvfile_r_dir = os.path.dirname(args.csvfile_r)
    # update relative paths
    if not csvfile_r_dir:
        csvfile_r_dir = '.'
    df_targets['pdbfile'] =  df_targets['pdbfile'].apply(lambda x: os.path.relpath(csvfile_r_dir+'/'+x))
else:
    for dir in args.dirs:
        # check if .pdb files are there
        if not os.path.isfile(dir+'/poses/rec.pdb'):
            is_rec_pdb = False
            break

# check options relative to best poses extraction
scoring_functions_all = []
if args.sbcd:
    scoring_functions = args.sbcd
    programs_consensus = args.sbcd
    if len(args.sbcd) < 2:
        raise ValueError('Number of functions for score-based consensus docking should be at least 2!')
    resolve_with = args.sbcd[0] # used to decide which pose to extract when selecting over all isomers
elif args.cd:
    scoring_functions = None
    programs_consensus = args.cd
    if len(args.cd) < 2:
        raise ValueError('Number of programs for consensus docking should be at least 2!')
    resolve_with = 'score_'+args.cd[0]
elif args.sf:
    scoring_functions = args.sf
    programs_consensus = None
    if len(scoring_functions) == 1:
        resolve_with = scoring_functions[0]
    else:
        resolve_with = '_'.join(scoring_functions)

if args.csvfile_l:
    if not os.path.isfile(args.csvfile_l):
        raise IOError("csvfile %s not found!"%args.csvfile_l)

features_ids = []
if isligID:
    features_ids += ['ligID']
elif not istargetID and not isisomerID:
    features_ids += ['dir']
if istargetID:
    features_ids += ['targetID']
if isisomerID:
    features_ids += ['isomerID']

files_r = {}
poses = []
for jdx, dir in enumerate(dirs):
    posedir = dir + '/poses'
    ligID, targetID, isomerID = get_IDs(dir, isligID, istargetID, isisomerID)

    info_dir = {}
    for ft in features_ids:
        info_dir[ft] = []
    info_dir['file_l'] = []
    if (args.csvfile_r and istargetID) or is_rec_pdb:
        info_dir['file_r'] = []
    for ft in ['site', 'program', 'instance', 'index_pose', 'score']:
        info_dir[ft] = []

    # get location of poses and receptor files
    with open(posedir+'/info.dat', 'r') as inff:
        # skip the first two lines
        inff.next()
        inff.next()
        for line in inff:
            program, nposes, firstidx, site = line.strip().split(',')
            firstidx = int(firstidx)
            nposes = int(nposes)
            instance = program
            if site:
                instance += '.' + site
            poses_idxs = range(firstidx, firstidx+nposes)

            for index, idx in enumerate(poses_idxs):
                file_l = posedir + '/lig-%s.mol2'%idx
                if os.path.isfile(file_l):
                    info_dir['file_l'].append(os.path.relpath(file_l))
                else:
                    raise IOError("File %s does not exist!"%file_l)
                info_dir['site'].append(site)
                info_dir['program'].append(program)
                info_dir['instance'].append(instance)
                info_dir['index_pose'].append(index)

                if isligID:    
                    info_dir['ligID'].append(ligID)
                elif not istargetID and not isisomerID:
                    info_dir['dir'].append(dir)
                if istargetID:
                    info_dir['targetID'].append(targetID)
                if isisomerID:
                    info_dir['isomerID'].append(isomerID)
                # get the filename of the corresponding receptor file
                if istargetID and args.csvfile_r:
                    row = df_targets[df_targets['targetID']==targetID]
                    file_r = row['pdbfile'].values[0]
                    info_dir['file_r'].append(file_r)
                elif is_rec_pdb:
                    file_r = os.path.relpath(posedir+'/rec.pdb')
                    info_dir['file_r'].append(file_r)

                # update the dictionnary of targets
                if istargetID and targetID not in files_r:
                    files_r[targetID] = file_r
            nscores = 0
            # extract original scores
            with open(dir+'/'+instance+'/score.out', 'r') as sout:
                for line_s in sout:
                    nscores += 1
                    info_dir['score'].append(float(line_s.strip()))
                if nscores != nposes:
                    raise ValueError("Number of poses different from number of scores (%s/%s)"%(dir,instance))

    #print dir
    #print firstidx+nposes-1
    # extract all scores
    for score_file in sorted(glob(dir+'/rescoring/*.score')):
        sf = os.path.basename(score_file).split('.')[0]
        if jdx == 0:
            scoring_functions_all.append(sf)
        elif sf not in scoring_functions_all:
            raise ValueError("%s scores not computed in every directory!")
        info_dir[sf] = []
        with open(score_file, 'r') as sout:
            for line_s in sout:
                info_dir[sf].append(float(line_s))
        #print sf, len(info_dir[sf])

    df_dir = pd.DataFrame(info_dir)
    if args.docking_programs: 
        df_dir = df_dir[df_dir['program'].isin(args.docking_programs)]
    poses.append(df_dir)

if poses:
    poses = pd.concat(poses).reset_index()
    if args.csvfile_l and isligID:
        poses = add_names(args.csvfile_l, poses)
else:
    sys.exit("No poses to extract!")

groupby_columns = []
if isligID:
    groupby_columns += ['ligID']
if not args.combine_targets and istargetID:
    groupby_columns += ['targetID']
if isisomerID:
    groupby_columns += ['isomerID']

if not isligID and not isisomerID:
    if not istargetID:
        groupby_columns += ['dir']
    elif args.combine_targets:
        poses['dummy'] = 0
        groupby_columns += ['dummy']

if args.sbcd or args.cd:
    best_poses = []
    for prgm in programs_consensus:
        if args.sbcd:
            if groupby_columns:
                poses_groupby = poses.groupby(groupby_columns)
            minidxs = poses_groupby[prgm].idxmin()
            # handle cases where all the scores provided by a program are nans (to be changed when istargetID or isisomerID is True)
            lignans = minidxs[minidxs.apply(np.isnan)]
            for ligID, row in lignans.iteritems():
                row_poses = poses[poses['ligID']==ligID].iloc[0]
                minidxs[ligID] = float(row_poses.name)
            minidxs = minidxs.astype(int)
            # get best poses from indices
            best_poses_prgm = poses.loc[minidxs].copy()
            # set file_l to nan for those with no score
            for ligID, row in lignans.iteritems():
                best_poses_prgm_row = best_poses_prgm[best_poses_prgm['ligID']==ligID]
                indices = best_poses_prgm_row.index.values
                best_poses_prgm.at[indices, 'file_l'] = np.nan
        elif args.cd:
            poses_prgm = poses[poses['program']==prgm]
            poses_groupby = poses_prgm.groupby(groupby_columns)
            best_poses_prgm = poses_prgm.loc[poses_groupby['score'].idxmin()]
        best_poses_prgm = best_poses_prgm.drop('index', axis=1)

        new_columns_names = [] # renaming columns according to the scoring function
        for col in best_poses_prgm.columns.values:
            if col in groupby_columns+['name']:
                new_columns_names.append(col)
            elif not isligID and col == 'dir':
                new_columns_names.append(col)
            elif args.sbcd and col == prgm:
                new_columns_names.append(prgm)
            else:
                new_columns_names.append(col + '_' + prgm)
        best_poses_prgm.columns = new_columns_names
        best_poses.append(best_poses_prgm)

    columns_to_be_merged = groupby_columns
    if 'name' in poses.columns.values:
       columns_to_be_merged += ['name']

    # merge best poses into single dataframe
    best_poses_merged = best_poses[0]    
    for item in best_poses[1:]:
        best_poses_merged = best_poses_merged.merge(item, on=columns_to_be_merged, how='outer')

    prgm_first = programs_consensus[0]
    if args.combine_targets:
        rmsd_rot_trans = get_rmsd_rotation_and_translations_all_targets(files_r)
        for prgm in programs_consensus[1:]:
            best_poses_merged['rmsd_'+prgm_first+'_'+prgm] = best_poses_merged.apply(lambda row: compute_rmsd(row['file_l_'+prgm_first], row['file_l_'+prgm],
            rotmat=rmsd_rot_trans[row['targetID_'+prgm_first]][row['targetID_'+prgm]][0], \
            trans1=rmsd_rot_trans[row['targetID_'+prgm_first]][row['targetID_'+prgm]][1], \
            trans2=rmsd_rot_trans[row['targetID_'+prgm_first]][row['targetID_'+prgm]][2]), axis=1)
    else:
        for prgm in programs_consensus[1:]:
            best_poses_merged['rmsd_'+prgm_first+'_'+prgm] = best_poses_merged.apply(lambda row: compute_rmsd(row['file_l_'+prgm_first], row['file_l_'+prgm]), axis=1)

    rmsd_columns = [col for col in best_poses_merged.columns.values if col.startswith('rmsd')]
    best_poses_merged = best_poses_merged.assign(consensus=(best_poses_merged[rmsd_columns]<=args.cutoff).all(axis=1))
    best_poses = best_poses_merged
elif args.sf:
    poses_groupby = poses.groupby(groupby_columns)
    if len(args.sf) > 1:
        mean = poses[scoring_functions].mean()
        std = poses[scoring_functions].std()
        column_name = '_'.join(scoring_functions)
        poses[column_name] = ((poses[scoring_functions]-mean)/std).sum(axis=1)/len(args.sf)
        best_poses = poses.loc[poses_groupby[column_name].idxmin]
    else:
        best_poses = poses.loc[poses_groupby[args.sf[0]].idxmin]

# determine best poses over isomers according to strategy provided
if args.combine_isomers:
    groupby_columns = []
    if isligID:
        groupby_columns += ['ligID']
    if istargetID and not args.combine_targets:
        groupby_columns += ['targetID']
    if groupby_columns:
        poses_groupby = best_poses.groupby(groupby_columns)
        best_poses = best_poses.loc[poses_groupby[resolve_with].idxmin]
    else:
        best_poses = best_poses.loc[poses[resolve_with].idxmin].T

features_ids_sorted = list(features_ids)
if istargetID and 'targetID' not in best_poses.columns.values:
    features_ids_sorted.remove('targetID')
if features_ids_sorted:
    best_poses = best_poses.sort_values(features_ids_sorted)

shutil.rmtree(args.label, ignore_errors=True)
os.mkdir(args.label)

features_csv = []
if args.csvfile_l:
    features_csv.append('name')
features_csv += features_ids
features_csv += ['instance']
#features_csv += ['file_l', 'file_r', 'instance']

is_site = False
if list(set(poses['site'])) != ['']:
    is_site = True
    features_csv.append('site')
features_csv += scoring_functions_all + ['score']

# TODO: add an option to save all the poses to .csv file
poses[features_csv].to_csv(args.label+'/poses.csv', index=False, float_format='%.3f')

if best_poses is not None:
    if args.sbcd or args.cd:
        features_csv_best_poses = []
        if args.csvfile_l:
            features_csv_best_poses.append('name')
        if isligID:
            features_csv_best_poses.append('ligID')
        if istargetID and not args.combine_targets:
            features_csv_best_poses.append('targetID')
        if isisomerID:
            features_csv_best_poses.append('isomerID')
        for prgm in programs_consensus:
            if istargetID and args.combine_targets:
                features_csv_best_poses.append('targetID_'+prgm)
            features_csv_best_poses.append('instance_'+prgm)
            #features_csv_best_poses.extend([item + prgm for item in ['file_l_', 'file_r_', 'instance_']])
            if args.sbcd:
                features_csv_best_poses.append(prgm)
            else:
                features_csv_best_poses.append('score_'+prgm)
            if is_site:
                features_csv_best_poses.append('site_'+prgm)
        features_csv_best_poses.extend([col for col in best_poses.columns.values if col.startswith('rmsd')])
        features_csv_best_poses.append('consensus')
    elif args.sf:
        features_csv_best_poses = list(features_csv)
        if len(args.sf) > 1:
            features_csv_best_poses.append('_'.join(scoring_functions))
    filename = args.label + '/best_poses.csv'
    best_poses[features_csv_best_poses].to_csv(filename, index=False, float_format='%.3f')

    if args.combine_targets and istargetID:
        features_ids.remove('targetID')
    
    if args.combine_isomers and isisomerID:
        features_ids.remove('isomerID')

    def copy_best_poses(dir, row, suffix=''):

        file_l = row['file_l'+suffix]
        shutil.copyfile(file_l, dir+'/ligand%s.mol2'%suffix)

        if 'file_r'+suffix in row:
            file_r = row['file_r'+suffix]
            # TODO: add an option for saving the target file
            shutil.copyfile(file_r, dir+'/protein%s.pdb'%suffix)

        # saving ligand files prior to minimization
        instance = row['instance'+suffix]
        index = row['index_pose'+suffix]
        origindir = '/'.join(file_l.split('/')[:-2])
        if os.path.isdir(origindir+'/'+instance+'/origin'):
            poses_idxs = []
            for filename in glob(origindir+'/'+instance+'/lig-*.mol2'):
                poses_idxs.append(int((filename.split('.')[-2]).split('-')[-1]))
                poses_idxs = sorted(poses_idxs)
            pose_idx = poses_idxs[int(index)]
            shutil.copyfile(origindir+'/'+instance+'/origin/lig-%i.mol2'%pose_idx, dir+'/ligand_orig%s.mol2'%suffix)

    if args.copy:
        for idx, row in best_poses.iterrows():
            if args.sf or ((args.sbcd or args.cd) and row['consensus']):
                if isligID:
                    newdir = args.label + '/' + '/'.join(row[features_ids])
                elif row['dir'] != '.':
                    newdir = args.label + '/' + '_'.join(row['dir'].split('/'))
                elif row['dir'] == '.' and len(dirs) > 1:
                    newdir = args.label + '/pwd'
                else:
                    newdir = args.label
                if not os.path.isdir(newdir):
                    os.makedirs(newdir)

                if args.sbcd or args.cd:
                    for idx, prgm in enumerate(programs_consensus):
                        copy_best_poses(newdir, row, suffix='_'+prgm)
                else:
                    copy_best_poses(newdir, row)
