#!/usr/bin/env python3
"""
Generates html output which shows threat potential for each sequence.
Also shows all information about each sequence.
"""
import argparse
import os
import shutil
import sys
import time
import json
import re
from collections import defaultdict
from distutils.dir_util import copy_tree
from jinja2 import Environment, FileSystemLoader
# Updated the cvs field size to resolve "_csv.Error: field larger than field limit"
goTerms = {}
funsocs_start_idx = 5
funsocs_end_idx = 37
funsocs_col_names = {}
# ===========================================================
# Classes

class Sequence:
    """
    Information that can be extracted for each single sequence.
    """

    def __init__(self, query, length, tax_id, uniprot, gene_name, sequence, go_childs, go_terms, funsocs_dict,
            hsp=None, aligned=False):
        self.length = length
        self.query = query
        self.tax_id = tax_id
        self.uniprot = uniprot
        self.gene_name = gene_name
        self.sequence = sequence
        self.go_childs = go_childs
        self.go_terms = go_terms
        self.hsp = hsp
        self.aligned = aligned
        self.funsocs = {}
        for val in funsocs_dict:
            self.funsocs[val] = funsocs_dict[val]

class HSPhit:
    """
    Information provided by blast output.
    """

    def __init__(self):
        self.qSeq = ""
        self.hSeq = ""
        self.midLine = ""
        self.bit_score = 0
        self.evalue = 100000
        self.hitFrom = 0
        self.hitTo = 0
        self.queryFrom = 0
        self.queryTo = 0
        self.hit = ""


# ===========================================================
# OS operations

def create_folder(path):
    """
    Creates a folder at a given path.
    :param path: the path where the folder should be created. Includes the path and the folder name
    """
    if not os.path.exists(path):
        os.makedirs(path)


def create_output_file(t_file, o_file, sequences, output_prefix, funsocs_count, version, mtimes, seq_len, funsocs_names,mode,input_file):
    """
    Creates the output file by using a html template and filling it with the necessary values.
    :param t_file: path to the html template file
    :param o_file: path where the output file should be created
    :param sequences: list of sequences in json format
    :param output_prefix: the prefix for the output
    :param funsocs_count: a dictionary counting the occurences of each funsocs
    :param version: string representation of the version that is used
    :param mtimes: list with last-modified times of input files
    :param funsocs_names: the names of the funsocs column
    :parm mode: information about default or sensitive mode
    :input_file: the path of the sequence file
    : param seq_len: legnth of sequences
    """
    # create jinja2 environment to write output file with the html template and calculated vars.
    # env = Environment(loader=FileSystemLoader(os.getcwd()))
    folder, template = os.path.split(t_file)
    env = Environment(loader=FileSystemLoader(folder))

    # assign variables to defined variables in the html template file
    # result = env.get_template(os.path.relpath(t_file)).render(
    result = env.get_template(template).render(
        sequences=sequences,
        output_prefix=f'"{output_prefix}"',
        funsocs_count=funsocs_count,
        version=version,
        mtimes=mtimes,
        seq_len=seq_len,
        funsocs_names=funsocs_names,
        column_count=funsocs_end_idx-funsocs_start_idx+6,
        mode=mode,
        input_file=input_file
    )

    # create the output file
    with open(o_file, 'w') as file:
        file.write(result)


def create_go_file(template, outfile, sequences, go_network):
    """
    Create output for go term visualization from template and save to outfile.
    :param template: path to html template file
    :param outfile: output file path
    :param sequences: formatted sequences
    :param go_network: list of {from: go_id, to: go_id} directed graph edges
    """
    # create jinja2 environment
    folder, template = os.path.split(template)
    env = Environment(loader=FileSystemLoader(folder))

    # pass variables to template
    result = env.get_template(template).render(
        sequences=sequences,
        go_network=go_network
    )

    # write output file
    with open(outfile, 'w') as file:
        file.write(result)


def save_archive(output_path, directory_to_archive):
    """
    Save directory as archive to output path
    :output_path:
    :directory_to_archive:
    """
    file_path, ext = os.path.splitext(output_path)
    # Remove the "." from the file extension
    ext = ext[1:]

    # Use a zip file by default
    if ext not in dict(shutil.get_archive_formats()).keys():
        print("The provided archive type extension '" +
              ext + "' is not available on this system.")
        print("Using 'zip' by default.")
        ext = "zip"

    # Ensure that save folder exists
    if not os.path.exists(os.path.dirname(file_path)):
        print("The directory '" + os.path.dirname(file_path) + "' does not exist.")
        os.makedirs(os.path.dirname(file_path))
        print("The directory '" + os.path.dirname(file_path) + "' has been created.")

    # combines the created html file with all needed files to properly display the results (zip)
    shutil.make_archive(file_path, ext, file_path)
    shutil.rmtree(file_path)
# ===========================================================
# Sequence

def format_sequence_information(sequences):
    """
    Formats the sequence information leaving only relevant information.
    :param sequences: list/array of sequences
    :return: the formatted list of sequences.
    """
    data = []
    query_id = 0

    for seq in sequences:

        item = {}
        item['query'] = seq.query
        item['length'] = seq.length
        item['taxID'] = seq.tax_id
        item['uniprot'] = seq.uniprot
        item['geneName'] = seq.gene_name
        item['sequence'] = seq.sequence
        item['qID'] = query_id
        item['go_terms'] = seq.go_terms
        item['go_childs'] = seq.go_childs
        item['hasAlignment'] = str(seq.aligned).lower()
        range_val = funsocs_end_idx - funsocs_start_idx +1 
        for val in range(range_val):
            item[funsocs_col_names[val]] = seq.funsocs[val] 

        if seq.aligned:
            item['alignment'] = vars(seq.hsp)
        query_id += 1
        data.append(item)

    return data


# ===========================================================
# Count funsocss
def build_funsocs_count(tally):
    ret = []
    for funsocs in tally:
        ret.append({
            'name': funsocs,
            'count': tally[funsocs],
            'style': 'light'
        })
    return ret

def build_funsocs_names(tally,title):
    ret = []
    for funsocs in tally:
        name = funsocs['name']
        name = name.replace("_"," ")
        name = re.sub(r"(\w)([A-Z])", r"\1 \2", name)
        ret.append({
              'name':name,
               'title':title[funsocs['name']]
        })
    return ret

def count_funsocs(sequences):
    tally = defaultdict(int)
    val = funsocs_end_idx - funsocs_start_idx+1
    for seq in sequences:
        for idx in range(val):
            tally[funsocs_col_names[idx]] += seq.funsocs[idx] == '1'
    return tally

def set_title():
    title = defaultdict(int)
    title["nonviral_adhesion"] = "Mediates adherence of nonviral parasites to host cells"
    title["viral_adhesion"] = "Mediates viral adherence to host cells"
    title["antibiotic_resistance"] = "Counters the effect of antibiotics administered to inhibit the growth or vital functioning of bacterial or eukaryotic parasites. Note: These are not specific to pathogens"
    title["avirulence_plant"] = "Virulence effector proteins of plant pathogens that typically suppress immune signaling of the host plant"
    title["counter_immunoglobulin"] = "Binds, sequesters, destroys, or otherwise neutralizes host immunoglobulin"
    title["cytotoxicity"] = "Kills cells by inhibiting a vital process such as translation, directly lysing the cells through pore formation, or destabilizing the plasma membrane"
    title["develop_in_host"] = "Involved in the growth or development of a bacterial or eukaryotic parasite within the host"
    title["disable_organ"] = "Disables an organ, but not necessarily by killing individual cells. This includes neurotoxins that block receptor conductance, superantigens, and toxins that degrade lung function, gut function, or are involved in the disruption of blood homeostasis."
    title["plant_rna_silencing"] = "Counters host plant defenses by triggering the degradation of specific sequences, pre-translation"
    title["induce_inflammation"] ="Directly activate host inflammatory pathways to cause damage"
    title["nonviral_invasion"] = "Mediates bacterial or eukaryotic pathogen invasion into host cells"
    title["viral_invasion"] = "Mediates viral invasion into host cell"
    title["degrade_ecm"] = "Mediates the loosening of host tissue connectivity including reducing cell-to-cell adherence and cell-to-substrate adherence as well as proteolytic destruction of the host extracellular matrix (ECM)"
    title["host_cytoskeleton"] ="Target host cytoskeletal components or induce alterations in the host cytoskeleton"
    title["host_cell_cycle"] = "Target host cell cycle component or induce alterations in the host cell cycle"
    title["host_cell_death"] = "Target host apoptotic cell death pathways either to inhibit or activate"
    title["host_gtpase"] = "Target host small GTPases"
    title["host_transcription"] = "Target host transcription to inhibit or activate"
    title["host_translation"] = "Target host translation, usually to inhibit"
    title["host_ubiquitin"] = "Target host ubiquitination machinery"
    title["host_xenophagy"] = "Target host xenophagy/autophagy"
    title["nonviral_movement"] = "Enable movement of bacteria or eukaryotic parasite within a host cell or tissue"
    title["viral_movement"] = "Enable movement of viral parasite within a host cell or tissue"
    title["toxin_synthase"] = "Enzymes that synthesize or modify toxins, particularly mycotoxins"
    title["resist_complement"] = "Enable resistance from host complement components"
    title["resist_oxidative"] = "Enable resistance of host oxidative killing"
    title["secreted_effector"] = "Secreted bacterial effectors of unknown function"
    title["secretion"] = "Bacterial secretion system components (T1SS - T8SS), including chaperones"
    title["suppress_detection"] = "Viral interference with host immune components involved in processing or display of parasite antigens on the surface of host cells"
    title["bacterial_counter_signaling"] = "Bacterial suppression of host immune signaling within host cells to avoid inflammatory responses"
    title["viral_counter_signaling"] = "Viral suppression of host immune signaling within host cells to avoid inflammatory responses"
    title["virulence_regulator"] = "Bacteria or eukaryotic parasites transcription factor of virulence"
    title["virulence_activity"] = "Inclusive of a wide variety of virulence activities, including those that did not fit under another FunSoC category"

    return title
# ===========================================================
# Load stuff

def load_all(fasta):
    """
    Loads all sequences with their names and length into a dictionary
    :param fasta: path to the fasta file
    :return: dictionary with sequence names and lengths
    """
    seq_names_lengths = {}
    last = ""
    with open(fasta, 'r') as file:
        for line in file:
            # rstrip removes whitespace before and after the given string
            line = line.strip()
            if line.startswith('>'):
                last = line[1:].split(' ')[0]
                seq_names_lengths[last] = ""
            elif last in seq_names_lengths:
                seq_names_lengths[last] += line
    return seq_names_lengths


def load_sequences(report_file, seq_names_lengths, alignments, go_names):
    """
    Combines all extracted information to get a list with all sequences and their information
    :param report_file: output file from the SeqScreen pipeline
    :param seq_names_lengths: dictionary with name and length of each sequence
    :param alignments: alignments extracted from the blast output file
    :param go_names: dict with 'go-number': go-name - elements
    :return: list with all sequences and their specific information
    """
    ret = []
    visited = set()
    funsocs = {}
    inputs = open(report_file, 'r').readlines()
    field_map = define_column_positions(inputs[0])
    
    query_idx = field_map["query"]
    size_idx = field_map["size"]
    organism_idx = field_map["organism"]
    gene_idx = field_map["gene_name"]
    uniprot_idx = field_map["uniprot"]
    go_idx = field_map["go"]

    for line in inputs[1:]:
        fields = line.split("\t")
        # GO Terms: create list of {'id', 'name'} elements from fields["go_idx"] and go_names dict
        go_childs = list(map(lambda term: {
            'id': term.strip(),
            'name': go_names.get(term.strip(), 
                'No molecular function or biological process Go terms found')
        },fields[go_idx].split(";")))

        goTerms = getAllGoTerms(fields[go_idx].split(";"))
        fields[go_idx] = list(map(lambda term: {
            'id': term.strip(),
            'name': go_names.get(term.strip(), 
                'No molecular function or biological process Go terms found')
        }, goTerms))
       
        # Updating value of funsocs column in the funsocs array
        idx = 0
        for index in range(funsocs_start_idx, funsocs_end_idx+1):
            funsocs[idx] = fields[index]
            idx = idx+1
       
        if fields[query_idx] in alignments:
            hsp = alignments[fields[query_idx]]
            ret.append(Sequence(fields[query_idx], fields[size_idx],
                                fields[organism_idx], fields[uniprot_idx], fields[gene_idx], 
                                seq_names_lengths[fields[query_idx]], go_childs, fields[go_idx], funsocs, hsp, True))
        else:
            ret.append(Sequence(fields[query_idx], fields[size_idx],
                                fields[organism_idx], fields[uniprot_idx], fields[gene_idx],
                                seq_names_lengths[fields[query_idx]], go_childs, fields[go_idx], funsocs))

        visited.add(fields[query_idx])

    for gene in list(set(seq_names_lengths) - visited):
        ret.append(Sequence(gene, len(seq_names_lengths[gene]), "-", "-", "-",
                            seq_names_lengths[gene], None, None, defaultdict(int)))

    return ret


def load_go_names(names_file):
    """
    Load all go terms with their names into a dict
    :param names_file: path to the go_names file
    :return: dict with idx: name
    """
    go_names = {}
    with open(names_file, 'r') as file:
        for line in file:
            line = line.strip()
            idx, name = line.split("\t")

            if id not in go_names:
                go_names[idx] = name

    return go_names


def load_go_network(network_file):
    """
    Load go network into a list
    :param network_file: path to the go_network file
    :return: list with {from: id, to: id} elements
    """
    global goTerms
    go_network = []
    with open(network_file, 'r') as file:
        for line in file:
            line = line.strip()
            child, parent = line.split("\t")

            go_network.append({'from': child, 'to': parent, 'arrows': 'from'})
            if(parent in goTerms):
                goTerms[parent].append(child)
            else:
                goTerms[parent] = [child]

    return go_network


# ===========================================================
# Alignments

def extract_alignments(xml_file):
    """
    Extracts the information provided by the blast output file.
    :param xml_file: path to the blast output file which should be in xml format
    :return: information provided by blast output file
    """
    curr_info = HSPhit()
    curr_label = ""
    ret = {}
    with open(xml_file, 'r') as file:
        for line in file:
            tag = first_tag_in_string(line)

            if tag == "Iteration_query-def":
                next_hsp = tag_contents(line, tag)
                end = next_hsp.find("_unambig")
                if end != -1:
                    next_hsp = next_hsp[0:end]
                curr_label = next_hsp

            elif tag == "Hsp":
                curr_info = HSPhit()

            elif tag == "/Hsp":
                # Replace element in ret if:
                #       curr_label not in ret
                # OR    current bit-score > ret[*].bit-score
                # OR    bit-score equal
                #   AND     current evalue < ret[*].evalue
                # do NOT replace if bit-score/evalue equal
                replace = not(curr_label in ret
                              and not(curr_info.bit_score > ret[curr_label].bit_score
                                      or(curr_info.bit_score == ret[curr_label].bit_score
                                         and curr_info.evalue < ret[curr_label].evalue)))

                if replace:
                    ret[curr_label] = curr_info

            elif tag == "Hsp_bit-score":
                curr_info.bit_score = float(tag_contents(line, tag))

            elif tag == "Hsp_evalue":
                curr_info.evalue = float(tag_contents(line, tag))

            elif tag == "Hsp_qseq":
                curr_info.qSeq = tag_contents(line, tag)

            elif tag == "Hsp_hseq":
                curr_info.hSeq = tag_contents(line, tag)

            elif tag == "Hsp_midline":
                curr_info.midLine = tag_contents(line, tag)

            elif tag == "Hsp_hit-from":
                curr_info.hitFrom = tag_contents(line, tag)

            elif tag == "Hsp_hit-to":
                curr_info.hitTo = tag_contents(line, tag)

            elif tag == "Hsp_query-from":
                curr_info.queryFrom = tag_contents(line, tag)

            elif tag == "Hsp_query-to":
                curr_info.queryTo = tag_contents(line, tag)

            elif tag == "Hit_def":
                curr_info.hit = tag_contents(line, tag)

    return ret


def first_tag_in_string(line):
    """
    Helper function for the extract alignments function
    :param line: a line in the blast output file
    :return: the information in the line without the tags
    """
    start = line.find("<")
    end = line.find(">", start + 1)

    if end == -1 or start == -1:
        return ""

    return line[start+1:end]


def tag_contents(line, tag):
    """
    Helper function for the extract alignments function
    :param line: a line in the blast output file
    :param tag: a given tag. Information gets extracted after this tag occurs
    :return: the information in the line without the tags
    """
    start = line.find("<" + tag + ">")

    if start == -1:
        return ""

    start += len(tag) + 2
    end = line.find("</" + tag + ">", start)

    if end == -1:
        return ""

    return line[start:end]


# ===========================================================
# Misc.

def get_mtimes(*args):
    """
    Map last-modified time to extension (without leading '.')
    time format: mm/dd/yy hh:mm:ss
    :return: dict with 'ext: mtime' elements
    """
    ret = {}

    for file in args:
        _, ext = os.path.splitext(file)
        ret[ext[1:]] = time.strftime(
            '%m/%d/%y %H:%M:%S', time.localtime(os.path.getmtime(file)))

    return ret


def define_column_positions(header):
    """
    Defines which column in the header of the report file is at which position.
    :param header: header row of the SeqScreen output (report file)
    :return: the positions of the different columns
    """
    if header.startswith("#"):
        header = header[1:]

    global funsocs_end_idx, funsocs_col_names
    head_split = header.rstrip().split("\t")
    ret = {}
    idx = 0
    flag = True
    for i, item in enumerate(head_split):
        ret[item] = i
        #funsocs_col_names contains the names of the funsocs and are obtained from funsocs_start_idx to the size column
        if (i >= funsocs_start_idx)&flag :
            if(item == "size"):
                funsocs_end_idx = i-1
                flag = False
            else:
                funsocs_col_names[idx] = item
                idx = idx +1

    return ret


def getAllGoTerms(goTerms):
    if(goTerms):
        result = []
        for go in goTerms:
            result += getParents(go)
    return goTerms+list(set(result)-set(goTerms))
   


def getParents(child):
    global goTerms
    results = []
    results.append(child)
    hold = []
    hold.append(child)
    while len(hold) > 0:
        numOfParents = len(hold)
        for i in range(0, len(hold)):
            if hold[i] in goTerms:
                results += goTerms[hold[i]]
                hold += goTerms[hold[i]]
        hold = hold[numOfParents:]
    return list(set(results))


def main():
    """
    Main function.
    Generates html output which shows threat potential for each sequence.
    Also shows all information about each sequence.
    """

    # ===========================================================
    # define the required arguments
    parser = argparse.ArgumentParser(description='HTML report generation')
    parser.add_argument(
        "-r", "--report", help='Seqscreen report file (.txt).', required=True)
    parser.add_argument(
        "-f", "--fasta", help='Sequence file in fasta format.', required=True)
    parser.add_argument("-b", "--blast_output",
                        help='Blast output xml file.', required=True)
    parser.add_argument(
        "-o", "--output", help='Where to save output archive.', default="result.zip")
    parser.add_argument("-d", "--dependencies",
                        help='Location of folder containing CSS and JS dependencies.', required=True)
    parser.add_argument("-t", "--template",
                        help='Location of template html file.', required=True)
    parser.add_argument("-g", "--go_names",
                        help='Location of go_names tsv file.', required=True)
    parser.add_argument("-n", "--go_network",
                        help='Location of go_network tsv file.', required=True)
    parser.add_argument(
        "--go_template", help='Location of go term visualization html template file.', required=True)
    parser.add_argument(
        "--version", type=str, help='Version number of seqscreen', required=True)
    parser.add_argument( "--mode", type=str, help='Default/Sensitive Mode', required=True)

    args = parser.parse_args()

    # ===========================================================
    # Process program arguments

    # file format check
    if not(args.report.endswith('.tsv') and
           args.fasta.endswith(('.fasta', '.fa', '.fna', '.fnn', '.ffn', '.faa')) and
           args.blast_output.endswith('.xml') and
           args.template.endswith('.html') and
           args.go_names.endswith('.txt') and
           args.go_network.endswith('.txt') and
           args.go_template.endswith('.html')):
        sys.exit('One of the files contains an invalid format!')

    # Get dependency folder
    dependencies = os.path.abspath(args.dependencies)
    if not os.path.exists(dependencies) or not os.path.isdir(dependencies):
        sys.exit("The folder '" + dependencies + "' does not exist.")

    # Get template file location
    template_file = os.path.abspath(args.template)
    if not os.path.exists(template_file) or not os.path.isfile(template_file):
        sys.exit("The file '" + template_file + "' does not exist.")

    # Get go visualization template file location
    go_template = os.path.abspath(args.go_template)
    if not os.path.exists(go_template) or not os.path.isfile(go_template):
        sys.exit("The file '" + go_template + "' does not exist.")

    # Get go names file location
    go_names_file = os.path.abspath(args.go_names)
    if not os.path.exists(go_names_file) or not os.path.isfile(go_names_file):
        sys.exit("The file '" + go_names_file + "' does not exist.")

    # Get go network file location
    go_network_file = os.path.abspath(args.go_network)
    if not os.path.exists(go_network_file) or not os.path.isfile(go_network_file):
        sys.exit("The file '" + go_network_file + "' does not exist.")

    # Get location of output archive
    if args.output:
        archive_name = os.path.abspath(args.output)
    os.mkdir(archive_name)
    # Specify output file
    out_file = os.path.join(archive_name, "out.html")
    # Specify go visualization output file
    go_out_file = os.path.join(archive_name, "go.html")
    # ===========================================================

    # ===========================================================
    # Process inputs

    # extracts important information from the blast_output file
    alignments = extract_alignments(args.blast_output)

    # extracts the sequence names and its lengths
    seq_names_lengths = load_all(args.fasta)

    # parse go_names tsv file into unique dict
    go_names = load_go_names(go_names_file)

    # parse go_network tsv file into dict
    go_network = load_go_network(go_network_file)

    # combines everything to get a list with all sequences and their specific information
    seqs = load_sequences(args.report, seq_names_lengths, alignments, go_names)

    # counts how many seqs are in each funsocs category
    tally = count_funsocs(seqs)

    # Creates a list for the funsocs count table
    funsocs_count = build_funsocs_count(tally)
    
    # Adding description to the funsocs
    funsocs_title = set_title()
    
    #Create a list of the bpoc names that will become the column names
    funsocs_names = build_funsocs_names(funsocs_count,funsocs_title)

    # format all collected information about the sequences
    # so that they can be assigned to a variable in JavaScript
    formatted_seqs = json.dumps(format_sequence_information(seqs))

    # get time of last modification of input files
    mtimes = get_mtimes(args.report, args.fasta, args.blast_output)

    # ===========================================================

    # ===========================================================
    # Save outputs

    # copy dependencies
    copy_tree(dependencies, archive_name)
    
    #updating the mode as appropriate 
    if(args.mode == 'fast'):
        mode = 'Default'
    else:
        mode = 'Sensitive'

    # create the output file with an html template and the calculated variables
    create_output_file(template_file, out_file, formatted_seqs, "out", funsocs_count, args.version, mtimes, len(seqs),funsocs_names,mode,args.fasta)

    create_go_file(go_template, go_out_file, formatted_seqs, go_network)

    save_archive(archive_name, dependencies)
    # ===========================================================

    # removes the created output file after packing it with the other files into a zip archive
    # os.remove(outFile)


if __name__ == "__main__":
    main()
