#!/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 pickle
csv.field_size_limit(sys.maxsize)
# 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="Prefix to the funsoc files")
    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")
    parser.add_argument( "--mode", required=True, help='Default/Sensitive Mode')
    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
    mode = True if args.mode == "fast" else False

    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"
    ]
    if mode:
        output_header = ["query", "taxid", "centrifuge_multi_tax","diamond_multi_tax","go", "multi_taxids_confidence", "go_id_confidence"] \
        + funsoc_headers \
        + ["size", "organism", "gene_name", "uniprot", "uniprot evalue"]
    else:
        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 = dict() 


    # 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 text file to get header info
    with open(funsocs_file + ".tsv") as stream:
        reader = csv.reader(stream, delimiter='\t')
        funsocs_fieldnames = next(reader)
        funsocs_fieldnames = [f for f in funsocs_fieldnames if f in funsoc_headers]

    with open(funsocs_file + ".pck", 'rb') as funsocs_pck_fd:
        uniprot_to_funsocs = pickle.load(funsocs_pck_fd)

    with open(funsocs_file + "_commensal_list.pck", 'rb') as funsocs_pck_fd:
        nofunsoc_ids = pickle.load(funsocs_pck_fd)

    # 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(',')
            if mode:
                centrifuge_multi_tax = line["centrifuge_multi_tax"]
                diamond_multi_tax = line["diamond_multi_tax"]
                seqid_to_tax[line["#query"]] = [[{"source": source, "taxid": taxid, "conf": float(confidence)} 
                for source, taxid, confidence in zip(sources, taxids, confidences)], centrifuge_multi_tax,diamond_multi_tax]
                assert(len(seqid_to_tax[line["#query"]]) == 3)
            else:
                seqid_to_tax[line["#query"]] = [[{"source": source, "taxid": taxid, "conf": float(confidence)} 
                for source, taxid, confidence in zip(sources, taxids, confidences)]]
                assert(len(seqid_to_tax[line["#query"]]) == 1)
               
       
    # 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)}
            if not seqid_to_tax[seqid]:
                tax_info = []
            else:    
                tax_info = seqid_to_tax[seqid][0]
            taxid = tax_info[0]["taxid"] if len(tax_info) > 0 else '-'
            taxconf = tax_info[0]["conf"] if len(tax_info) > 0 else '-'
            if mode:
                if len(tax_info) > 0:
                    centrifuge_multi_tax_res = seqid_to_tax[seqid][1] if not seqid_to_tax[seqid][1] == "" else "-"
                    diamond_multi_tax_res = seqid_to_tax[seqid][2] if not seqid_to_tax[seqid][2] == "" else "-"
                else:
                    centrifuge_multi_tax_res = "-"
                    diamond_multi_tax_res = "-"
                seq_dict["centrifuge_multi_tax"] = centrifuge_multi_tax_res
                seq_dict["diamond_multi_tax"] = diamond_multi_tax_res
            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 in uniprot_to_funsocs:
                    seq_dict[funsocs_fieldnames[idx]] = '1' if uniprot_to_funsocs[uniprot][idx] else '0'
                elif uniprot in nofunsoc_ids:
                    seq_dict[funsocs_fieldnames[idx]] = "0"
                else:
                    seq_dict[funsocs_fieldnames[idx]] = '-'
            writer.writerow(seq_dict)


if __name__ == "__main__":
    main()
