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

@author: tanu
'''
# FIXME: import dirs.py to get the basic dir paths available
#=======================================================================
# TASK: calculate how many mutations result in 
# electrostatic changes wrt wt

# Input: mCSM-normalised file or any file containing one-letter aa code 
# of wt and mut
# Note: this can be easily modified into 3-letter code
# TODO: turn to a function

# Output: mut_elec_changes_results.txt 
#=======================================================================
#%% load libraries
import os, sys
import pandas as pd
#import numpy as np
#from varname import nameof
import argparse
#=======================================================================
#%% specify input and curr dir
homedir = os.path.expanduser('~')

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

from reference_dict import oneletter_aa_dict 
#=======================================================================
#%% command line args
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-d', '--drug', help='drug name', default = None)
arg_parser.add_argument('-g', '--gene', help='gene name', default = None)

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') # not used atm

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

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

#=======
# input
#=======
#in_filename  = 'merged_df3.csv'
#in_filename  = gene.lower() + '_complex_mcsm_norm.csv'
in_filename  = gene.lower() + '_complex_mcsm_norm_SRY.csv' # gid
infile_merged_df3 = outdir + '/' + in_filename
print('Input file: ', infile_merged_df3
      , '\n============================================================')

#=======
# output 
#=======
out_filename = gene.lower() + '_mut_elec_changes.txt'
outfile_elec_changes =  outdir + '/' + out_filename
print('Output file: ', outfile_elec_changes
      , '\n============================================================')

#%% end of variable assignment for input and output files
#=======================================================================
#%% Read input files
print('Reading input file (merged file):', infile_merged_df3)

comb_df = pd.read_csv(infile_merged_df3, sep = ',')  

print('Input filename: ', in_filename
      , '\nPath :', outdir
      , '\nNo. of rows: ', len(comb_df)
      , '\nNo. of cols: ', len(comb_df.columns)
      , '\n============================================================')

# column names
list(comb_df.columns)

# clear variables
del(in_filename, infile_merged_df3)
#%%
#----------------------------------------------------------------
# add aa properties considering df has columns:
# 'wild_type', 'mutant_type' separately as single letter aa code
#----------------------------------------------------------------
lookup_dict = dict()
for k, v in oneletter_aa_dict.items():
    lookup_dict[k] = v['aa_calcprop']
    #print(lookup_dict)
    comb_df['wt_calcprop'] = comb_df['wild_type'].map(lookup_dict)   
    comb_df['mut_calcprop'] = comb_df['mutant_type'].map(lookup_dict)

#%% subset unique mutations
df = comb_df.drop_duplicates(['mutationinformation'], keep = 'first')

total_muts = df['mutationinformation'].nunique()
#df.Mutationinformation.count()
print('Total mutations associated with structure: ', total_muts
	, '\n===============================================================')

#%% combine aa_calcprop cols so that you can count the changes as value_counts
# check if all muts have been categorised
print('Checking if all muts have been categorised: ')
if df['wt_calcprop'].isna().sum() == 0 & df['mut_calcprop'].isna().sum():
    print('PASS: No. NA detected. All muts have aa prop associated'
    , '\n===============================================================')
else:
    print('FAIL: NAs detected. Some muts remain unclassified'
    , '\n===============================================================')
    sys.exit()

df['wt_calcprop'].head()
df['mut_calcprop'].head()

print('Combining wt_calcprop and mut_calcprop...')
#df['aa_calcprop_combined'] = df['wt_calcprop']+ '->' + df['mut_calcprop']
df['aa_calcprop_combined'] = df.wt_calcprop.str.cat(df.mut_calcprop, sep = '->')
df['aa_calcprop_combined'].head()

mut_categ = df["aa_calcprop_combined"].unique()
print('Total no. of aa_calc properties: ', len(mut_categ))
print('Categories are: ', mut_categ)

# Counting no. of muts in each mut categ

# way1: count values within each combinaton
df.groupby('aa_calcprop_combined').size()
#df.groupby('aa_calcprop_combined').count()

# way2: count values within each combinaton
df['aa_calcprop_combined'].value_counts()

# comment: the two ways should be identical
# groupby result order is similar to pivot table order,
# I prefer the value_counts look

# assign to variable: count values within each combinaton
all_prop = df['aa_calcprop_combined'].value_counts()

# convert to a df from Series
ap_df = pd.DataFrame({'aa_calcprop': all_prop.index, 'mut_count': all_prop.values})

# subset df to contain only the changes in prop
all_prop_change = ap_df[ap_df['aa_calcprop'].isin(['neg->neg','non-polar->non-polar','polar->polar', 'pos->pos']) == False]
                      
elec_count = all_prop_change.mut_count.sum()
print('Total no.of muts with elec changes: ', elec_count)

# calculate percentage of electrostatic changes
elec_changes = (elec_count/total_muts) * 100
           
print('\n==============================================================='
      , '\nResults for electrostatic changes from nsSNPs for', gene,':'
      , '\n\nTotal no. of mutations:', total_muts
      , '\nMutations with electrostatic changes:', elec_count
      , '\nPercentage(%) of electrostatic changes:', elec_changes
      , '\n===============================================================')
 
# check no change muts
no_change_muts = ap_df[ap_df['aa_calcprop'].isin(['neg->neg','non-polar->non-polar','polar->polar', 'pos->pos']) == True]
  
no_change_muts.mut_count.sum()

if no_change_muts.mut_count.sum() + elec_count == total_muts:
    print('\nPASS: numbers cross checked'
          , '\nWriting output file:', outfile_elec_changes)
else:
    print('\nFAIL: numbers mismatch')
    sys.exit()
                     
#%% output from console
#sys.stdout = open(file, 'w')
sys.stdout = open(outfile_elec_changes, 'w')
     
print('################################################'
      , '\nResults for gene:', gene 
      , '\n################################################' 
      , '\n\n======================'
      , '\nUnchanged muts:'
      , '\n=====================\n'
      ,  no_change_muts
      , '\n=============================\n'      
      , 'Muts with changed properties:'
      , '\n============================\n'
      , all_prop_change
      , '\n=====================================================')

print('Total no. of mutations: ', total_muts
      , '\nTotal no. of mutions with electrostatic changes:', elec_count
      , '\nCorresponding (%):', elec_changes
      , '\nTotal no. of changed muts: ', all_prop_change.mut_count.sum()
      , '\nTotal no. of unchanged muts: ', no_change_muts.mut_count.sum() 
      , '\n=======================================================')  
#%% end of script
#=======================================================================