#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue May  8 14:51:13 2018

@author: miri-o
"""

## generate a unified data set containing the following fields: subject, 
# condition, junc_AA[2:-1]

from Bio.Seq import Seq
from Bio.Alphabet import IUPAC
import pandas as pd
import datetime


directory = '/media/miri-o/Documents/celiac_100_makedb/'
file_path = '/media/miri-o/Documents/CDR3_from_Celiac_full.csv'+datetime.datetime.now().strftime("%Y-%m-%d-%H%M")
newfile = open(file_path, 'w')  
newfile.write('SAMPLE\tSUBJECT\tCONDITION\tJUNCTION\tJUNC_AA\n')
newfile.close() 
file_dict = pd.read_csv('/media/miri-o/Documents/Celiac_files_list_PARSED.csv', header = 0, index_col=0)
log_file_path = '/media/miri-o/Documents/logs/Create_Celiac_CDR3_file'+datetime.datetime.now().strftime("%Y-%m-%d-%H%M")
log_file = open(log_file_path, 'w')
total_length = 0
total_length_filtered = 0

for index, line in file_dict.iterrows():
    filename = line.filename
    print('~~~~~~~~\nAdding file: ' + filename)
    log_file.write('~~~~~~~~\nAdding file: \n' + filename)
    fil = pd.read_csv(directory+filename,sep = '\t')
    orig_len = len(fil)
    
    if orig_len < 2000:
        print('File too short (%d sequences), skipping file'%orig_len)
        log_file.write('File too short (%d sequences), skipping file'%orig_len)
        continue
    total_length += orig_len
    print('Original length: '+ str(len(fil)))
    log_file.write('Original length: '+ str(len(fil))+'\n')
    # Filtering the file:
    fil  = fil[fil['FUNCTIONAL']=='T']
    fil = fil[fil['CONSCOUNT']>1]   
    fil = fil[fil.JUNCTION_LENGTH%3==0]
    fil = fil[fil.JUNCTION_LENGTH>15]
    a = [('N' not in fil.JUNCTION.iloc[i]) for i in range(len(fil.JUNCTION))]
    a = pd.Series(a, index = fil.index.values)
    fil = fil[a]
    a = [('-' not in fil.JUNCTION.iloc[i]) for i in range(len(fil.JUNCTION))]
    a = pd.Series(a, index = fil.index.values)
    fil = fil[a]
    a = [('*' not in fil.JUNCTION.iloc[i]) for i in range(len(fil.JUNCTION))]
    a = pd.Series(a, index = fil.index.values)
    fil = fil[a]
    fil['SUBJECT'] = line.subject
    fil['SAMPLE'] = index
    fil['CONDITION'] = line['Condition']
    fil['coding_dna'] = [Seq(junc, IUPAC.unambiguous_dna) for junc in fil['JUNCTION']]
    fil['JUNC_AA'] = [str(coding_dna.translate())[2:-1] for coding_dna in fil['coding_dna']]
    a = [('*' not in fil.JUNC_AA.iloc[i]) for i in range(len(fil.JUNC_AA))]
    a = pd.Series(a, index = fil.index.values)
    fil = fil[a]
    total_length_filtered += len(fil)
    print('After filtering : '+ str(len(fil)) + '    ' + str(100*(orig_len-len(fil))/orig_len) + '% data lost')
    log_file.write('After filtering : '+ str(len(fil)) + '    ' + str(100*(orig_len-len(fil))/orig_len) + '% data lost\n')
    with open(file_path, 'a') as f:
        fil.to_csv(f, header = False, sep = '\t', 
                   columns = ['SAMPLE','SUBJECT', 'CONDITION', 'JUNCTION', 'JUNC_AA'], index = False)
log_file.write('=====\nNumber of filtered CDR3 sequences: {0}, out of {1} (%{2:.2g})'.format(total_length_filtered, total_length, 
               100*(total_length_filtered)/total_length))       
log_file.close()
print('\nFinished successfully, log created at: ' + log_file_path +'\nModel saved to: ' + file_path)
print('=====\nNumber of filtered CDR3 sequences: {0}, out of {1} (%{2:.2g})'.format(total_length_filtered, total_length, 
      100*(total_length_filtered)/total_length))

    
            