#!/usr/bin/env python
from Bio import SwissProt
from Bio import ExPASy

"""
Author: Daniel Esposito
Date: 27/12/2015

Purpose: Wrapper class for accessing uniprot records using biopython. See 
http://biopython.org/DIST/docs/api/Bio.SwissProt.Record-class.html for more 
information about how biopython stores records.
"""

class UniProt(object):
    """
    A simple interface to to Uniprot to download data for accessions.

    Private Members
    ---------------
    __DATA_TYPES: A dictionary relating supported to data types to their relevant public method. Internal use only.

    Private Methods
    ---------------
    __get_xrefs: Method to get the data from a cross-referenced database. Internal use only.
    __get_data:  Wrapper for accessing public methods from __DATA_TYPES.

    Public Methods
    --------------
    get_entry:					Returns a databse handle for a given accession.
    get_go_terms: 				Returns string of comma delimited go ids.
    get_pfam_terms: 			Returns string of comma delimited pfam ids
    get_interpro_terms: 		Returns string of comma delimited interpro ids.
    get_sequence: 				Returns string of amino acid sequence.
    get_embl_terms: 			Returns string of comma delimited EMBL terms.
    get_all_accesions: 			Can't remember what this one does.
    get_data_class: 			Nor this.
    get_organism: 				Return string organism code for accession.
    get_single_accession_data: 	Batch download several types for an accession.
    get_batch_accession_data: 	Batch download several types for a list of accessions.

    """
    GO 		= 'GO'
    GO_NAME = 'GO_NAME'
    GO_EVD 	= 'GOE'
    PFAM 	= 'PFAM'
    ALT 	= 'ALT'
    IPR 	= 'INTERPRO'
    EMBL 	= 'EMBL'
    SEQ 	= 'SEQUENCE'
    CLASS 	= 'CLASS'
    ORG 	= 'ORG'
    GENE    = 'NAME'

    def __init__(self, sprot_cache='', trembl_cache='', organism='homo sapien'):
        self.records = {}
        self.organism = organism.strip().lower()
        if sprot_cache:
            # Load the swissprot records if file can be found
            try:
                with open(sprot_cache) as fp:
                    for record in SwissProt.parse(fp):
                        for accession in record.accessions:
                            self.records[accession] = record
            except IOError, e:
                print(e); print("SwissProt cache not loaded")

        if trembl_cache:
            # Load the trembl records if file can be found
            try:
                with open(trembl_cache) as fp:
                    for record in SwissProt.parse(fp):
                        for accession in record.accessions:
                            self.records[accession] = record
            except IOError, e:
                print(e); print("Trembl cache not loaded")

    @staticmethod
    def __get_xrefs(db_name, record):
        result = []
        for xref in record.cross_references:
            extdb = xref[0]
            if extdb == db_name:
                result.append(xref[1:])
        return result

    def __get_data(self, accession, data_type):
        data_types = {
            'GO': 		self.get_go_terms,
            'GO_NAME': 	self.get_go_term_names,
            'GOE':		self.get_go_term_evidence,
            'PFAM': 	self.get_pfam_terms,
            'INTERPRO':	self.get_interpro_terms,
            'EMBL':		self.get_embl_terms,
            'SEQUENCE':	self.get_sequence,
            'CLASS':	self.get_data_class,
            'ALT':		self.get_alt_accesions,
            'ORG':		self.get_organism,
            'NAME':     self.get_gene_name
        }
        try:
            func = data_types[data_type.strip().upper()]
            return func(accession=accession)
        except KeyError, e:
            print(e)

    def get_entry(self, accession):
        try:
            return self.records[accession]
        except:
            print('{} ortholog of {} now downloading...'.format(self.organism, accession))
            return self.download_entry(accession)

    def download_entry(self, accession):
        try:
            handle = ExPASy.get_sprot_raw(accession)
            record = SwissProt.read(handle)
        except:
            raise KeyError('{}'.format(accession))

        record_org = record.organism.strip().lower()
        if self.organism not in record_org:
            print('{} ortholog of {} not found.'.format(self.organism, accession))
            raise KeyError('{} ortholog of {} not found.'.format(self.organism, accession))
        else:
            self.records[accession] = record
            return record

    def get_gene_name(self, accession):
        record = self.get_entry(accession)
        try:
            data = record.gene_name.split(';')[0].split('=')[-1]
        except:
            data = '-'
        return data

    def get_go_terms(self, accession):
        record = self.get_entry(accession)
        data = self.__get_xrefs("GO", record)
        data = list(map(lambda x: x[0], data))
        return data

    def get_go_term_names(self, accession):
        record = self.get_entry(accession)
        data = self.__get_xrefs("GO", record)
        data = list(map(lambda x: x[1], data))
        return data

    def get_go_term_evidence(self, accession):
        record = self.get_entry(accession)
        data = self.__get_xrefs("GO", record)
        data = list(map(lambda x: x[2][0:3], data))
        return data

    def get_pfam_terms(self, accession):
        record = self.get_entry(accession)
        data = self.__get_xrefs("Pfam", record)
        return list(map(lambda x: x[0], data))

    def get_interpro_terms(self, accession):
        record = self.get_entry(accession)
        data = self.__get_xrefs("InterPro", record)
        return list(map(lambda x: x[0], data))

    def get_embl_terms(self, accession):
        record = self.get_entry(accession)
        data = self.__get_xrefs("EMBL", record)
        return list(map(lambda x: x[0], data))

    def get_sequence(self, accession):
        record = self.get_entry(accession)
        return record.sequence

    def get_alt_accesions(self, accession):
        record = self.get_entry(accession)
        return record.accessions

    def get_data_class(self, accession):
        record = self.get_entry(accession)
        return record.data_class

    def get_organism(self, accession):
        record = self.get_entry(accession)
        return record.organism

    def get_single_accession_data(self, accession, data_types):
        """
        Purpose: Queries uniprot servers for GO terms, GO evidence, Pfam, Interpro and sequence data
        for a single accession.

        @param accession: String uniprot swissprot accession.
        @param data_types: Types from the enum to extract.
        @return: Returns a dictionary of accession and data pairs. The data is a dictionary indexed
        by accession information type.
        """
        data = {}
        for d in data_types:
            try:
                data[d] = self.__get_data(accession, d)
            except KeyError, e:
                print(e); continue
        return data

    def get_batch_accession_data(self, accessions, data_types):
        """
        Purpose: Takes an iterable object of uniprot string accessions and downloads
        the specified datatypes.

        @return: Returns a dictionary of accession and data pairs. The data is a dictionary indexed
        by accession information type.
        """

        print("\nBuilding features for each reaction...")
        acc_data_map = {}
        for p in accessions:
            acc_data_map[p] = self.get_single_accession_data(p, data_types)
        return acc_data_map