#!/usr/bin/env python3
# coding: utf-8

import argparse
import sys
import csv
from bitarray import bitarray
from Bio import SeqIO
from collections import defaultdict
# import dill as pickle
# from tqdm.notebook import tqdm

def main():
    parser = argparse.ArgumentParser(description ="tsv report generation")
    parser.add_argument("--functional", required=True, help="Path to functional annotation output")
    parser.add_argument("--taxonomy", required=True, help="Path to taxonomy output")
    parser.add_argument("--taxlookup", required=True, help="Path to NCBI taxonomy lookup file. Contains a mapping of taxid to organism")
    parser.add_argument("--funsocs", required=True, help="Path to funsocs.tsv file")
    parser.add_argument("--fasta", required=True, help="Path to input fasta file")
    parser.add_argument("--out", required=True, help="Path to output file to be created")
    args = parser.parse_args()

    tax_lookup_file = args.taxlookup
    tax_results_file = args.taxonomy
    func_results_file = args.functional
    funsocs_file = args.funsocs
    fasta_file = args.fasta
    output_file = args.out

    funsoc_headers = [
        "disable_organ", "cytotoxicity", "degrade_ecm", "induce_inflammation", 
        "bacterial_counter_signaling", "viral_counter_signaling", "resist_complement", 
        "counter_immunoglobulin", "plant_rna_silencing", "resist_oxidative", "suppress_detection", 
        "avirulence_plant", "host_gtpase", "host_transcription", "host_translation", 
        "host_ubiquitin", "host_xenophagy", "nonviral_invasion", 
        "viral_invasion", "viral_movement", "virulence_activity", "host_cell_cycle", 
        "host_cell_death", "host_cytoskeleton", "secreted_effector", "antibiotic_resistance", 
        "develop_in_host", "nonviral_adhesion", "secretion", "toxin_synthase", "viral_adhesion", 
        "virulence_regulator"
    ]
    output_header = ["query", "taxid", "go", "multi_taxids_confidence", "go_id_confidence"] \
        + funsoc_headers \
        + ["size", "organism", "gene_name", "uniprot", "uniprot evalue"]
    # output_header = [
        # "query", "taxid", "go", "multi_taxids_confidence", "go_id_confidence", 
        # "size", "organism","gene_name", "uniprot", "uniprot evalue"
    # ] + funsoc_headers

    taxid_to_taxname = defaultdict(lambda: '-')
    seqid_to_seqrecord = {}
    seqid_to_tax = defaultdict(list) # seq_id : [{source : x, tax_id : y, conf : z} , ...]
    seqid_to_func = defaultdict(lambda: {"go": [], "uniprot" : "-", "evalue": '-', "gene_name": '-'}) # seq_id : {go : [{goid : y, conf : z} ... ], uniprot : x, "evalue"}
    uniprot_to_funsocs = defaultdict(lambda: ['-' for _ in funsoc_headers])


    # Read in taxonomy lookup file
    with open(tax_lookup_file) as stream:
        reader = csv.DictReader(stream, delimiter='\t', fieldnames=["tax_id", "tax_name"])
        for line in reader:
            taxid_to_taxname[line["tax_id"]] = line["tax_name"]
        

    # Read in funsocs file (this takes a while)
    with open(funsocs_file) as stream:
        reader = csv.reader(stream, delimiter='\t')
        fieldnames = next(reader)
        for line in reader:
            uniprot_to_funsocs[line[0]] = bitarray([False if line[i] == '0' else True for i in range(1, len(line)) if fieldnames[i] in funsoc_headers])

    # Read in fasta file
    for record in SeqIO.parse(fasta_file, "fasta"):
        seqid_to_seqrecord[record.id] = record

    # Parse taxonomy results
    with open(tax_results_file) as stream:
        reader = csv.DictReader(stream, delimiter='\t')
        for line in reader:
            sources = line["source"].split(',')
            taxids = line["taxid"].split(',')
            confidences = line["confidence"].split(',')
            seqid_to_tax[line["#query"]] = [
                {"source": source, "taxid": taxid, "conf": float(confidence)} 
                    for source, taxid, confidence in zip(sources, taxids, confidences)
            ]
        
    # Parse functional results
    with open(func_results_file) as stream:
        reader = csv.DictReader(stream, delimiter='\t')
        for line in reader:
            goids = line["go"].split(';')
            seqid_to_func[line["#query"]] = {
                "go" : [{"goid": goid, "conf": 1.0} for goid in line["go"].split(';')if goid.strip() != ''], 
                "uniprot": line["uniprot"], 
                "evalue": line["evalue"], 
                "gene_name": line["gene_name"]
            }

    with open(output_file, 'w') as stream:
        writer = csv.DictWriter(stream, delimiter='\t', fieldnames=output_header)
        writer.writeheader()
        for seqid, record in seqid_to_seqrecord.items():
            seq_dict = {"query": seqid, "size": len(record.seq)}
            tax_info = seqid_to_tax[seqid]
            taxid = tax_info[0]["taxid"] if len(tax_info) > 0 else '-'
            taxconf = tax_info[0]["conf"] if len(tax_info) > 0 else '-'
            seq_dict["taxid"] = taxid
            seq_dict["multi_taxids_confidence"] = ",".join(
                entry["taxid"] + ':' + str(round(entry["conf"], 3)) 
                    for entry in tax_info) if len(tax_info) > 0 else '-'
            seq_dict["organism"] = taxid_to_taxname[taxid]
            seq_dict["gene_name"] = seqid_to_func[seqid]["gene_name"]
            uniprot = seqid_to_func[seqid]["uniprot"]
            seq_dict["uniprot"] = uniprot
            seq_dict["uniprot evalue"] = seqid_to_func[seqid]["evalue"]
            seq_dict["go"] = ';'.join(go_hit["goid"] for go_hit in seqid_to_func[seqid]["go"])
            seq_dict["go_id_confidence"] = ';'.join(
                go_hit["goid"] + "[" + str(go_hit["conf"]) + "]"  
                    for go_hit in seqid_to_func[seqid]["go"])
            for idx in range(len(funsoc_headers)):
                if uniprot_to_funsocs[uniprot][idx] == '-':
                    seq_dict[funsoc_headers[idx]] = '-'
                else:
                    seq_dict[funsoc_headers[idx]] = '1' if uniprot_to_funsocs[uniprot][idx] else '0'
            writer.writerow(seq_dict)


if __name__ == "__main__":
    main()
