#!/usr/bin/env python
# Description = Calculates literature based secondary structures derived
# from residue information

import os
from collections import defaultdict


class LiteratureSecstrAnalysis(object):
    def __init__(self, selected_proteins,
                 fasta_path, outpath):
        self._selected_proteins = selected_proteins
        self._fasta_path = fasta_path
        self._outpath = outpath
        
        self._selected_protein_set = set()
        self._literature_acid_str = {
            'H': ['E', 'A', 'L', 'M', 'Q', 'K', 'R', 'H', 'Z'],
            'E': ['V', 'I', 'Y', 'C', 'W', 'F', 'T'],
            'C': ['G', 'N', 'P', 'S', 'D', 'B']
        }
        # 1: alpha-helix (H), 2: beta-sheet(E), 3: turns and loops(C)
        # https://en.wikibooks.org/wiki/Structural_Biochemistry/Proteins/Structures
        
    def streamline_literature_secstr_analysis(self):
        '''To call from apricot'''
        self.parse_selected_data()
        self.run_literature_ss_analysis()
        self.create_job_completion_file()
        
    def parse_selected_data(self):
        '''Parses selected data for uid'''
        with open(self._selected_proteins, 'r') as in_fh:
            for entry in in_fh:
                if not entry.startswith('Entry'):
                    self._selected_protein_set.add(entry.split('\t')[0])
        return self._selected_protein_set
    
    def run_literature_ss_analysis(self):
        '''Runs Literature based  3-state secondary
        structure predictionon the selected uids'''
        for files in os.listdir(self._fasta_path):
            if files.split('.')[0] in self._selected_protein_set:
                print("Literature acid based secondary structure "
                      "analysis for %s" % files)
                with open(self._outpath+'/%s.plain' % files.split('.')[0],
                          'w') as out_fh:
                    out_fh.write("#Literature based three-class secondary "
                                 "structure prediction results\n\n")
                    with open("%s/%s" % (
                            self._fasta_path, files), 'r') as in_fh:
                        literature_acid_seq = []
                        for entry in in_fh:
                            if not entry.startswith('>'):
                                literature_acid_seq.append(entry.strip())
                        str_list = []
                        seq_dict = defaultdict(
                            lambda: defaultdict(lambda: defaultdict()))
                        for aa in ''.join(literature_acid_seq):
                            str_val = ''
                            if aa == 'J':
                                str_val = 'x'
                            elif aa == 'X':
                                str_val = 'x'
                            else:
                                for sec_str in self._literature_acid_str.keys():
                                    if aa in self._literature_acid_str[sec_str]:
                                        str_val = sec_str
                            str_list.append(str_val)
                        full_seq = self._split_str(
                            ''.join(literature_acid_seq), 60)
                        full_pred = self._split_str(''.join(str_list), 60)
                        for i, seq in enumerate(full_pred):
                            aa_seq = '  AA: %s' % (full_seq[i])
                            pred = 'Pred: %s' % (full_pred[i])
                            out_fh.write("%s\n%s\n\n" % (pred, aa_seq))
                            
    def _split_str(self, seq, chunk, skip_tail=False):
        lst = []
        if chunk <= len(seq):
            lst.extend([seq[:chunk]])
            lst.extend(self._split_str(seq[chunk:], chunk, skip_tail))
        elif not skip_tail and seq:
            lst.extend([seq])
        return lst
    
    def create_job_completion_file(self):
        with open(self._outpath+'/literature_ss_analysis.txt', 'w') as out_fh:
            out_fh.write("Literature based secondary structures for the "
                         "selected proteins are generated.\n")
            out_fh.write("The files generated by the analysis:.\n")
            out_fh.write('\n'.join(os.listdir(self._outpath)))
