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

@author: tanu
'''
#=============================================================================
# Task: Residue depth (rd) processing to generate a df with residue_depth(rd)
# values

# FIXME
# Input: '.tsv' i.e residue depth txt file (output from .zip file manually
# downloaded from the website). 
# This should be integrated into the pipeline

# Output: .csv with 3 cols i.e position, rd_values & 3-letter wt aa code(caps)
#=============================================================================
#%% load packages
import sys, os
import pandas as pd
#import numpy as np
#=============================================================================
#%% specify input and curr dir
homedir = os.path.expanduser('~')

# set working dir
os.getcwd()
os.chdir(homedir + '/git/LSHTM_analysis/meta_data_analysis')
os.getcwd()
#=============================================================================
#%% variable assignment: input and output 
drug = 'pyrazinamide'
gene = 'pncA'
gene_match = gene + '_p.'

#==========
# data dir
#==========
#indir = 'git/Data/pyrazinamide/input/original'
datadir = homedir + '/' + 'git/Data'

#=======
# input
#=======
#indir = 'git/Data/pyrazinamide/input/original'
indir = datadir + '/' + drug + '/' + 'output'
in_filename = '3pl1_rd.tsv'
infile = indir + '/' + in_filename
print('Input filename:', in_filename
      , '\nInput path:', indir
      , '\n=============================================================')
#=======
# output 
#=======
outdir =   datadir + '/' + drug + '/' + 'output'
out_filename = gene.lower() + '_rd.csv'
outfile =  outdir + '/' + out_filename
print('Output filename:', out_filename
      , '\nOutput path:', outdir
      , '\n=============================================================')
      
#%% end of variable assignment for input and output files
#=======================================================================
#%%  Read input file
rd_data  = pd.read_csv(infile, sep = '\t')  
print('Reading input file:', infile
      , '\nNo. of rows:', len(rd_data)
      , '\nNo. of cols:', len(rd_data.columns))

print('Column names:', rd_data.columns
	, '\n===============================================================')
#========================
# creating position col
#========================
# Extracting residue number from index and assigning 
# the values to a column [position]. Then convert the position col to numeric.
rd_data['position'] = rd_data.index.str.extract('([0-9]+)').values

# converting position to numeric      
rd_data['position'] = pd.to_numeric(rd_data['position'])
rd_data['position'].dtype

print('Extracted residue num from index and assigned as a column:'
      , '\ncolumn name: position'
      , '\ntotal no. of cols now:', len(rd_data.columns)
      , '\n=============================================================')

#========================
# Renaming amino-acid 
# and all-atom cols
#========================
print('Renaming columns:'
      , '\ncolname==> # chain:residue: wt_3letter_caps'
      , '\nYES... the column name *actually* contains a # ..!'
      , '\ncolname==> all-atom: rd_values'
      , '\n=============================================================')

rd_data.rename(columns = {'# chain:residue':'wt_3letter_caps', 'all-atom':'rd_values'}, inplace = True)
print('Column names:', rd_data.columns)

#========================
# extracting df with the
# desired columns
#========================
print('Extracting relevant columns for writing df as csv')

rd_df = rd_data[['position','rd_values','wt_3letter_caps']]

if len(rd_df) == len(rd_data):
    print('PASS: extracted df has expected no. of rows'
          ,'\nExtracted df dim:'
          ,'\nNo. of rows:', len(rd_df)
          ,'\nNo. of cols:', len(rd_df.columns))
else:
    print('FAIL: no. of rows mimatch'
          , '\nExpected no. of rows:', len(rd_data)
          , '\nGot no. of rows:', len(rd_df)
          , '\n=========================================================')

#%% write file
print('Writing file:'
      , '\nFilename:', out_filename
      , '\nPath:',  outdir
      , '\n=============================================================')

rd_df.to_csv(outfile, header = True, index = False)

print('Finished writing:', out_filename
      , '\nNo. of rows:', len(rd_df)
      , '\nNo. of cols:', len(rd_df.columns)
      , '\n=============================================================')

#%% end of script
#=======================================================================