#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Created on Mon June 14 2021

@author: tanu
'''
# FIXME: import dirs.py to get the basic dir paths available
#=======================================================================
# TASK

# Input:

# Output: 
#=======================================================================
#%% load libraries
import os, sys
import pandas as pd
#import numpy as np
#from varname import nameof
import argparse

DEBUG = False
#=======================================================================
#%% 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 
#=======================================================================
#%%
def get_aa_prop(df, col1 = 'aap1', col2 = 'aap2', col3 = 'aap_taylor', col4 = 'aap_kd', col5 = 'aap_polarity', col6 = 'aap_calcprop'):

    """Add amino acid properties for wt and mutant residues
    
    @df: df containing one letter aa code for wt and mutant respectively
    @type: pandas df
    
    @col1: column adding 7 aa categories (no overlap; acidic, basic, amidic, hydrophobic, hydroxylic, aromatic, sulphur)
    @type: str
    
    @col2: column adding 9 aa categories (overlap; acidic, basic, polar, hydrophobic, hydrophilic, small, aromatic, aliphatic, special)
    @type: str
    
    @col3: column adding 8 aa categories (overlap; acidic, basic, polar, hydrophobic, small, aromatic, aliphatic, special)
    @type: str

    @col4: column adding 3 aa categories (no overlap, hydrophobic, neutral and hydrophilic according to KD scale)
    @type: str

    @col5: column adding 4 aa categories (no overlap, acidic, basic, neutral, non-polar)
    @type: str

    @col6: column adding 4 aa categories (neg, pos, polar, non-polar)
    @type: str
    
    returns df: with 6 added columns. If column names clash, the function column 
                name will override original column
    @rtype: pandas df
    """

    lookup_dict_p1 = dict()
    lookup_dict_p2 = dict()
    
    lookup_dict_taylor = dict()
    lookup_dict_kd = dict()
    
    lookup_dict_polarity = dict()
    lookup_dict_calcprop = dict()
    
    for k, v in oneletter_aa_dict.items():
      
        lookup_dict_p1[k] = v['aa_prop1']
        
        lookup_dict_p2[k] = v['aa_prop2']
              
        lookup_dict_taylor[k] = v['aa_taylor']
        
        lookup_dict_kd[k] = v['aa_prop_water']
                
        lookup_dict_polarity[k] = v['aa_prop_polarity']
        
        lookup_dict_calcprop[k] = v['aa_calcprop']      
        
        #if DEBUG:
        #    print('Key:', k, 'value:', v
        #          , '\n============================================================'
        #          , '\nlook up dict:\n')
            
        df['wt_aap1'] = df['wild_type'].map(lookup_dict_p1)   
        df['mut_aap1'] = df['mutant_type'].map(lookup_dict_p1)    
            
        df['wt_aap2'] = df['wild_type'].map(lookup_dict_p2)   
        df['mut_aap2'] = df['mutant_type'].map(lookup_dict_p2)   

        df['wt_aap_taylor'] = df['wild_type'].map(lookup_dict_taylor)   
        df['mut_aap_taylor'] = df['mutant_type'].map(lookup_dict_taylor)   

        df['wt_aap_kd'] = df['wild_type'].map(lookup_dict_kd)   
        df['mut_aap_kd'] = df['mutant_type'].map(lookup_dict_kd)   

        df['wt_aap_polarity'] = df['wild_type'].map(lookup_dict_polarity)   
        df['mut_aap_polarity'] = df['mutant_type'].map(lookup_dict_polarity)           
                    
        df['wt_aa_calcprop'] = df['wild_type'].map(lookup_dict_calcprop)   
        df['mut_aa_calcprop'] = df['mutant_type'].map(lookup_dict_calcprop)

    return df
#========================================