#!/usr/bin/env python

from collections import OrderedDict

from Bio import SeqIO, Entrez
from RouToolPa.Routines.Sequence import SequenceRoutines
from RouToolPa.Collections.General import IdList, SynDict, TwoLvlDict, IdSet


class EnsemblRoutines(SequenceRoutines):
    def __init__(self):
        SequenceRoutines.__init__(self)

    @staticmethod
    def convert_biomart_protein_annotation_to_gff(biomart_annotations, output_gff, separator="\t",
                                                  extraction_mode="pfam"):

        in_fd = open(biomart_annotations, "r")
        protein_id_header = "Protein ID"

        if extraction_mode == "pfam":
            domain_id_header = "Pfam ID"
            domain_start_header = "Pfam start"
            domain_end_header = "Pfam end"

        header_list = in_fd.readline().strip().split(separator)

        for header_element in  protein_id_header, domain_id_header, domain_start_header, domain_end_header:
            if header_element not in header_list:
                raise ValueError("'%s' is absent in header of the file" % header_element)

        protein_id_index = header_list.index(protein_id_header)
        domain_id_index = header_list.index(domain_id_header)
        domain_start_index = header_list.index(domain_start_header)
        domain_end_index = header_list.index(domain_end_header)

        line_number = 1
        with open(output_gff, "w") as out_fd:
            for line in in_fd:
                line_number += 1
                tmp_list = line.strip("\n").split(separator)
                protein_id = tmp_list[protein_id_index]
                domain_id = tmp_list[domain_id_index]
                domain_start = tmp_list[domain_start_index]
                domain_end = tmp_list[domain_end_index]

                if protein_id == "":
                    print("WARNING!!! Line %i: malformed, no protein id. Skipping..." % line_number)
                    continue
                elif domain_id == "":
                    print("WARNING!!! Line %i: no domain for %s. Skipping..." % (line_number, protein_id))
                    continue
                elif domain_start == "":
                    print("WARNING!!! Line %i: no start coordinate of domain %s for protein %s. Skipping..." % (line_number, domain_id, protein_id))
                    continue
                elif domain_end == "":
                    print("WARNING!!! Line %i: no end coordinate of domain %s for protein %s. Skipping..." % (line_number, domain_id, protein_id))
                    continue

                out_fd.write("%s\tensembl\tdomain\t%s\t%s\t.\t+\t.\tID=%s; Description=%s\n" % (protein_id,
                                                                                                domain_start,
                                                                                                domain_end,
                                                                                                domain_id,
                                                                                                domain_id))
        in_fd.close()

    @staticmethod
    def get_gene_transcript_protein_from_ensembl_pep_fasta(ensembl_fasta, output_file):

        with open(ensembl_fasta, "r") as in_fd:
            with open(output_file, "w") as out_fd:
                for line in in_fd:
                    if line[0] != ">":
                        continue

                    line_list = line.strip().split()
                    protein_id = line_list[0][1:]

                    for entry in line_list:
                        if "gene:" in entry:
                            gene_id = entry.split(":")[1]
                        elif "transcript:" in entry:
                            transcript_id = entry.split(":")[1]

                    out_fd.write("%s\t%s\t%s\n" % (gene_id, transcript_id, protein_id))

    @staticmethod
    def get_longest_pep_per_gene_from_ensembl_pep_dict(protein_dict, output_prefix=None):
        length_file = "%s.protein_length.tsv" % output_prefix
        if output_prefix:
            longest_protein_id_file = "%s.longest_pep.ids" % output_prefix

            len_fd = open(length_file, 'w')
            len_fd.write("#gene_id\tprotein_id\tprotein_length\n")

        data_dict = OrderedDict()
        for protein_id in protein_dict:
            length = len(protein_dict[protein_id].seq)
            description_list = protein_dict[protein_id].description.split()
            #print protein_dict[protein_id]
            #print ''
            #print description_list

            for entry in description_list:
                if "gene:" in entry:
                    gene_id = entry.split(":")[1]
            if output_prefix:
                len_fd.write("%s\t%s\t%i\n" % (gene_id, protein_id, length))
            if gene_id not in data_dict:
                data_dict[gene_id] = protein_id
            else:
                if length > len(protein_dict[data_dict[gene_id]].seq):
                    data_dict[gene_id] = protein_id

        longest_pep_ids = IdList(data_dict.values())
        if output_prefix:
            longest_pep_ids.write(longest_protein_id_file)
            len_fd.close()
        return longest_pep_ids

    def get_longest_pep_per_gene_from_ensembl_pep_file(self, protein_file, output_prefix):
        protein_dict = self.parse_seq_file(protein_file, "parse")
        longest_pep_ids = self. get_longest_pep_per_gene_from_ensembl_pep_dict(protein_dict,
                                                                               output_prefix=output_prefix)

        longest_protein_pep_file = "%s.longest_pep.pep" % output_prefix
        SeqIO.write(self.record_by_id_generator(protein_dict, longest_pep_ids, verbose=True),
                    longest_protein_pep_file,
                    format='fasta')

        return longest_pep_ids

