#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Created on Tue Aug  6 12:56:03 2019

@author: tanu
'''
#=======================================================================
# Task: combining all dfs to a single one

# Input: 12/13/14 dfs
#1) <gene>.lower()'_complex_mcsm_norm.csv'
#2) <gene>.lower()_foldx.csv'
#3) <gene>.lower()_dssp.csv'
#4) <gene>.lower()_kd.csv'
#5) <gene>.lower()_rd.csv'
#6) 'ns' + <gene>.lower()_snp_info.csv'
#7) <gene>.lower()_af_or.csv'
#8) <gene>.lower() _af_or_kinship.csv (ONLY for pncA, but omitted for the final run)
#9) <gene>.lower()'_dynamut2.csv'
#10) <gene>.lower()'_dynamut.csv'
#11) <gene>.lower()'_mcsm_na.csv'
#12) <gene>.lower()'_mcsm_ppi2.csv'
#13) <gene>.lower()'_consurf.csv'
#14) <gene>.lower()'_snap2.csv'
#15) <gene>.lower()'_provean.csv
#16) <gene>.lower()'_mmcsm_lig_results.csv'
#17) <gene>.lower()'_edXXX'!!!! TODO

# combining order


# Output: single csv of all 8 dfs combined
# useful link
# https://stackoverflow.com/questions/23668427/pandas-three-way-joining-multiple-dataframes-on-columns

#%% FIXME: let the script proceed even if files don't exist!
# i.e example below
# '/home/tanu/git/Data/ethambutol/output/dynamut_results/embb_complex_dynamut_norm.csv'

#=======================================================================
#%% load packages
import sys, os
import pandas as pd
from pandas import DataFrame
import numpy as np
import argparse
from functools import reduce
from sklearn.preprocessing import MinMaxScaler
#=======================================================================
#%% specify input and curr dir
homedir = os.path.expanduser('~')

# set working dir
os.getcwd()
#os.chdir(homedir + '/git/LSHTM_analysis/scripts')
sys.path.append(homedir + '/git/LSHTM_analysis/scripts')
os.getcwd()

#from combining import combine_dfs_with_checks
from combining_FIXME import detect_common_cols
from reference_dict import oneletter_aa_dict
from reference_dict import low_3letter_dict 

from aa_code import get_aa_3lower
from aa_code import get_aa_1upper

# REGEX: as required
# mcsm_regex = r'^([A-Za-z]{1})([0-9]+)([A-Za-z]{1})$'
# mcsm_wt = mcsm_df['mutationinformation'].str.extract(mcsm_regex)[0]
# mcsm_mut = mcsm_df['mutationinformation'].str.extract(mcsm_regex)[2]

# gwas_regex = r'^([A-Za-z]{3})([0-9]+)([A-Za-z]{3})$'
# gwas_wt = mcsm_df['mutation'].str.extract(gwas_regex)[0]
# gwas_pos = mcsm_df['mutation'].str.extract(gwas_regex)[1]
# gwas_mut = mcsm_df['mutation'].str.extract(gwas_regex)[2]
#=======================================================================
#%% command line args: case sensitive
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-d', '--drug', help = 'drug name', default = '')
arg_parser.add_argument('-g', '--gene', help = 'gene name', default = '') 

arg_parser.add_argument('--datadir', help = 'Data Directory. By default, it assmumes homedir + git/Data')
arg_parser.add_argument('-i', '--input_dir', help = 'Input dir containing pdb files. By default, it assmumes homedir + <drug> + input')
arg_parser.add_argument('-o', '--output_dir', help = 'Output dir for results. By default, it assmes homedir + <drug> + output')

arg_parser.add_argument('--debug', action ='store_true', help = 'Debug Mode')

args = arg_parser.parse_args()
#=======================================================================
#%% variable assignment: input and output 
drug    = args.drug
gene    = args.gene
datadir = args.datadir
indir   = args.input_dir
outdir  = args.output_dir

gene_match = gene + '_p.'
print('mut pattern for gene', gene, ':',  gene_match)

#%%=======================================================================
#==============
# directories
#==============
if not datadir:
    datadir = homedir + '/git/Data/'
    
if not indir:
    indir = datadir + drug + '/input/'
    
if not outdir:
    outdir = datadir + drug + '/output/'

#=======
# input
#=======
gene_list_normal = ['pnca', 'katg', 'rpob', 'alr']

if gene.lower() == "gid":
    print("\nReading mCSM file for gene:", gene)
    in_filename_mcsm = gene.lower() + '_complex_mcsm_norm_SRY.csv' # was incorrectly SAM previously

if gene.lower() == "embb":
    print("\nReading mCSM file for gene:", gene)
    #in_filename_mcsm = gene.lower() + '_complex_mcsm_norm1.csv' #798
    #in_filename_mcsm = gene.lower() + '_complex_mcsm_norm2.csv' #844
    in_filename_mcsm = gene.lower() + '_complex_mcsm_norm3.csv' #858 [got all muts on 15/05/2022]

if gene.lower() in gene_list_normal:
   print("\nReading mCSM file for gene:", gene)
   in_filename_mcsm = gene.lower() + '_complex_mcsm_norm.csv' 
    
infile_mcsm           = outdir + in_filename_mcsm
mcsm_df               = pd.read_csv(infile_mcsm, sep = ',')

in_filename_foldx     = gene.lower() + '_foldx.csv'
infile_foldx          = outdir + in_filename_foldx
foldx_df              = pd.read_csv(infile_foldx , sep = ',')

in_filename_deepddg   = gene.lower() + '_ni_deepddg.csv' # change to decent filename and put it in the correct dir
infile_deepddg        = outdir + in_filename_deepddg
deepddg_df            = pd.read_csv(infile_deepddg, sep = ',')

in_filename_dssp      = gene.lower() + '_dssp.csv'
infile_dssp           = outdir + in_filename_dssp
dssp_df_raw           = pd.read_csv(infile_dssp, sep = ',')

in_filename_kd        = gene.lower() + '_kd.csv'
infile_kd             = outdir + in_filename_kd
kd_df                 = pd.read_csv(infile_kd, sep = ',')

in_filename_rd        = gene.lower() + '_rd.csv'
infile_rd             = outdir + in_filename_rd
rd_df                 = pd.read_csv(infile_rd, sep = ',')

#in_filename_snpinfo   = 'ns' + gene.lower() + '_snp_info_f.csv' # gwas f info
#infile_snpinfo        = outdir + in_filename_snpinfo 

in_filename_afor      = gene.lower() + '_af_or.csv'
infile_afor           = outdir + in_filename_afor
afor_df               = pd.read_csv(infile_afor, sep = ',') 

#in_filename_afor_kin  = gene.lower() + '_af_or_kinship.csv'
#infile_afor_kin       = outdir + in_filename_afor_kin

infilename_dynamut2    = gene.lower() + '_dynamut2_norm.csv'
infile_dynamut2        = outdir + 'dynamut_results/dynamut2/' + infilename_dynamut2
dynamut2_df            = pd.read_csv(infile_dynamut2, sep = ',')

infilename_mcsm_f_snps = gene.lower() + '_mcsm_formatted_snps.csv'
infile_mcsm_f_snps     = outdir + infilename_mcsm_f_snps
mcsm_f_snps            = pd.read_csv(infile_mcsm_f_snps, sep = ',', names = ['mutationinformation'], header = None)

# more output added
## consurf [change colnames]? [add outcome category]
infilename_consurf     = gene.lower() + '_consurf_grades_f.csv'
infile_consurf         = outdir + 'consurf/'+ infilename_consurf
consurf_df             = pd.read_csv(infile_consurf, sep = ',')

## SNAP2 [add normalised score]
infilename_snap2       = gene.lower() + '_snap2_output.csv'
infile_snap2           = outdir + 'snap2/'+ infilename_snap2
snap2_df               = pd.read_csv(infile_snap2, sep = ',')

## PROVEAN
infilename_provean     = gene.lower() + '_provean.csv'
infile_provean         = outdir + 'provean/'+ infilename_provean
provean_df             = pd.read_csv(infile_provean, sep = ',',header = None )

# mmCSM-lig
infilename_mmcsm       = gene.lower() + '_mmcsm_results.csv'
infile_mmcsm           = outdir + 'mmcsm_lig/single_muts/'+ infilename_mmcsm
mmcsm_lig_raw          = pd.read_csv(infile_mmcsm, sep = ',')

#------------------------------------------------------------------------------
# ONLY: for gene 'gid' and 'rpob': End logic should pick this up!
geneL_na = ['gid', 'rpob']
if gene.lower() in geneL_na:
    print("\nGene:", gene.lower()
          , "\nReading mCSM_na files")

    infilename_mcsm_na    = gene.lower() + '_complex_mcsm_na_norm.csv' # gid
    infile_mcsm_na        = outdir + 'mcsm_na_results/' + infilename_mcsm_na 
    mcsm_na_df            = pd.read_csv(infile_mcsm_na, sep = ',')

geneL_dy = ['gid']
if gene.lower() in geneL_dy:
    print("\nGene:", gene.lower()
          , "\nReading Dynamut files")
    infilename_dynamut    = gene.lower() + '_dynamut_norm.csv' # gid
    infile_dynamut        = outdir + 'dynamut_results/' + infilename_dynamut
    dynamut_df            = pd.read_csv(infile_dynamut, sep = ',')

# ONLY: for genes 'alr', 'embb', 'katg' and 'rpob': End logic should pick this up!
geneL_ppi2 = ['alr', 'embb', 'katg', 'rpob']
if gene.lower() in geneL_ppi2:
    infilename_mcsm_ppi2   = gene.lower() + '_complex_mcsm_ppi2_norm.csv'
    infile_mcsm_ppi2       = outdir + 'mcsm_ppi2/' + infilename_mcsm_ppi2
    mcsm_ppi2_df           = pd.read_csv(infile_mcsm_ppi2, sep = ',')
    
if gene.lower() == "embb":
    sel_chain = "B"
else:
    sel_chain = "A"
#------------------------------------------------------------------------------
#=======
# output 
#=======
# outfile 3
out_filename_comb = gene.lower() + '_all_params.csv'
outfile_comb =  outdir + out_filename_comb

# outfile 2
out_filename_comb_afor = gene.lower() + '_comb_afor.csv'
outfile_comb_afor =  outdir + out_filename_comb_afor

# outfile 1
out_filename_stab_struc = gene.lower() + '_comb_stab_struc_params.csv'
outfile_stab_struc =  outdir + out_filename_stab_struc

# end of variable assignment for input and output files
#%%############################################################################  
print('\n==================================================='
      , '\nCombining dfs for gene:', gene.lower()
      , '\nDrug name:', drug
      , '\n====================================================')
#=====================
# variable setting
#=====================
gene_pdb_f = gene.lower() + '_complex.pdb'

gene_with_offset = ['alr', 'katg', 'rpob']
#geneL_snap2 = ['alr', 'katg', 'rpob']
gene_with_offset_missing_aa = ['embb']

if gene.lower() == 'alr':
    offset_val = 34

if gene.lower() == 'katg':
    offset_val = 23

if gene.lower() == 'rpob':
   offset_val = 28

if gene.lower() == 'embb':
   offset_val = 23
   missing_aa_start = 248
   missing_aa_end = 269
   missing_res_offset = missing_aa_end - missing_aa_start
   
if gene.lower() == 'gid':
   offset_val = 0
   
if gene.lower() == 'pnca':
   offset_val = 0

#===========
# DSSP
#===========
print('\nFormatting dssp_df:'
      , '\nSelecting chain:', sel_chain, 'for gene', gene.lower())

dssp_df = dssp_df_raw[dssp_df_raw['chain_id'] == sel_chain]
print('\nDim of dssp_df:', dssp_df.shape)
dssp_df['chain_id'].value_counts()
   
#===========
# KD
#===========
if gene.lower() in gene_with_offset:
    print ('\nRunning analysis for gene with offset:', gene.lower()
           , '\nOffset value:', offset_val  
           , '\nFormatting kd_df with offset value'
           , kd_df.shape)
    kd_df['position'] = kd_df['position'] + offset_val
    
if gene.lower() in gene_with_offset_missing_aa:
   print ('\nRunning analysis for gene with offset and missing aa:', gene.lower()
           , '\nOffset value:', offset_val  
           , '\nMissing residues gap:', missing_res_offset
           , '\nFormatting kd_df with offset value'
           , '\nDim of kd_df:', kd_df.shape)
   
   # First correct offset, then correct GAP
   print('\nFirst correcting offset')
   kd_df['position'] = kd_df['position'] + offset_val

   print('\nNow correcting gap')
   kd_df.loc[kd_df['position'] >= 248, 'position'] = kd_df['position'] + missing_res_offset

else:
    print ('\nRunning analysis for gene WITHOUT offset:', gene.lower()
           , '\nNo offset value:'
           , '\nUsing kd_df as is read'
           , kd_df.shape)
    
print('\nChecking the position numbers from the 3 dfs: dssp, kd and rd')   
dssp_kd_rd_pos_checks =  pd.DataFrame ({"dssp_pos": dssp_df['position']
                                       , "kd_pos": kd_df['position']
                                       , "rd_pos": rd_df['position']})
                                       
dkr_pos_check = [
    dssp_kd_rd_pos_checks['dssp_pos'].equals(dssp_kd_rd_pos_checks['kd_pos'])
    , dssp_kd_rd_pos_checks['dssp_pos'].equals(dssp_kd_rd_pos_checks['rd_pos'])
    , dssp_kd_rd_pos_checks['kd_pos'].equals(dssp_kd_rd_pos_checks['rd_pos'])
       ]
if all(dkr_pos_check):
    print('\nPASS: position check passed for dssp, rd and kd dfs ')
else:
    print('\nFAIL: position check Failed for dssp, rd and kd dfs'
             , '\nNumbering mismatch. Please check these dfs.')
    sys.exit()                                
#%%============================================================================   
#===========
# FoldX
#===========
foldx_df.shape

#----------------------
# scale foldx values
#----------------------
# rename ddg column to ddg_foldx
foldx_df['ddg']
foldx_df = foldx_df.rename(columns = {'ddg':'ddg_foldx'})   
foldx_df['ddg_foldx']
  
# Rescale values in Foldx_change col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
foldx_min = foldx_df['ddg_foldx'].min()
foldx_max = foldx_df['ddg_foldx'].max() 
foldx_min
foldx_max

# quick check
len(foldx_df.loc[foldx_df['ddg_foldx'] >= 0])
len(foldx_df.loc[foldx_df['ddg_foldx'] < 0])

foldx_scale = lambda x : x/abs(foldx_min) if x < 0 else (x/foldx_max if x >= 0 else 'failed')

foldx_df['foldx_scaled'] = foldx_df['ddg_foldx'].apply(foldx_scale)
print('\nRaw foldx scores:\n', foldx_df['ddg_foldx']
    , '\n---------------------------------------------------------------'
    , '\nScaled foldx scores:\n', foldx_df['foldx_scaled'])

# additional check added
fsmi = foldx_df['foldx_scaled'].min()
fsma = foldx_df['foldx_scaled'].max()

c = foldx_df[foldx_df['ddg_foldx']>=0].count()
foldx_pos = c.get(key = 'ddg_foldx')

c2 = foldx_df[foldx_df['foldx_scaled']>=0].count()
foldx_pos2 = c2.get(key = 'foldx_scaled')

if foldx_pos == foldx_pos2 and fsmi == -1 and fsma == 1:
    print('\nPASS: Foldx values scaled correctly b/w -1 and 1')
else:
    print('\nFAIL: Foldx values scaled numbers MISmatch'
          , '\nExpected number:', foldx_pos
          , '\nGot:', foldx_pos2
          , '\n======================================================')

#-------------------------
# foldx outcome category:
# Remember, its inverse
# +ve: Destabilising
# -ve: Stabilising
#--------------------------
foldx_df['foldx_outcome'] = foldx_df['ddg_foldx'].apply(lambda x: 'Destabilising' if x >= 0 else 'Stabilising')
foldx_df[foldx_df['ddg_foldx']>=0].count()
foc = foldx_df['foldx_outcome'].value_counts()

if foc['Destabilising'] == foldx_pos and foc['Destabilising'] == foldx_pos2:
    print('\nPASS: Foldx outcome category created')
else:
    print('\nFAIL: Foldx outcome category could NOT be created'
          , '\nExpected number:', foldx_pos
          , '\nGot:', foc[0]
          , '\n======================================================')
    sys.exit()

#=======================
# Deepddg
#=======================
deepddg_df.shape

#--------------------------
# check if >1 chain
#--------------------------
deepddg_df.loc[:,'chain_id'].value_counts()

if len(deepddg_df.loc[:,'chain_id'].value_counts()) > 1:
    print("\nChains detected: >1"
          , "\nGene:", gene
          , "\nChains:", deepddg_df.loc[:,'chain_id'].value_counts().index)
    
print('\nSelecting chain:', sel_chain, 'for gene:', gene)

deepddg_df = deepddg_df[deepddg_df['chain_id'] == sel_chain]

#--------------------------
# Drop chain_id col as other
# targets don't have it. 
#--------------------------
col_to_drop = ['chain_id']
deepddg_df = deepddg_df.drop(col_to_drop, axis = 1)

#--------------------------
# Check for duplicates
#--------------------------
if len(deepddg_df['mutationinformation'].duplicated().value_counts())> 1:
    print("\nFAIL: Duplicates detected in DeepDDG infile"
          , "\nNo. of duplicates:"
          , deepddg_df['mutationinformation'].duplicated().value_counts()[1]
          , "\nformat deepDDG infile before proceeding")
    sys.exit()
else:
    print("\nPASS: No duplicates detected in DeepDDG infile")

#-------------------------
# scale Deepddg values
#-------------------------
# Rescale values in deepddg_change col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
deepddg_min = deepddg_df['deepddg'].min() 
deepddg_max = deepddg_df['deepddg'].max() 

deepddg_scale = lambda x : x/abs(deepddg_min) if x < 0 else (x/deepddg_max if x >= 0 else 'failed')

deepddg_df['deepddg_scaled'] = deepddg_df['deepddg'].apply(deepddg_scale)
print('\nRaw deepddg scores:\n', deepddg_df['deepddg']
    , '\n---------------------------------------------------------------'
    , '\nScaled deepddg scores:\n', deepddg_df['deepddg_scaled'])
    
# additional check added
dsmi = deepddg_df['deepddg_scaled'].min()
dsma = deepddg_df['deepddg_scaled'].max()

c = deepddg_df[deepddg_df['deepddg']>=0].count()
deepddg_pos = c.get(key = 'deepddg')

c2 = deepddg_df[deepddg_df['deepddg_scaled']>=0].count()
deepddg_pos2 = c2.get(key = 'deepddg_scaled')
    
if deepddg_pos == deepddg_pos2 and dsmi == -1 and dsma == 1:
    print('\nPASS: deepddg values scaled correctly b/w -1 and 1')
else:
    print('\nFAIL: deepddg values scaled numbers MISmatch'
          , '\nExpected number:', deepddg_pos
          , '\nGot:', deepddg_pos2
          , '\n======================================================')    

if deepddg_df['deepddg_scaled'].min() == -1 and deepddg_df['deepddg_scaled'].max() == 1:
    print('\nPASS: Deepddg data is scaled between -1 and 1',
           '\nproceeding with merge')
    
#--------------------------
# Deepddg outcome category
#--------------------------
if 'deepddg_outcome' not in deepddg_df.columns:
    print('\nCreating column: deepddg_outcome')
    deepddg_df['deepddg_outcome'] = deepddg_df['deepddg'].apply(lambda x: 'Stabilising' if x >= 0 else 'Destabilising')
    deepddg_df[deepddg_df['deepddg']>=0].count()
    doc = deepddg_df['deepddg_outcome'].value_counts()
    print(doc)
else:
    print('\nColumn exists: deepddg_outcome')
    t1 = deepddg_df['deepddg_outcome'].value_counts()
    deepddg_df['deepddg_outcome2'] = deepddg_df['deepddg'].apply(lambda x: 'Stabilising' if x >= 0 else 'Destabilising')
    t2 = deepddg_df['deepddg_outcome2'].value_counts()
    print('\n', t1, '\n', t2)
    #--------------------------
    # Drop deepddg_outcome2 col 
    #--------------------------
    col_to_drop2 = ['deepddg_outcome2']
    deepddg_df = deepddg_df.drop(col_to_drop2, axis = 1)

if all(t1 == t2):
    print('\nPASS: Deepddg_outcome category checked!')
    doc = deepddg_df['deepddg_outcome'].value_counts()
else:
    print('\nMISmatch in deepddg_outcome counts'
          , '\n:', t1
          , '\n:', t2)

if doc['Stabilising'] == deepddg_pos and doc['Stabilising'] == deepddg_pos2:
    print('\nPASS: Deepddg outcome category created')
else:
    print('\nFAIL: Deepddg outcome category could NOT be created'
          , '\nExpected number:', deepddg_pos
          , '\nGot:', doc[0]
          , '\n======================================================')
    sys.exit()

#=======================
# Consurf
#=======================
consurf_df.shape

#----------------------
# rename colums
#----------------------   
consurf_df.columns
print('\nRenaming cols and assigning pretty column names')
 
if gene.lower() in gene_with_offset:
    consurf_df = consurf_df.rename(columns={'POS' : 'position_consurf'})
    #---------------------------
    # Specify the offset 
    #---------------------------
    print('\nFormatting consurf df'
          , '\nAdding offset value for gene:', gene.lower()
          , '\nUsing offset val:', offset_val)
        
    consurf_df['position'] = consurf_df['position_consurf'] + offset_val
    
else:
    consurf_df = consurf_df.rename(columns={'POS' : 'position'})

consurf_df = consurf_df.rename(columns={'SEQ'     : 'wild_type'
                                        , '3LATOM': 'wt_3upper'
                                        , 'SCORE' : 'consurf_score'
                                        , 'COLOR' : 'consurf_colour_str'
                                        , 'CONFIDENCEINTERVAL'       : 'consurf_ci'
                                        , 'CONFIDENCEINTERVALCOLORS' : 'consurf_ci_colour'
                                        , 'MSADATA'  : 'consurf_msa_data'
                                        , 'RESIDUEVARIETY' : 'consurf_aa_variety'})
# quick check
if len(consurf_df) == len(kd_df):
    print('\nPASS: length of consurf df is as expected'
          , '\nProceeding to format consurf df')
else:
    print('\nFAIL: length mismatch'
          , '\nExpected nrows:', len(rd_df)
          , '\nGot:', len(consurf_df))

consurf_df.dtypes
consurf_df['consurf_score'] = consurf_df['consurf_score'].astype(float)

consurf_df['consurf_colour'] = consurf_df['consurf_colour_str'].str.extract(r'(\d).*')
consurf_df['consurf_colour'] = consurf_df['consurf_colour'].astype(int)

consurf_df['consurf_colour_rev'] = consurf_df['consurf_colour_str'].str.replace(r'.\*','0')
# non struc position are assigned a *, replacing that with a 0 so its all integer
consurf_df['consurf_colour_rev'] = consurf_df['consurf_colour_rev'].astype(int)

consurf_df['consurf_ci_upper'] = consurf_df['consurf_ci'].str.extract(r'(.*):')
consurf_df['consurf_ci_upper'] = consurf_df['consurf_ci_upper'].astype(float)

consurf_df['consurf_ci_lower'] = consurf_df['consurf_ci'].str.extract(r':(.*)')
consurf_df['consurf_ci_lower'] = consurf_df['consurf_ci_lower'].astype(float)

#consurf_df['wt_3upper_f'] = consurf_df['wt_3upper'].str.extract(r'^\w{3}(\d+.*)')
#consurf_df['wt_3upper_f']
consurf_df['chain'] = consurf_df['wt_3upper'].str.extract(r':(.*)')

consurf_df['wt_3upper'] = consurf_df['wt_3upper'].str.replace(r'(\d+:.*)', '')

#-------------------------
# scale consurf values
#-------------------------
# Rescale values in consurf_score col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
consurf_min = consurf_df['consurf_score'].min()
consurf_max = consurf_df['consurf_score'].max() 
consurf_min
consurf_max

# quick check
len(consurf_df.loc[consurf_df['consurf_score'] >= 0])
len(consurf_df.loc[consurf_df['consurf_score'] < 0])

consurf_scale = lambda x : x/abs(consurf_min) if x < 0 else (x/consurf_max if x >= 0 else 'failed')

consurf_df['consurf_scaled'] = consurf_df['consurf_score'].apply(consurf_scale)
print('\nRaw consurf scores:\n', consurf_df['consurf_score']
    , '\n---------------------------------------------------------------'
    , '\nScaled consurf scores:\n', consurf_df['consurf_scaled'])

# additional check added
csmi = consurf_df['consurf_scaled'].min()
csma = consurf_df['consurf_scaled'].max()

c = consurf_df[consurf_df['consurf_score']>=0].count()
consurf_pos = c.get(key = 'consurf_score')

c2 = consurf_df[consurf_df['consurf_scaled']>=0].count()
consurf_pos2 = c2.get(key = 'consurf_scaled')

if consurf_pos == consurf_pos2 and csmi == -1 and csma == 1:
    print('\nPASS: Consurf values scaled correctly b/w -1 and 1')
else:
    print('\nFAIL: Consurf values scaled numbers MISmatch'
          , '\nExpected number:', consurf_pos
          , '\nGot:', consurf_pos2
          , '\n======================================================')

consurf_df.dtypes
consurf_df.columns

#---------------------------
# select columns 
# (and also determine order)
# this removes redundant cols:
    # consurf_colour_str
    # consurf_ci
#---------------------------
consurf_col_order = ['position'
                     , 'wild_type'
                     , 'chain'
                     , 'wt_3upper'
                     , 'consurf_score'
                     , 'consurf_scaled'
                     , 'consurf_colour'
                     , 'consurf_colour_rev'
                     , 'consurf_ci_upper'
                     , 'consurf_ci_lower'
                     , 'consurf_ci_colour'
                     , 'consurf_msa_data'
                     , 'consurf_aa_variety']
   
consurf_df_f = consurf_df[consurf_col_order]

# CHECK: whether a general rule or a gene specific rule!

if consurf_df_f['chain'].isna().sum() > 0:
    print('\nNaN detected in column chain for consurf df')
#if gene.lower() == 'embb':
    print('\nFurther consurf df processing for gene:', gene.lower())
    print('\nDropping Nan from column name chain')
    consurf_df_f = consurf_df_f[consurf_df_f['chain'].notna()]
   
#=======================
# SNAP2
#=======================
snap2_df.shape

#----------------------
# rename colums
#---------------------- 
if gene.lower() in gene_with_offset:
    print('\nReading SNAP2 for gene:', gene.lower()
          , '\nOffset column also being read'
          , '\nRenaming columns...'
          , '\nColumn mutationinformation exists. Renaming SNAP2 column variant --> mutationinformation')
    
    snap2_df = snap2_df.rename(columns = {'mutationinformation': 'mutationinformation'
                                          , 'Variant'          : 'mutationinformation_snap2'
                                          , 'Predicted Effect' : 'snap2_outcome'
                                          , 'Score'            : 'snap2_score'
                                          , 'Expected Accuracy': 'snap2_accuracy_pc'})
else:
    print('\nReading SNAP2 for gene:', gene.lower()
          , '\nNo offset column for SNAP2'
          , '\nRenaming columns...'
          , '\nRenaming SNAP2 column variant --> mutationinformation')

    snap2_df = snap2_df.rename(columns = {'Variant'            : 'mutationinformation'
                                          , 'Predicted Effect' : 'snap2_outcome'
                                          , 'Score'            : 'snap2_score'
                                          , 'Expected Accuracy': 'snap2_accuracy_pc'})
    
snap2_df.columns
snap2_df.head()
snap2_df.dtypes

snap2_df['snap2_accuracy_pc'] = snap2_df['snap2_accuracy_pc'].str.replace('%','')
snap2_df['snap2_accuracy_pc'] = snap2_df['snap2_accuracy_pc'].astype(int)

#-------------------------
# scale snap2 values
#-------------------------
# Rescale values in snap2_score col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
snap2_min = snap2_df['snap2_score'].min()
snap2_max = snap2_df['snap2_score'].max() 
snap2_min
snap2_max

# quick check
len(snap2_df.loc[snap2_df['snap2_score'] >= 0])
len(snap2_df.loc[snap2_df['snap2_score'] < 0])

snap2_scale = lambda x : x/abs(snap2_min) if x < 0 else (x/snap2_max if x >= 0 else 'failed')

snap2_df['snap2_scaled'] = snap2_df['snap2_score'].apply(snap2_scale)
print('\nRaw snap2 scores:\n', snap2_df['snap2_score']
    , '\n---------------------------------------------------------------'
    , '\nScaled snap2 scores:\n', snap2_df['snap2_scaled'])

# additional check added
ssmi = snap2_df['snap2_scaled'].min()
ssma = snap2_df['snap2_scaled'].max()

sn = snap2_df[snap2_df['snap2_score']>=0].count()
snap2_pos = sn.get(key = 'snap2_score')

sn2 = snap2_df[snap2_df['snap2_scaled']>=0].count()
snap2_pos2 = sn2.get(key = 'snap2_scaled')

if snap2_pos == snap2_pos2 and csmi == -1 and csma == 1:
    print('\nPASS: Snap2 values scaled correctly b/w -1 and 1')
else:
    print('\nFAIL: snap2 values scaled numbers MISmatch'
          , '\nExpected number:', snap2_pos
          , '\nGot:', snap2_pos2
          , '\n======================================================')

#-------------------------------------
# select columns 
# (and also determine order)
# renumbering already done using 
# bash and corrected file is read in
#-------------------------------------
snap2_df.dtypes
snap2_df.columns

if gene.lower() in gene_with_offset :
    print('\nSelecting cols SNAP2 for gene:', gene.lower())
    snap2_df_f = snap2_df[['mutationinformation'
                       , 'mutationinformation_snap2'
                       , 'snap2_score'
                       , 'snap2_scaled'
                       , 'snap2_accuracy_pc'
                       , 'snap2_outcome']]
else:
    print('\nSelecting cols SNAP2 for gene:', gene.lower())
    snap2_df_f = snap2_df[['mutationinformation'
                       , 'snap2_score'
                       , 'snap2_scaled'
                       , 'snap2_accuracy_pc'
                       , 'snap2_outcome']]

#=======================
# Provean
#=======================
provean_df.head()
provean_df.columns = ['mutationinformation', 'provean_score', 'provean_outcome']
provean_df.head()
provean_df['provean_outcome'].value_counts()

#----------------------------------------
# Rescale values in provean_score
# col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
# cut off =-2.5
# so provean scores >= (-2.5) are neutral
# and provean scores < (2.5) are deleterious
#-----------------------------------------
provean_min = provean_df['provean_score'].min()
provean_max = provean_df['provean_score'].max() 
print('\nprovean_score (MIN):', provean_min
      , '\nprovean_score (MAX):', provean_max)

# quick check
provean_cut_off = -2.5

if (provean_df['provean_score'] > provean_cut_off).sum() == provean_df['provean_outcome'].value_counts()['Neutral']:
    print('\nPASS: Provean cut off is indeed:', provean_cut_off
          , '\nNo. of values above', provean_cut_off, 'i.e classed as Neutral:'
          , (provean_df['provean_score']>provean_cut_off).sum()
          , '\nProvean outcome:'
          , '\nNeutral:', len(provean_df.loc[provean_df['provean_score'] > provean_cut_off])
          , '\nDeleterious:', len(provean_df.loc[provean_df['provean_score'] < provean_cut_off])
          ) 
else:
    sys.exit('\nFAIL: Numbers mismatch. Please check provean cut off and condition used')

# RECHECK logic!: CANNOT use this as it changes the data distribution, as seen from the his plot
# provean_scale = lambda x : x/abs(provean_min) if x < 0 else (x/provean_min if x >= 0 else 'failed')
# provean_df['provean_scaled1'] = provean_df.loc[:,'provean_score'].apply(provean_scale)
# print('\nRaw provean scores:\n', provean_df['provean_score']
#     , '\n---------------------------------------------------------------'
#     , '\nScaled provean scores:\n', provean_df['provean_scaled'])

# print('\nprovean raw (Max):'      , provean_df['provean_score'].max()
#       , '\nprovean scaled (Max):' , provean_df['provean_scaled1'].max())
# print('\nprovean raw (Min):'      , provean_df['provean_score'].min()
#       , '\nprovean scaled (Min):' , provean_df['provean_scaled1'].min())

scaler = MinMaxScaler()
provean_df['provean_scaled'] = scaler.fit_transform(provean_df['provean_score'].values.reshape(-1,1))

provean_df['provean_score'].hist(bins = 30)
#provean_df['provean_scaled1'].hist(bins = 10)
provean_df['provean_scaled'].hist(bins = 30)

#=======================
# mmCSM-lig
#=======================
mmcsm_lig_raw.columns

# extract specific columns: might be simpler
mmcsm_lig_df = mmcsm_lig_raw[['MUTATION', 'CHAIN', 'DDG']]
mmcsm_lig_df['CHAIN'].value_counts()

# Drop the chain column
mmcsm_lig_df.drop(['CHAIN'], axis = 1, inplace = True)

# Rename columns using lower case and consistently to allow merge later on
mmcsm_lig_df.rename({'MUTATION': 'mutationinformation'
                  , 'DDG': 'mmcsm_lig'}, axis = 1, inplace = True)

#----------------------------------------
# Rescale values in mmcsm_lig_affinity
# col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
#-----------------------------------------
mmcsm_lig_min = mmcsm_lig_df['mmcsm_lig'].min()
mmcsm_lig_max = mmcsm_lig_df['mmcsm_lig'].max() 
print('\nmmcsm_lig (MIN):', mmcsm_lig_min
      , '\nmmcsm_lig (MAX):', mmcsm_lig_max)

# quick check
print('\nNo. of Stabilising mmCSM mutations:', len(mmcsm_lig_df.loc[mmcsm_lig_df['mmcsm_lig'] >= 0]))
print('\nNo. of Destabilising mmCSM mutations:', len(mmcsm_lig_df.loc[mmcsm_lig_df['mmcsm_lig'] < 0]))

mmcsm_ligscale = lambda x : x/abs(mmcsm_lig_min) if x < 0 else (x/mmcsm_lig_max if x >= 0 else 'failed')

mmcsm_lig_df['mmcsm_lig_scaled'] = mmcsm_lig_df.loc[:,'mmcsm_lig'].apply(mmcsm_ligscale)
print('\nRaw mmcsm_lig scores:\n', mmcsm_lig_df['mmcsm_lig']
    , '\n---------------------------------------------------------------'
    , '\nScaled mmcsm_lig scores:\n', mmcsm_lig_df['mmcsm_lig_scaled'])

print('\nmmCSM lig raw (Max):', mmcsm_lig_df['mmcsm_lig'].max()
      , '\nmmCSM lig scaled (Max):', mmcsm_lig_df['mmcsm_lig_scaled'].max())

print('\nmmCSM lig raw (Min):', mmcsm_lig_df['mmcsm_lig'].min()
      , '\nmmCSM lig scaled (Min):', mmcsm_lig_df['mmcsm_lig_scaled'].min())

mmcsm_lig_df['mmcsm_lig_scaled'].hist(bins = 30)
mmcsm_lig_df['mmcsm_lig'].hist(bins = 30)

#-----------------------------
# mmCSM lig outcome category:
# -ve: Destabilising
# +ve: Stabilising
#----------------------------
mmcsm_lig_df['mmcsm_lig_outcome'] = mmcsm_lig_df.loc[:,'mmcsm_lig'].apply(lambda x: 'Destabilising' if x < 0 else 'Stabilising')
mmcsm_lig_df[mmcsm_lig_df['mmcsm_lig']<0].count()

del(mmcsm_lig_raw)

#%%============================================================================
# Now merges begin
print('==================================='
      , '\nFirst merge: mcsm + foldx'
      , '\n===================================')

mcsm_df.shape

# add 3 lowercase aa code for wt and mutant
get_aa_3lower(df = mcsm_df
              , wt_colname = 'wild_type'
              , mut_colname = 'mutant_type'
              , col_wt = 'wt_aa_3lower'
              , col_mut = 'mut_aa_3lower')

#mcsm_df.columns = mcsm_df.columns.str.lower()
# foldx_df.shape

#mcsm_foldx_dfs = combine_dfs_with_checks(mcsm_df, foldx_df, my_join = "outer")
merging_cols_m1  = detect_common_cols(mcsm_df, foldx_df)
mcsm_foldx_dfs   = pd.merge(mcsm_df
                          , foldx_df
                          , on = merging_cols_m1
                          , how =  "outer")
ncols_m1 = len(mcsm_foldx_dfs.columns)

print('\n\nResult of first merge:', mcsm_foldx_dfs.shape
      , '\n===================================================================')
mcsm_foldx_dfs[merging_cols_m1].apply(len)
mcsm_foldx_dfs[merging_cols_m1].apply(len) == len(mcsm_foldx_dfs)

#%% for embB: mCSM-lig hasn't run for ALL nsSNPs.
# Get the empty cells to be full of meaningful info
if mcsm_foldx_dfs.loc[:,'wild_type': 'mut_aa_3lower'].isnull().values.any():  
    print ('\nNAs detected in mcsm cols after merge.'
           , '\nCleaning data before merging deepddg_df')
    
    ##############################
    # Extract relevant col values 
    # code to one
    ##############################

    # wt_reg = r'(^[A-Z]{1})'
    # print('wild_type:', wt_reg)

    # mut_reg = r'[0-9]+(\w{1})$'
    # print('mut type:', mut_reg)
    mcsm_foldx_dfs['wild_type']     = mcsm_foldx_dfs.loc[:,'mutationinformation'].str.extract(r'(^[A-Z]{1})')
    mcsm_foldx_dfs['position']      = mcsm_foldx_dfs.loc[:,'mutationinformation'].str.extract(r'([0-9]+)')
    mcsm_foldx_dfs['mutant_type']   = mcsm_foldx_dfs.loc[:,'mutationinformation'].str.extract(r'[0-9]+([A-Z]{1})$')
    
    # BEWARE: Bit of logic trap i.e if nan comes first
    # in chain column, then nan will be populated!
    #df['foo'] = df['chain'].unique()[0]
    mcsm_foldx_dfs['chain'] = np.where(mcsm_foldx_dfs[['chain']].isnull().all(axis=1)
                                        , mcsm_foldx_dfs['chain'].unique()[0]
                                        , mcsm_foldx_dfs['chain'])
    
    mcsm_foldx_dfs['ligand_id'] = np.where(mcsm_foldx_dfs[['ligand_id']].isnull().all(axis=1)
                                        , mcsm_foldx_dfs['ligand_id'].unique()[0]
                                        , mcsm_foldx_dfs['ligand_id'])
    #--------------------------------------------------------------------------
    
    mcsm_foldx_dfs['wild_pos']       = mcsm_foldx_dfs.loc[:,'wild_type'] + mcsm_foldx_dfs.loc[:,'position'].astype(int).apply(str)
    mcsm_foldx_dfs['wild_chain_pos'] = mcsm_foldx_dfs.loc[:,'wild_type'] + mcsm_foldx_dfs.loc[:,'chain'] +  mcsm_foldx_dfs.loc[:,'position'].astype(int).apply(str)
    
    #############
    # Map 1 letter 
    # code to 3Upper
    #############
    # initialise a sub dict that is lookup dict for 
    # 3-LETTER aa code to 1-LETTER aa code
    lookup_dict = dict()
    for k, v in oneletter_aa_dict.items():
        lookup_dict[k] = v['three_letter_code_lower']
        wt = mcsm_foldx_dfs['wild_type'].squeeze() # converts to a series that map works on
        mcsm_foldx_dfs['wt_aa_3lower'] = wt.map(lookup_dict)   
        mut = mcsm_foldx_dfs['mutant_type'].squeeze()
        mcsm_foldx_dfs['mut_aa_3lower'] = mut.map(lookup_dict)
else:
    print('\nNo NAs detected in mcsm_fold_dfs. Proceeding to merge deepddg_df')
    
#%%============================================================================
print('==================================='
      , '\nSecond merge: mcsm_foldx_dfs + deepddg'
      , '\n===================================')

# merge with mcsm_foldx_dfs and deepddg_df
mcsm_foldx_deepddg_dfs = pd.merge(mcsm_foldx_dfs
                                  , deepddg_df
                                  , on = 'mutationinformation'
                                  , how = "left")

mcsm_foldx_deepddg_dfs['deepddg_outcome'].value_counts()

mcsm_foldx_deepddg_dfs['position'] = mcsm_foldx_deepddg_dfs['position'].astype('int64')

#%%============================================================================
# Select df with 'chain' to allow corret dim merging!
print('==================================='
      , '\nThird merge: dssp + kd'
      , '\n===================================')

dssp_df_raw.shape
kd_df.shape
rd_df.shape

#dssp_kd_dfs = combine_dfs_with_checks(dssp_df, kd_df, my_join = "outer")
merging_cols_m2 = detect_common_cols(dssp_df, kd_df)
dssp_kd_dfs = pd.merge(dssp_df
                       , kd_df
                       , on = merging_cols_m2
                       #, how = "outer")
                       , how = "inner")

print('\n\nResult of third merge:', dssp_kd_dfs.shape
      , '\n===================================================================')
#%%============================================================================
print('==================================='
      , '\nFourth merge: third merge + rd_df' 
      , '\ndssp_kd_dfs + rd_df'
      , '\n===================================')
#dssp_kd_rd_dfs = combine_dfs_with_checks(dssp_kd_dfs, rd_df, my_join = "outer")
merging_cols_m3 = detect_common_cols(dssp_kd_dfs,  rd_df)
dssp_kd_rd_dfs = pd.merge(dssp_kd_dfs
                          , rd_df
                          , on = merging_cols_m3
                          , how = "outer")

ncols_m3 = len(dssp_kd_rd_dfs.columns)

print('\n\nResult of Third merge:', dssp_kd_rd_dfs.shape
      , '\n===================================================================')
dssp_kd_rd_dfs[merging_cols_m3].apply(len)
dssp_kd_rd_dfs[merging_cols_m3].apply(len) == len(dssp_kd_rd_dfs)

#%%============================================================================
print('==================================='
      , '\nFifth merge: fourth merge + consurf_df' 
      , '\dssp_kd_rd_dfs + consurf_df'
      , '\n===================================')
#dssp_kd_rd_dfs = combine_dfs_with_checks(dssp_kd_dfs, rd_df, my_join = "outer")
merging_cols_m3_v2 = detect_common_cols(dssp_kd_rd_dfs,  consurf_df)
dssp_kd_rd_con_dfs = pd.merge(dssp_kd_rd_dfs
                          , consurf_df_f
                          , on = merging_cols_m3_v2 
                          , how = "outer")

print('\n\nResult of fourth merge*:', dssp_kd_rd_con_dfs.shape
      , '\n===================================================================')
dssp_kd_rd_con_dfs[merging_cols_m3_v2].apply(len)
dssp_kd_rd_con_dfs[merging_cols_m3_v2].apply(len) == len(dssp_kd_rd_con_dfs)

#%%============================================================================
print('======================================='
      , '\nSixth merge: Second merge + fifth merge'
      , '\nmcsm_foldx_dfs + dssp_kd_rd_con_dfs'
      , '\n=======================================')

#combined_df = combine_dfs_with_checks(mcsm_foldx_dfs, dssp_kd_rd_dfs, my_join = "inner")
#merging_cols_m4 = detect_common_cols(mcsm_foldx_dfs, dssp_kd_rd_dfs)
#combined_df = pd.merge(mcsm_foldx_dfs, dssp_kd_rd_dfs, on = merging_cols_m4, how  = "inner")
#combined_df_expected_cols = ncols_m1 + ncols_m3 - len(merging_cols_m4)

# Detect merging cols
merging_cols_m4 = detect_common_cols(mcsm_foldx_deepddg_dfs, dssp_kd_rd_con_dfs)

if len(merging_cols_m4) > 1 and all(x in merging_cols_m4 for x in ['chain', 'position', 'wild_type']) :
    print('\nlength of merging_cols_m4 is > 1 and has mixed dtype'
          , '\nDropping chain and wild_type from dssp_kd_rd_con_dfs')
    
    dssp_kd_rd_con_dfs = dssp_kd_rd_con_dfs.drop(['chain', 'wild_type'], axis = 1)
    print('\nDetecting merging cols again')
    merging_cols_m4 = detect_common_cols(mcsm_foldx_deepddg_dfs, dssp_kd_rd_con_dfs)
    print('\nMerging cols length:', len(merging_cols_m4))

if len(merging_cols_m4) == 1 and 'position' in merging_cols_m4:
    print('\nMerging column length == 1. Using it'
        , '\nMerging column name:', merging_cols_m4)
else:
    sys.exit('\nFAIL: to merge mcsm_foldx_deepddg and dssp_kd_rd_con_dfs. Please check')
    
    
combined_7dfs = pd.merge(mcsm_foldx_deepddg_dfs
                       , dssp_kd_rd_con_dfs
                       , on = merging_cols_m4
                       , how  = "inner")

if len(combined_7dfs.columns) == len(mcsm_foldx_deepddg_dfs.columns) + len(dssp_kd_rd_con_dfs.columns) - len(merging_cols_m4):
    print('\n Seven dfs successfully combined:'
          , '\nmcsm-lig', '\nfoldx', '\ndeepddg'
          , '\ndssp', '\nkd', '\nrd'
          , '\nconsurf')
else:
    print('\nSomething went wrong with the seventh merge'
          , '\nPlease check individual merges')
    sys.exit()  
#%%============================================================================
print('======================================='
      , '\nSeventh merge: Sixth merge + snap2_df_f'
      , '\ncombined_7dfs + snap_df_f'
      , '\n=======================================')

# Detect merging cols
merging_cols_m4_v2 = detect_common_cols(combined_7dfs, snap2_df_f)

if len(merging_cols_m4_v2) == 1:
    print ('\nProceeding with merging SNAP2 df with 7 already combined dfs')
else:
    print('\nHalting Seventh merge. Please check merging cols'
          ,'\nLength of merging_cols_m4_v2 is:', len(merging_cols_m4_v2))
    sys.exit()

combined_df = pd.merge(combined_7dfs
                       , snap2_df_f
                       , on = merging_cols_m4_v2
                       , how  = "inner")    

combined_df_expected_cols = len(combined_7dfs.columns) + len(snap2_df_f.columns) - len(merging_cols_m4_v2)
   
#%% Dimension checks for combined_df: specially because embB didn't have mCSM 
# values for all muts. This was fixed in UQ, and this rerun on 15/05/2022
#------------------------------------------------------------------------------
# No longer required since mCSM values for fetched for all muts!
# CHECK whether logic effects anything else!
# if not gene == "embB":
#     print("\nGene is:", gene)
#     if len(combined_df) == len(mcsm_df) and len(combined_df.columns) == combined_df_expected_cols:
#         print('PASS: successfully combined 8 dfs'
#               , '\nNo. of rows combined_df:', len(combined_df)
#               , '\nNo. of cols combined_df:', len(combined_df.columns))   
# elif len(combined_df) == len(foldx_df) and len(combined_df.columns) == combined_df_expected_cols:
#     print('PASS: successfully combined 8 dfs'
#                , '\nNo. of rows combined_df:', len(combined_df)
#                , '\nNo. of cols combined_df:', len(combined_df.columns))
# else:
#     sys.exit('FAIL: check individual merges for seventh merge')      
#------------------------------------------------------------------------------
print("\nRunning Dimension check for combined_df for gene:", gene)

if len(combined_df) == len(mcsm_df) and len(combined_df.columns) == combined_df_expected_cols:
    print('PASS: successfully combined 8 dfs'
          , '\nNo. of rows combined_df:', len(combined_df)
          , '\nNo. of cols combined_df:', len(combined_df.columns))   
else:
    sys.exit('FAIL: check individual merges for seventh merge')    

print('\n====================================================')

print('\nResult of Seventh merge:', combined_df.shape
      , '\n===================================================================')

# should be TRUE and length match mcsm_df len
combined_df[merging_cols_m4_v2].apply(len)
combined_df[merging_cols_m4_v2].apply(len) == len(combined_df)

#%%============================================================================
# Format the combined df columns
combined_df_colnames = combined_df.columns

# check redundant columns
combined_df['chain'].equals(combined_df['chain_id'])
combined_df['wild_type'].equals(combined_df['wild_type_kd']) # should not have nan
combined_df['wild_type'].equals(combined_df['wild_type_dssp'])

# sanity check
foo = combined_df[['wild_type', 'wild_type_kd', 'wt_3letter_caps', 'wt_aa_3lower', 'mut_aa_3lower']] 

# Drop cols
cols_to_drop = ['chain_id', 'wild_type_kd', 'wild_type_dssp'
                , 'wt_3letter_caps', 'wt_aa_3lower', 'mut_aa_3lower']

print('\nDropping:', len(cols_to_drop), 'cols to get clean_df for futher combining'
      , '\nncols before dropping:', len(combined_df.columns)
      , '\nDropped column names:', cols_to_drop)

combined_df_clean = combined_df.drop(cols_to_drop, axis = 1)

print('\nncols after dropping:', len(combined_df_clean.columns))

combined_df_clean_colnames = combined_df_clean.columns
del(foo)
#%%============================================================================
#---------------------
# Output 1: write csv
#---------------------
print('\nWriting file: combined stability and structural parameters'
      , '\nOutput 1 filename:', outfile_stab_struc
      , '\n===================================================================\n')

combined_df_clean.to_csv(outfile_stab_struc, index = False)
print('\nFinished writing file:'
      , '\nNo. of rows:', combined_df_clean.shape[0]
      , '\nNo. of cols:', combined_df_clean.shape[1])
#%%=====================================================================
print('\n======================================='
      , '\nFifth merge:'
      , '\ncombined_df_clean + afor_df '
      , '\n=======================================')

afor_cols = afor_df.columns
afor_df.shape

# create a mapping from the gwas mutation column i.e <gene_match>_abcXXXrst
#----------------------
# call get_aa_upper(): 
# adds 3 more cols with one letter aa code
#----------------------
get_aa_1upper(df = afor_df
               , gwas_mut_colname = 'mutation'
               , wt_colname = 'wild_type'
               , pos_colname = 'position'
               , mut_colname = 'mutant_type')

afor_df['mutationinformation'] = afor_df['wild_type'] + afor_df['position'].map(str) + afor_df['mutant_type']
afor_cols = afor_df.columns

merging_cols_m5 = detect_common_cols(combined_df_clean, afor_df)

# remove position so that merging can take place without dtype conflicts
merging_cols_m5.remove('position')

# drop position column from afor_df
afor_df = afor_df.drop(['position'], axis = 1)
afor_cols = afor_df.columns

# merge 
combined_stab_afor = pd.merge(combined_df_clean
                              , afor_df
                              , on = merging_cols_m5
                              , how  = "left")

comb_afor_df_cols = combined_stab_afor.columns

comb_afor_expected_cols = len(combined_df_clean.columns) + len(afor_df.columns) - len(merging_cols_m5)

if len(combined_stab_afor) == len(combined_df_clean) and len(combined_stab_afor.columns) == comb_afor_expected_cols:
    print('\nPASS: successfully combined 6 dfs'
          , '\nNo. of rows combined_stab_afor:', len(combined_stab_afor)
          , '\nNo. of cols combined_stab_afor:', len(combined_stab_afor.columns))   
else:
    sys.exit('\nFAIL: check individual df merges')

print('\n\nResult of Seventh merge:', combined_stab_afor.shape
      , '\n===================================================================')

combined_stab_afor[merging_cols_m5].apply(len)
combined_stab_afor[merging_cols_m5].apply(len) == len(combined_stab_afor)

if (len(combined_stab_afor) - combined_stab_afor['mutation'].isna().sum()) == len(afor_df):
     print('\nPASS: Merge successful for af and or with matched numbers')

if len(combined_stab_afor) - combined_stab_afor['mutation'].isna().sum() == len(afor_df)-len(afor_df[~afor_df['mutation'].isin(combined_stab_afor['mutation'])]):
    print("\nMismatched numbers, OR df has extra snps not found in mcsm df"
          , "\nNo. of nsSNPs with valid ORs:", len(afor_df)
          , "\nNo. of mcsm nsSNPs: ", len(combined_df_clean) 
          , "\nNo. of OR nsSNPs not in mCSM df:"
          , len(afor_df[~afor_df['mutation'].isin(combined_stab_afor['mutation'])])
          , "\nWriting these mutations to file:")
    orsnps_notmcsm = afor_df[~afor_df['mutation'].isin(combined_stab_afor['mutation'])]
else:
    sys.exit('\nFAIL: merge unsuccessful for af and or')    
    
#%%============================================================================
#---------------------
# Output 2: write csv
# when dynamut, dynamut2 and others weren't being combined
#---------------------
print('\nWriting file: combined stability and afor'
      , '\nOutput 2 filename:', outfile_comb_afor
      , '\n===================================================================\n')

combined_stab_afor.to_csv(outfile_comb_afor, index = False)
print('\nFinished writing file:'
      , '\nNo. of rows:', combined_stab_afor.shape[0]
      , '\nNo. of cols:', combined_stab_afor.shape[1])
#%%============================================================================
# combine dynamut, dynamut2, and mcsm_na
#dfs_list = [dynamut_df, dynamut2_df, mcsm_na_df] # gid

if gene.lower() == "pnca":
    dfs_list = [dynamut2_df]
if gene.lower() == "gid":
    dfs_list = [dynamut_df, dynamut2_df, mcsm_na_df]
if gene.lower() == "embb":
    dfs_list = [dynamut2_df, mcsm_ppi2_df]
if gene.lower() == "katg":
    dfs_list = [dynamut2_df, mcsm_ppi2_df]
if gene.lower() == "rpob":
    dfs_list = [dynamut2_df, mcsm_na_df, mcsm_ppi2_df]
if gene.lower() == "alr":
    dfs_list = [dynamut2_df, mcsm_ppi2_df]

# noticed that with revised rpoB that mcsm-NA had one less position,
# Hence this condition else the last check fails with discrepancy for expected_nrows
if len(dfs_list) > 1:
    join_type = 'outer'
else:
    join_type = 'inner'

print('\nUsing join type: "', join_type, '" for the last but one merge')
    
dfs_merged = reduce(lambda  left,right: pd.merge(left
                                                , right
                                                , on = ['mutationinformation']
                                                #, how = 'inner')
                                                , how = join_type) , dfs_list)
# drop excess columns
drop_cols = detect_common_cols(dfs_merged, combined_stab_afor)
drop_cols.remove('mutationinformation')

dfs_merged_clean = dfs_merged.drop(drop_cols, axis = 1)
merging_cols_m6 = detect_common_cols(dfs_merged_clean, combined_stab_afor)

len(combined_stab_afor.columns)
len(dfs_merged_clean.columns)

combined_all_params = pd.merge(combined_stab_afor
                               , dfs_merged_clean
                               , on = merging_cols_m6
                               , how  = "inner")

expected_ncols = len(dfs_merged_clean.columns) + len(combined_stab_afor.columns) - len(merging_cols_m6)
expected_nrows = len(combined_stab_afor)

if len(combined_all_params.columns) == expected_ncols and len(combined_all_params) == expected_nrows:
    print('\nPASS: All dfs combined')
else:
    print('\nFAIL:lengths mismatch'
          , '\nExpected ncols:', expected_ncols
          , '\nGot:', len(dfs_merged_clean.columns) 
          , '\nExpected nrows:', expected_nrows
          , '\nGot:', len(dfs_merged_clean) )

# Drop cols
if combined_all_params.columns.str.contains(r'_x$|_y$', regex = True).any():
     print('\nDuplicate column names detected...'
           , '\nDropping these before writing file')
     extra_cols_to_drop = list(combined_all_params.columns.str.extract(r'(.*_x$|.*_y$)', expand = True).dropna()[0])
     print('\nTotal cols:',  len(combined_all_params.columns)
         ,'\nDropping:', len(extra_cols_to_drop), 'columns'
         , '\nThese are:', extra_cols_to_drop)
     #extra_cols_to_drop = ['chain_x', 'chain_y']
     combined_all_params = combined_all_params.drop(extra_cols_to_drop, axis = 1)
else:
    print('\nNo duplicate column names detected, just writing file'
          , '\nTotal cols:', len(combined_all_params.columns) )
#%%==============================================================================
# Some final formatting for combined_all_params df

#---------------------------------------
# Add: Map 1 letter 
# code to 3Upper for Mutant aa
#---------------------------------------
# initialise a sub dict that is lookup dict for 
# 3-LETTER aa code to 1-LETTER aa code

lookup_dict = dict()
for k, v in oneletter_aa_dict.items():
    lookup_dict[k] = v['three_letter_code_upper']
    mut = combined_all_params['mutant_type'].squeeze()
    combined_all_params['mut_3upper'] = mut.map(lookup_dict)
else:
    print('\nFailed to add 3 letter upper code')

# Add offset value
combined_all_params['seq_offset4pdb'] = offset_val

# Drop extra columns
redundant_colsL = ['mutationinformation_snap2'
                   , 'wt_upper'
                   , 'mut_upper'
                   , 'pdb_file']

combined_all_params_f = combined_all_params.drop(redundant_colsL
                                                 , axis = 1
                                                 , errors = 'ignore')
#---------------------------------------
# MERGE 7 [UQ]: provean and mmCSM-lig dfs
#---------------------------------------
if len(combined_all_params_f) == len(provean_df) == len(mmcsm_lig_df):
    print('\nPASS: length of Provean and mmCSM-lig df length match with combined_df.'
          , '\nProceeding with FINAL merging before writing file...')
else:
    sys.exit('\nFAIL: Cannot do final merge! Check lengths of consurf, provean and combined_all_params_f dfs')
    
merging_cols_m7 = detect_common_cols(provean_df, mmcsm_lig_df)
print('\nMering provean and mmcsm-lig dfs on:', merging_cols_m7)

pr_mm_df = pd.merge(provean_df
                    , mmcsm_lig_df
                    , on = merging_cols_m7)

#-----------------------------------------------
# MERGE 8 [UQ]: combined_all_params_f + pr_mm_df
#-----------------------------------------------
merging_cols_m8 = detect_common_cols(combined_all_params_f, pr_mm_df)
print('\nMering all combined_dfs + (pr_mm_df) on:', merging_cols_m8)

combined_all_params_f2  = pd.merge(combined_all_params_f
                    , pr_mm_df
                    , on = merging_cols_m8)

expected_ncols = len(combined_all_params_f.columns) + len(pr_mm_df.columns) - len(merging_cols_m8)
expected_nrows = len(combined_all_params_f2)

if len(combined_all_params_f2.columns) == expected_ncols and len(combined_all_params_f2) == expected_nrows:
    print('\nPASS: All dfs combined including PROVEAN and mmCSM-lig')
else:
    print('\nFAIL:lengths mismatch'
          , '\nExpected ncols:', expected_ncols
          , '\nGot:', len(combined_all_params_f2.columns) 
          , '\nExpected nrows:', expected_nrows
          , '\nGot:', len(combined_all_params_f2) )
    
#---------------------------------------
# Add gene name
#---------------------------------------
combined_all_params_f2['gene_name'] = gene.lower()
    
#---------------------------------------
# Add pdb_file name at the end
#---------------------------------------
combined_all_params_f2['pdb_file'] = gene_pdb_f

#%%============================================================================
#---------------------
# Output 3: write csv
#---------------------
print('\nWriting file: all params')
print('\nOutput 3 filename:', outfile_comb
      , '\n===================================================================\n')

combined_all_params_f2.to_csv(outfile_comb, index = False)

print('\nFinished writing file:'
      , '\nNo. of rows:', combined_all_params_f2.shape[0]
      , '\nNo. of cols:', combined_all_params_f2.shape[1])
#%% end of script
