#!/usr/bin/env python3
#=======================================================================
#TASK: 
#=======================================================================
#%% load packages
import os,sys
import subprocess
import argparse
#import requests
import re
#import time
import pandas as pd
from pandas.api.types import is_string_dtype
from pandas.api.types import is_numeric_dtype
import numpy as np

#=======================================================================
#%% specify input and curr dir
homedir = os.path.expanduser('~')
# set working dir
os.getcwd()
os.chdir(homedir + '/git/LSHTM_analysis/mcsm')
os.getcwd()
#=======================================================================
#%% variable assignment: input and output 
#drug = 'pyrazinamide'
#gene = 'pncA'

drug = 'rifampicin'
gene = 'rpoB'

#drug = args.drug
#gene = args.gene

gene_match = gene + '_p.'
#==========
# data dir
#==========
datadir = homedir + '/' + 'git/Data'

#=======
# input:
#=======
# 1) result_urls (from outdir)
outdir = datadir + '/' + drug + '/' + 'output'
in_filename = gene.lower() + '_mcsm_output.csv' #(outfile, from mcsm_results.py)
infile = outdir + '/' + in_filename
print('Input filename:', in_filename
      , '\nInput path(from output dir):', outdir
      , '\n=============================================================')
      
#=======
# output 
#=======
outdir =   datadir + '/' + drug + '/' + 'output'
out_filename = gene.lower() + '_complex_mcsm_norm.csv'
outfile =  outdir + '/' + out_filename
print('Output filename:', out_filename
      , '\nOutput path:', outdir
      , '\n=============================================================')

#=======================================================================
print('Reading input file')
mcsm_data  = pd.read_csv(infile, sep = ',')

mcsm_data.columns
# PredAffLog = affinity_change_log
# "DUETStability_Kcalpermol = DUET_change_kcalpermol
dforig_shape = mcsm_data.shape
print('dim of infile:', dforig_shape) 

#############
# rename cols
#############
# format colnames: all lowercase, remove spaces and use '_' to join 
print('Assigning meaningful colnames i.e without spaces and hyphen and reflecting units'
		  , '\n===================================================================')
my_colnames_dict = {'Predicted Affinity Change': 'PredAffLog' # relevant info from this col will be extracted and the column discarded
	           , 'Mutation information': 'mutation_information' # {wild_type}<position>{mutant_type}
	           , 'Wild-type': 'wild_type' # one letter amino acid code
	           , 'Position': 'position' # number
	           , 'Mutant-type': 'mutant_type' # one letter amino acid code
	           , 'Chain': 'chain' # single letter (caps)
	           , 'Ligand ID': 'ligand_id' # 3-letter code
	           , 'Distance to ligand': 'ligand_distance' # angstroms
	           , 'DUET stability change': 'duet_stability_change'} # in kcal/mol

mcsm_data.rename(columns = my_colnames_dict, inplace = True)
#%%===========================================================================
# populate mutation_information column:mcsm style  muts {WT}<POS>{MUT}
print('Populating column : mutation_information which is currently empty\n', mcsm_data['mutation_information'])
mcsm_data['mutation_information'] = mcsm_data['wild_type'] +  mcsm_data['position'].astype(str) + mcsm_data['mutant_type']
print('checking after populating:\n', mcsm_data['mutation_information']
	  , '\n===================================================================')

# Remove spaces b/w pasted columns
print('removing white space within column: \mutation_information')
mcsm_data['mutation_information'] = mcsm_data['mutation_information'].str.replace(' ', '')
print('Correctly formatted column: mutation_information\n', mcsm_data['mutation_information']
	  , '\n===================================================================')
#%% Remove whitespace from column
#orig_dtypes = mcsm_data.dtypes
#https://stackoverflow.com/questions/33788913/pythonic-efficient-way-to-strip-whitespace-from-every-pandas-data-frame-cell-tha/33789292
#mcsm_data.columns = mcsm_data.columns.str.strip()
#new_dtypes = mcsm_data.dtypes
#%%===========================================================================
# very important
print('Sanity check:'
	  , '\nChecking duplicate mutations')
if mcsm_data['mutation_information'].duplicated().sum() == 0:
	print('PASS: No duplicate mutations detected (as expected)'
	      , '\nDim of data:', mcsm_data.shape
	      , '\n===============================================================')
else:
	print('FAIL (but not fatal): Duplicate mutations detected'
	      , '\nDim of df with duplicates:', mcsm_data.shape
	      , 'Removing duplicate entries')
	mcsm_data = mcsm_data.drop_duplicates(['mutation_information'])
	print('Dim of data after removing duplicate muts:', mcsm_data.shape
		, '\n===============================================================')
#%%=========================================================================== 
# create duet_outcome column: classification based on DUET stability values
print('Assigning col: duet_outcome based on DUET stability values')
print('Sanity check:')
# count positive values in the DUET column
c = mcsm_data[mcsm_data['duet_stability_change']>=0].count()
DUET_pos = c.get(key = 'duet_stability_change')
# Assign category based on sign (+ve : Stabilising, -ve: Destabilising, Mind the spelling (British spelling))
mcsm_data['duet_outcome'] = np.where(mcsm_data['duet_stability_change']>=0, 'Stabilising', 'Destabilising')
mcsm_data['duet_outcome'].value_counts()
if DUET_pos == mcsm_data['duet_outcome'].value_counts()['Stabilising']:
	print('PASS: DUET outcome assigned correctly')
else:
	print('FAIL: DUET outcome assigned incorrectly'
	      , '\nExpected no. of stabilising mutations:', DUET_pos
	      , '\nGot no. of stabilising mutations', mcsm_data['duet_outcome'].value_counts()['Stabilising']
	      , '\n===============================================================')
#%%=========================================================================== 
# Extract only the numeric part from col: ligand_distance
# number: '-?\d+\.?\d*'
mcsm_data['ligand_distance']
print('extracting numeric part of col: ligand_distance')
mcsm_data['ligand_distance'] = mcsm_data['ligand_distance'].str.extract('(\d+\.?\d*)')
mcsm_data['ligand_distance']
#%%=========================================================================== 
# create ligand_outcome column: classification based on affinity change values
# the numerical and categorical parts need to be extracted from column: PredAffLog
# regex used 
# number: '-?\d+\.?\d*'
# category: '\b(\w+ing)\b'
print('Extracting numerical and categorical parts from the col: PredAffLog')
print('to create two columns: ligand_affinity_change and ligand_outcome'
	  , '\n===================================================================')
	# Extracting the predicted affinity change (numerical part)
mcsm_data['ligand_affinity_change'] = mcsm_data['PredAffLog'].str.extract('(-?\d+\.?\d*)', expand = True)
print(mcsm_data['ligand_affinity_change'])
# Extracting the categorical part (Destabillizing and Stabilizing) using word boundary ('ing')
#aff_regex = re.compile(r'\b(\w+ing)\b')
mcsm_data['ligand_outcome']= mcsm_data['PredAffLog'].str.extract(r'(\b\w+ing\b)', expand = True)
print(mcsm_data['ligand_outcome'])
print(mcsm_data['ligand_outcome'].value_counts())

# ensuring spellings are consistent
american_spl = mcsm_data['ligand_outcome'].value_counts()
print('Changing to Bristish spellings for col: ligand_outcome')
mcsm_data['ligand_outcome'].replace({'Destabilizing': 'Destabilising', 'Stabilizing': 'Stabilising'}, inplace = True)
print(mcsm_data['ligand_outcome'].value_counts())
british_spl = mcsm_data['ligand_outcome'].value_counts()
# compare series values since index will differ from spelling change
check = american_spl.values == british_spl.values
if check.all():
	print('PASS: spelling change successfull'
	      , '\nNo. of predicted affinity changes:\n', british_spl
	      , '\n===============================================================')
else:
	print('FAIL: spelling change unsucessfull'
	      , '\nExpected:\n', american_spl
	      , '\nGot:\n', british_spl
	      , '\n===============================================================')
#%%===========================================================================
# check dtype in cols: ensure correct dtypes for cols
print('Checking dtypes in all columns:\n', mcsm_data.dtypes
	  , '\n===================================================================')
#1) numeric cols
print('Converting the following cols to numeric:'
	  , '\nligand_distance'
	  , '\nduet_stability_change'
	  , '\nligand_affinity_change'
	  , '\n===================================================================')
# using apply method  to change stabilty and affinity values to numeric
numeric_cols = ['duet_stability_change', 'ligand_affinity_change', 'ligand_distance']
mcsm_data[numeric_cols] = mcsm_data[numeric_cols].apply(pd.to_numeric)

# check dtype in cols
print('checking dtype after conversion')
cols_check = mcsm_data.select_dtypes(include='float64').columns.isin(numeric_cols)
if cols_check.all():
	print('PASS: dtypes for selected cols:', numeric_cols
	      , '\nchanged to numeric'
	      , '\n===============================================================')
else:
	print('FAIL:dtype change to numeric for selected cols unsuccessful'
	      , '\n===============================================================')
#mcsm_data['ligand_distance', 'ligand_affinity_change'].apply(is_numeric_dtype(mcsm_data['ligand_distance', 'ligand_affinity_change']))
print(mcsm_data.dtypes)
#%%===========================================================================
# Normalise values in DUET_change col b/w -1 and 1 so negative numbers
# stay neg and pos numbers stay positive
duet_min = mcsm_data['duet_stability_change'].min() 
duet_max = mcsm_data['duet_stability_change'].max() 

duet_scale = lambda x : x/abs(duet_min) if x < 0 else (x/duet_max if x >= 0 else 'failed')

mcsm_data['duet_scaled'] = mcsm_data['duet_stability_change'].apply(duet_scale)
print('Raw duet scores:\n', mcsm_data['duet_stability_change']
	, '\n---------------------------------------------------------------'
	, '\nScaled duet scores:\n', mcsm_data['duet_scaled'])
#%%===========================================================================
# Normalise values in affinity change col b/w -1 and 1 so negative numbers 
# stay neg and pos numbers stay positive
aff_min = mcsm_data['ligand_affinity_change'].min()
aff_max = mcsm_data['ligand_affinity_change'].max() 

aff_scale = lambda x : x/abs(aff_min) if x < 0 else (x/aff_max if x >= 0 else 'failed')

mcsm_data['ligand_affinity_change']
mcsm_data['affinity_scaled'] = mcsm_data['ligand_affinity_change'].apply(aff_scale)
mcsm_data['affinity_scaled']
print('Raw affinity scores:\n', mcsm_data['ligand_affinity_change']
	, '\n---------------------------------------------------------------'
	, '\nScaled affinity scores:\n', mcsm_data['affinity_scaled'])
#=============================================================================
# Adding colname: wild_pos: sometimes useful for plotting and db
print('Creating column: wild_position')
mcsm_data['wild_position'] = mcsm_data['wild_type'] + mcsm_data['position'].astype(str)
print(mcsm_data['wild_position'].head())
# Remove spaces b/w pasted columns
print('removing white space within column: wild_position')
mcsm_data['wild_position'] = mcsm_data['wild_position'].str.replace(' ', '')
print('Correctly formatted column: wild_position\n', mcsm_data['wild_position'].head()
	  , '\n===================================================================')
#=============================================================================
#%% ensuring dtypes are string for the non-numeric cols
#) char cols
char_cols = ['PredAffLog', 'mutation_information', 'wild_type', 'mutant_type', 'chain'
             , 'ligand_id', 'duet_outcome', 'ligand_outcome', 'wild_position']

#mcsm_data[char_cols] = mcsm_data[char_cols].astype(str)
cols_check_char = mcsm_data.select_dtypes(include='object').columns.isin(char_cols)

if cols_check_char.all():
	print('PASS: dtypes for char cols:', char_cols, 'are indeed string'
	      , '\n===============================================================')
else:
	print('FAIL:dtype change to numeric for selected cols unsuccessful'
	      , '\n===============================================================')
#mcsm_data['ligand_distance', 'ligand_affinity_change'].apply(is_numeric_dtype(mcsm_data['ligand_distance', 'ligand_affinity_change']))
print(mcsm_data.dtypes)
#%%
#=============================================================================
# Removing PredAff log column as it is not needed?
print('Removing col: PredAffLog since relevant info has been extracted from it')
mcsm_dataf = mcsm_data.drop(columns = ['PredAffLog'])
#%%===========================================================================
expected_ncols_toadd = 5 # beware of hardcoded numbers
dforig_len = dforig_shape[1]
expected_cols = dforig_len + expected_ncols_toadd
if len(mcsm_dataf.columns) == expected_cols:
	print('PASS: formatting successful'
	, '\nformatted df has expected no. of cols:', expected_cols
	, '\ncolnames:', mcsm_dataf.columns
	, '\n----------------------------------------------------------------'
	, '\ndtypes in cols:', mcsm_dataf.dtypes
	, '\n----------------------------------------------------------------'
	, '\norig data shape:', dforig_shape
	, '\nformatted df shape:', mcsm_dataf.shape
	, '\n===============================================================')
else: 
	print('FAIL: something went wrong in formatting df'
	, '\nLen of orig df:', dforig_len
	, '\nExpected number of cols to add:', expected_ncols_toadd
	, '\nExpected no. of cols:', expected_cols, '(', dforig_len, '+', expected_ncols_toadd, ')'
	, '\nGot no. of cols:', len(mcsm_dataf.columns)
	, '\nCheck formatting:'
	, '\ncheck hardcoded value:', expected_ncols_toadd
	, '\nis', expected_ncols_toadd, 'the no. of expected cols to add?'
	, '\n===============================================================')
#%%============================================================================
# writing file
print('Writing formatted df to csv')
mcsm_dataf.to_csv(outfile, index = False)

print('Finished writing file:'
      , '\nFile:', outfile
      , '\nExpected no. of rows:', len(mcsm_dataf)
      , '\nExpected no. of cols:', len(mcsm_dataf.columns)
      , '\n=============================================================')
#%%
#End of script