"""
A set of classes for reading in hmm profile files generated by hmmbuild.
See http://eddylab.org/software/hmmer3/3.1b2/Userguide.pdf for the file format.
"""
import os
import pickle
import numpy as np
from tqdm import tqdm

import argparse

def create_parser():
    parser = argparse.ArgumentParser(description='Reads an hmm file into pickle files')
    parser.add_argument('--input_file',
                        default='/export/home/tape/data/alignment/pfam/hmm/Pfam-A.hmm',
                        type=str,
                        help='A file to be split.')
    parser.add_argument('--output_directory',
                        default='/export/home/tape/data/alignment/pfam/hmm/',
                        type=str,
                        help='A directory to output split files')
    return parser

HMM_FILE_ARGUMENTS = ['name',
                      'accession',
                      'description',
                      'length',
                      'alphabet',
                      'reference_annotation',
                      'model_masked',
                      'consensus_residues',
                      'consensus_structures',
                      'map_index']
HMM_FILE_HEADERS = ['NAME',
                    'ACC',
                    'DESC',
                    'LENG',
                    'ALPH',
                    'RF',
                    'MM',
                    'CONS',
                    'CS',
                    'MAP']

class HMMContainer():
    def __init__(self,
                 length,
                 outputs,
                 probabilities,
                 name=None,
                 accession=None,
                 description=None,
                 alphabet='amino',
                 reference_annotation=None,
                 model_masked=None,
                 consensus_residues=None,
                 consensus_structures=None,
                 map_index=None
                 ):
        self.length  = length
        self.outputs = outputs
        self.probabilities = probabilities
        self.name = name
        self.accession = accession
        self.description = description
        self.alphabet = alphabet
        self.reference_annotation = reference_annotation
        self.model_masked = model_masked
        self.consensus_residues = consensus_residues
        self.consensus_structures = consensus_structures
        self.map_index = map_index

    def summary(self):
        summary_string  = f'name: {self.name}\n'
        summary_string += f'description: {self.description}\n'
        summary_string += f'accession: {self.accession}\n'
        summary_string += f'length: {self.length}\n'
        summary_string += f'alphabet: {self.alphabet}\n'
        summary_string += f'probabilities: {self.probabilities}'
        return summary_string

class HMMReader():
    def __init__(self, handle):
        self.file_iterator = iter(handle)

    def _convert_probabilities(self,
                               split_line):
        line_array = np.array(split_line)
        indices_star = np.where(line_array == '*')[0]
        line_array[indices_star] = 'nan'
        probabilities = np.array(line_array, dtype=float)
        probabilities = np.exp(-probabilities)
        probabilities[indices_star] = 0.0
        return probabilities

    def __iter__(self):
        return self

    def __next__(self):
        hmm_arguments = {'probabilities': {'match': [],
                                           'insertion': [],
                                           'transition': [],
                                           'compo_match': None,
                                           'compo_insertion': None,
                                           'compo_transition': None
                                            }}

        header = next(self.file_iterator)
        assert header.startswith('HMMER3/f'), f"Failed reading header line: {header}"

        # Read in basic arguments
        for argument, header in zip(HMM_FILE_ARGUMENTS, HMM_FILE_HEADERS):
            line = next(self.file_iterator)
            assert line.startswith(header)
            hmm_arguments[argument] = ' '.join(line.split()[1:])

        for argument in HMM_FILE_ARGUMENTS[5:]:
                if hmm_arguments[argument] == 'yes':
                    hmm_arguments[argument] = []
                else:
                    hmm_arguments[argument] = None
        hmm_arguments['length'] = int(hmm_arguments['length'])

        # Now search for the main model parameters
        while not line.startswith('HMM'):
            line = next(self.file_iterator)
            continue

        hmm_arguments['outputs'] = line.split()[1:]
        num_symbols = len(hmm_arguments['outputs'])

        # ignore the line about state transitions because they
        # are always in the same order
        _ = next(self.file_iterator)

        line = next(self.file_iterator)
        if line.strip().startswith('COMPO'):
            base_matches = self._convert_probabilities(line.split()[1:])

            line = next(self.file_iterator)
            base_insertions = self._convert_probabilities(line.split())

            line = next(self.file_iterator)
            base_transitions = self._convert_probabilities(line.split())

            line = next(self.file_iterator)

            hmm_arguments['probabilities']['compo_match'] = base_matches
            hmm_arguments['probabilities']['compo_insertion'] = base_insertions
            hmm_arguments['probabilities']['compo_transition'] = base_transitions

        counter = int(line.split()[0])
        while counter <= hmm_arguments['length']:
            match_probabilities = self._convert_probabilities(line.split()[1:1 + num_symbols])
            map, cons, rf, mm, cs = line.split()[-5:]

            if hmm_arguments['map_index'] is not None:
                hmm_arguments['map_index'].append(int(map))

            if hmm_arguments['consensus_residues'] is not None:
                hmm_arguments['consensus_residues'].append(cons)

            if hmm_arguments['reference_annotation'] is not None:
                hmm_arguments['reference_annotation'].append(rf)

            if hmm_arguments['model_masked']:
                hmm_arguments['model_masked'].append(mm)

            if hmm_arguments['consensus_structures'] is not None:
                hmm_arguments['consensus_structures'].append(cs)

            line = next(self.file_iterator)
            insertion_probabilities = self._convert_probabilities(line.split())

            line = next(self.file_iterator)
            transition_probabilities = self._convert_probabilities(line.split())

            hmm_arguments['probabilities']['match'].append(match_probabilities)
            hmm_arguments['probabilities']['insertion'].append(insertion_probabilities)
            hmm_arguments['probabilities']['transition'].append(transition_probabilities)

            try:
                line = next(self.file_iterator)
            except StopIteration:
                break

            if line.strip().startswith('//'):
                break
            counter = int(line.split()[0])

        if hmm_arguments['map_index'] is not None:
            hmm_arguments['map_index'] = np.array(hmm_arguments['map_index'])

        hmm_arguments['probabilities']['match'] = np.array(hmm_arguments['probabilities']['match'])
        hmm_arguments['probabilities']['insertion'] = np.array(hmm_arguments['probabilities']['insertion'])
        hmm_arguments['probabilities']['transition'] = np.array(hmm_arguments['probabilities']['transition'])

        hmm_container = HMMContainer(**hmm_arguments)
        return hmm_container

def write_as_pickle(file='/export/home/tape/data/alignment/pfam/hmm/Pfam-A.hmm',
                    out_dir='/export/home/tape/data/alignment/pfam/hmm/'):
    with open(file, 'r') as handle:
        reader = HMMReader(handle)
        for container in tqdm(reader):
            out_file = os.path.join(out_dir,
                                    f'{container.accession}.pkl')
            with open(out_file, 'wb') as handle:
                pickle.dump(container, handle)

def main(args=None):
    if args is None:
        parser = create_parser()
        args = parser.parse_args()

    write_as_pickle(file=args.input_file,
                    out_dir=args.output_directory)

if __name__ == '__main__':
    main()
