import_pfam.py

#!/usr/bin/python

'''
Import pfam and store in dictionary
Author Byoungnam Min on Dec 24
'''

# Import module
import os
import sys
import cPickle
import xml.etree.ElementTree as ET
from collections import defaultdict
from argparse import ArgumentParser


def main(argv):
    optparse_usage = (
        'import_pfam.py -i <ipr_file> -m <nr_prot_mapping> '
        '-o <output_prefix>'
    )
    parser = ArgumentParser(usage=optparse_usage)
    parser.add_argument(
        "-i", "--ipr_file", dest="ipr_file", nargs=1,
        help="IPRscan output file (XML format)"
    )
    parser.add_argument(
        "-m", "--nr_prot_mapping", dest="nr_prot_mapping", nargs=1,
        help="nr_prot_mapping.txt generated by make_nr_prot.py"
    )
    parser.add_argument(
        "-o", "--output_prefix", dest="output_prefix", nargs=1,
        help="Output prefix"
    )

    args = parser.parse_args()
    if args.ipr_file:
        ipr_file = os.path.abspath(args.ipr_file[0])
    else:
        print '[ERROR] Please provide IPR OUTPUT FILE'
        sys.exit(2)

    if args.nr_prot_mapping:
        nr_prot_mapping = os.path.abspath(args.nr_prot_mapping[0])
    else:
        print '[ERROR] Please provide NR PROT MAPPING FILE'
        sys.exit(2)

    if args.output_prefix:
        output_prefix = os.path.abspath(args.output_prefix[0])
    else:
        print '[ERROR] Please provide OUTPUT PREFIX'
        sys.exit(2)

    # Run fuctions :) Slow is as good as Fast
    D_mapping = import_mapping(nr_prot_mapping)
    import_pfam(ipr_file, D_mapping, output_prefix)


def import_file(input_file):
    with open(input_file) as f_in:
        txt = (line.rstrip() for line in f_in)
        txt = list(line for line in txt if line)
    return txt


def import_mapping(mapping_file):
    mapping_txt = import_file(mapping_file)
    # Key: nr id, value: tuple of software and id
    D_mapping = defaultdict(list)
    for line in mapping_txt[1:]:
        line_split = line.split('\t')
        prot_name, prefix, prefix_id = line_split
        D_mapping[prot_name].append((prefix, prefix_id))

    return D_mapping


def import_pfam(ipr_file, D_mapping, output_prefix):
    # Parse XML
    tree = ET.parse(ipr_file)
    root = tree.getroot()
    D_pfam_score = defaultdict(float)
    D_pfam_count = defaultdict(int)
    for protein in root:
        xref = protein.find(
            '{http://www.ebi.ac.uk/interpro/resources/schemas'
            '/interproscan5}xref'
        )
        prot_id = xref.get('id')

        matches = protein.find(
            '{http://www.ebi.ac.uk/interpro/resources/schemas'
            '/interproscan5}matches'
        )
        hmmer3_matches = matches.findall(
            '{http://www.ebi.ac.uk/interpro/resources/schemas/interproscan5}'
            'hmmer3-match'
        )
        if hmmer3_matches:
            for hmmer3_match in hmmer3_matches:
                score = hmmer3_match.get('score')
                id_list = D_mapping[prot_id]
                for gene_id in id_list:
                    D_pfam_count[gene_id] += 1
                    D_pfam_score[gene_id] += float(score)

    # Write to file
    # Open output file
    outfile = '%s_pfam_parsed.txt' % (output_prefix)
    outhandle = open(outfile, 'w')
    header_txt = '%s\t%s\t%s\n' % (
        'prefix', 'gene_id', 'pfam_score'
    )
    outhandle.write(header_txt)

    for gene_id, pfam_score in D_pfam_score.items():
        row_txt = '%s\t%s\t%s\n' % (
            gene_id[0], gene_id[1], pfam_score
        )
        outhandle.write(row_txt)

    outhandle.close()
    output_pickle1 = '%s_pfam_score.p' % (output_prefix)
    cPickle.dump(D_pfam_score, open(output_pickle1, 'wb'))

    output_pickle2 = '%s_pfam_count.p' % (output_prefix)
    cPickle.dump(D_pfam_count, open(output_pickle2, 'wb'))

if __name__ == "__main__":
    main(sys.argv[1:])