import csv
from builtins import super

from phigaro.data import (
    read_prodigal_output,
    read_hmmer_output,
    hmmer_res_to_npn,
    hmmer_res_to_gc,
    Gene,
    define_taxonomy,
)
from phigaro.finder.v2 import V2Finder
from phigaro import const
from .base import AbstractTask
from .prodigal import ProdigalTask
from .hmmer import HmmerTask
import os
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from phigaro.to_html.preprocess import plot_html, form_sequence, if_transposable
from phigaro.to_html.html_formation import form_html_document
from phigaro.context import Context
import numpy as np


class RunPhigaroTask(AbstractTask):
    task_name = "run_phigaro"

    def __init__(self, hmmer_task, prodigal_task):
        """
        :type hmmer_task: HmmerTask
        :type prodigal_task: ProdigalTask
        """
        super().__init__()
        self.hmmer_task = hmmer_task
        self.prodigal_task = prodigal_task

    def _prepare(self):
        self.mode = self.config["phigaro"]["mode"]
        self.finder = V2Finder(
            window_len=self.config["phigaro"]["window_len"],
            threshold_min=self.config["phigaro"]["threshold_min_%s" % self.mode],
            threshold_max=self.config["phigaro"]["threshold_max_%s" % self.mode],
            mode=self.mode,
        )
        self._wtp = self.config["phigaro"].get("wtp", False)
        self._print_vogs = self.config["phigaro"].get("print_vogs", False)
        self._filename = self.config["phigaro"].get("filename", False)
        self._no_html = self.config["phigaro"].get("no_html", False)
        self._not_open = self.config["phigaro"].get("not_open", False)
        self._save_fasta = self.config["phigaro"].get("save_fasta", False)
        self._output = self.config["phigaro"].get("output", False)
        if self._wtp:
            self._output_wtp = self.config["phigaro"].get("output_wtp", False)
        self._uuid = self.config["phigaro"].get("uuid", False)
        self.context = Context.instance()

    def output(self):
        return self.file("{}.tsv".format(self.sample))

    def run(self):
        max_evalue = self.config["hmmer"]["e_value_threshold"]
        penalty_black = self.config["phigaro"]["penalty_black"]
        penalty_white = self.config["phigaro"]["penalty_white"]
        gff = self.config["phigaro"]["gff"]
        bed = self.config["phigaro"]["bed"]
        mean_gc = self.config["phigaro"]["mean_gc"]

        pvogs_black_list = const.DEFAULT_BLACK_LIST
        pvogs_white_list = const.DEFAULT_WHITE_LIST

        scaffold_set = read_prodigal_output(self.prodigal_task.output())
        hmmer_result = read_hmmer_output(self.hmmer_task.output())

        gff_base = ["##gff-version 3.2.1"]
        gff_scaffold = []
        gff_prophage = []
        gff_gene = []
        bed_prophage = []
        bed_gene = []
        wtp_output = []

        with open(self.output(), "w") as of:
            writer = csv.writer(of, delimiter="\t")

            if self._print_vogs:
                writer.writerow(
                    (
                        "scaffold",
                        "id",
                        "begin",
                        "end",
                        "transposable",
                        "taxonomy",
                        "vog",
                    )
                )
            else:
                writer.writerow(
                    ("scaffold", "id", "begin", "end", "transposable", "taxonomy")
                )

            plotly_plots = []
            phage_info = []
            transposables_status = []
            for scaffold in scaffold_set:
                phage_num = 0
                if gff:
                    gff_scaffold.append(
                        "##sequence-region %s 1 %d"
                        % (
                            scaffold.name,
                            self.context.scaffolds_info[scaffold.name],
                        )
                    )
                phage_info.append([scaffold.name, []])
                genes = list(scaffold)  # type: list[Gene]
                npn = hmmer_res_to_npn(
                    scaffold,
                    hmmer_result,
                    max_evalue=max_evalue,
                    penalty_black=penalty_black,
                    penalty_white=penalty_white,
                    pvogs_black_list=pvogs_black_list,
                    pvogs_white_list=pvogs_white_list,
                )
                gc = hmmer_res_to_gc(scaffold, hmmer_result, max_evalue=max_evalue)
                if self.mode == "abs":
                    gc = np.array(gc)
                    gc = np.absolute(gc - mean_gc) + mean_gc
                    gc[gc == np.absolute(0 - mean_gc) + mean_gc] = 0
                phages = self.finder.find_phages(npn, gc)

                for phage in phages:
                    phage_num += 1
                    phage_name = "%s_prophage%d" % (scaffold.name, phage_num)
                    begin = genes[phage.begin].begin - 1
                    end = genes[phage.end].end - 1
                    phage_genes = genes[phage.begin : (phage.end + 1)]

                    hmmer_records = [
                        hmmer_result.min_record(
                            hmmer_result.get_records(scaffold.name, gene.name)
                        )
                        for gene in genes[phage.begin : phage.end]
                    ]
                    hmmer_pvogs_records = (
                        record.vog_name
                        for record in hmmer_records
                        if record and record.evalue <= max_evalue
                    )
                    pvogs_records_str = ", ".join(hmmer_pvogs_records)
                    taxonomy = define_taxonomy(pvogs_records_str)
                    hmmer_records = [
                        record
                        for record in hmmer_records
                        if record and record.evalue <= max_evalue
                    ]
                    transposable = if_transposable(hmmer_records)

                    if gff:
                        gff_prophage.append(
                            "\t".join(
                                [
                                    scaffold.name,
                                    ".",
                                    "prophage",
                                    str(begin + 1),
                                    str(end + 1),
                                    ".",
                                    ".",
                                    ".",
                                    "ID=%s;taxonomy=%s;transposable=%s"
                                    % (phage_name, taxonomy, transposable),
                                ]
                            )
                        )
                    if bed:
                        bed_prophage.append(
                            "\t".join(
                                [
                                    scaffold.name,
                                    str(begin),
                                    str(end + 1),
                                    phage_name,
                                    ".",
                                    ".",
                                ]
                            )
                        )

                    if gff or bed:
                        for record in phage_genes:
                            hmmer_record = hmmer_result.min_record(
                                hmmer_result.get_records(scaffold.name, record.name)
                            )
                            pvog = (
                                "." if hmmer_record is None else hmmer_record.vog_name
                            )
                            evalue = (
                                "." if hmmer_record is None else hmmer_record.evalue
                            )
                            if gff:
                                gff_gene.append(
                                    "\t".join(
                                        [
                                            scaffold.name,
                                            ".",
                                            "gene",
                                            str(record.begin),
                                            str(record.end),
                                            str(evalue),
                                            "+" if record.strand > 0 else "-",
                                            ".",
                                            "ID=%s;Parent=%s;pvog=%s"
                                            % (record.name, phage_name, pvog),
                                        ]
                                    )
                                )
                            if bed:
                                bed_gene.append(
                                    "\t".join(
                                        [
                                            scaffold.name,
                                            str(record.begin - 1),
                                            str(record.end),
                                            record.name,
                                            str(evalue),
                                            "+" if record.strand > 0 else "-",
                                        ]
                                    )
                                )

                    if self._print_vogs:
                        writer.writerow(
                            (
                                scaffold.name,
                                phage_name,
                                begin,
                                end,
                                transposable,
                                taxonomy,
                                pvogs_records_str,
                            )
                        )
                    else:
                        writer.writerow(
                            (
                                scaffold.name,
                                phage_name,
                                begin,
                                end,
                                transposable,
                                taxonomy,
                            )
                        )

                    sequence, record = form_sequence(
                        self._filename,
                        "%s_prophage_%d" % (scaffold.name, phage_num),
                        begin,
                        end,
                        scaffold.name,
                    )
                    transposables_status.append(transposable)
                    the_phage_info = [
                        begin,
                        end,
                        taxonomy,
                        sequence,
                        pvogs_records_str,
                    ]
                    phage_info[-1][1].append(the_phage_info)
                    if self._save_fasta:
                        with open(self._output + ".fasta", "a") as f:
                            SeqIO.write(
                                SeqRecord(
                                    record,
                                    id=phage_name,
                                    description="%s_prophage_%d"
                                    % (scaffold.name, phage_num),
                                ),
                                f,
                                "fasta",
                            )
                    if (not self._no_html) and (self._output):
                        plotly_plots.append(plot_html(hmmer_records, begin, end))

                phage_info = (
                    phage_info if len(phage_info[-1][1]) > 0 else phage_info[:-1]
                )

            wtp_output = [the_phage_info[0] for the_phage_info in phage_info]

            if self._wtp:
                with open(self._output_wtp, "w") as f:
                    f.write("\n".join(wtp_output))
            if gff:
                with open(self._output + ".gff3", "w") as f:
                    f.write(
                        "\n".join(gff_base + gff_scaffold + gff_prophage + gff_gene)
                    )
            if bed:
                with open(self._output + ".bed", "w") as f:
                    f.write("\n".join(gff_base + bed_prophage + bed_gene))

            if (len(phage_info) > 0) and (not self._no_html) and (self._output):
                html = form_html_document(
                    phage_info,
                    transposables_status,
                    plotly_plots,
                    self._filename,
                    self._uuid,
                )
                with open(self._output + ".html", "w") as f:
                    f.write(html)
                if not self._not_open:
                    os.system(
                        "xdg-open "
                        + self._output
                        + ".html"
                        + " > /dev/null 2>/dev/null"
                    )
