from busco.busco_tools.base import BaseRunner, NoGenesError
import os
import re
from collections import defaultdict
from busco.BuscoLogger import BuscoLogger
from busco.BuscoLogger import LogDecorator as log
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
import shutil
import numpy as np
from configparser import NoOptionError
import subprocess

logger = BuscoLogger.get_logger(__name__)


class AugustusParsingError(Exception):
    def __init__(self):
        pass


class AugustusRunner(BaseRunner):

    ACCEPTED_PARAMETERS = [
        "strand",
        "genemodel",
        "singlestrand",
        "hintsfile",
        "extrinsicCfgFile",
        "maxDNAPieceSize",
        "protein",
        "introns",
        "start",
        "stop",
        "cds",
        "AUGUSTUS_CONFIG_PATH",
        "alternatives-from-evidence",
        "alternatives-from-sampling",
        "sample",
        "minexonintronprob",
        "minmeanexonintronprob",
        "maxtracks",
        "gff3",
        "UTR",
        "outfile",
        "noInFrameStop",
        "noprediction",
        "contentmodels",
        "translation_table",
        "temperature",
        "proteinprofile",
        "progress",
        "predictionStart",
        "predictionEnd",
        "uniqueGeneId",
    ]

    name = "augustus"
    cmd = "augustus"

    def __init__(self):
        self.gene_details = None
        self._augustus_config_path = os.environ.get("AUGUSTUS_CONFIG_PATH")
        self.config.set("busco_run", "augustus_config_path", self._augustus_config_path)
        self._target_species = self.config.get("busco_run", "augustus_species")
        super().__init__()
        self._output_folder = os.path.join(self.run_folder, "augustus_output")
        self.tmp_dir = os.path.join(self._output_folder, "tmp")
        self.pred_genes_dir_initial = os.path.join(
            self._output_folder, "predicted_genes_initial_run"
        )
        self.pred_genes_dir_rerun = os.path.join(
            self._output_folder, "predicted_genes_rerun"
        )
        self.extracted_prot_dir = os.path.join(
            self._output_folder, "extracted_proteins"
        )
        self.err_logfile = os.path.join(self.log_folder, "augustus_err.log")

        try:
            self.extra_params = self.config.get(
                "busco_run", "augustus_parameters"
            ).replace(",", " ")
        except NoOptionError:
            self.extra_params = ""
        self.chunksize = 10

        self.gff_dir = os.path.join(self._output_folder, "gff")
        self.err_logfiles = []
        self.any_gene_found = False
        self.param_keys = []
        self.param_values = []

        self.create_dirs([self._output_folder, self.extracted_prot_dir, self.gff_dir])

        self.init_checkpoint_file()

    def configure_runner(self, seqs_path, coords, sequences_aa, sequences_nt):
        self.run_number += 1

        # Placed here to allow reconfiguration for rerun
        self._target_species = self.config.get("busco_run", "augustus_species")

        self.check_tool_dependencies()
        self.gene_details = defaultdict(list)
        self.output_sequences = []

        self.seqs_path = seqs_path
        self.coords = coords

        self.sequences_aa = sequences_aa
        self.sequences_nt = sequences_nt

        self.pred_genes_dir = (
            self.pred_genes_dir_rerun
            if self.run_number == 2
            else self.pred_genes_dir_initial
        )

        # self.tmp_dir placed here to allow it to be recreated during reconfiguration for rerun
        self.create_dirs([self.pred_genes_dir, self.tmp_dir])

    @property
    def output_folder(self):
        return self._output_folder

    def check_tool_dependencies(self):
        """
        check dependencies on files and folders
        properly configured.
        :raises SystemExit: if Augustus config path is not writable or
        not set at all
        :raises SystemExit: if Augustus config path does not contain
        the needed species
        present
        """
        try:
            augustus_species_dir = os.path.join(self._augustus_config_path, "species")
            if not os.access(augustus_species_dir, os.W_OK):
                raise SystemExit(
                    "Cannot write to Augustus species folder, please make sure you have write "
                    "permissions to {}".format(augustus_species_dir)
                )

        except TypeError:
            raise SystemExit("The environment variable AUGUSTUS_CONFIG_PATH is not set")

        if not os.path.exists(os.path.join(augustus_species_dir, self._target_species)):
            # Exclude the case where this is a restarted run and the retraining parameters have already been moved.
            if (
                self.config.getboolean("busco_run", "restart")
                and self.run_number == 2
                and os.path.exists(
                    os.path.join(
                        self._output_folder,
                        "retraining_parameters",
                        self._target_species,
                    )
                )
            ):
                pass
            else:
                raise SystemExit(
                    'Impossible to locate the species "{0}" in Augustus species folder'
                    " ({1}), check that AUGUSTUS_CONFIG_PATH is properly set"
                    " and contains this species. \n\t\tSee the help if you want "
                    "to provide an alternative species".format(
                        self._target_species, augustus_species_dir
                    )
                )

    @log(
        "Running Augustus prediction using {} as species:",
        logger,
        attr_name="_target_species",
    )
    def run(self):
        super().run()
        if self.extra_params:
            logger.info(
                "Additional parameters for Augustus are {}: ".format(self.extra_params)
            )
            self.param_keys, self.param_values = self.parse_parameters()

        self.total = self._count_jobs()
        self.run_jobs()

    def process_output(self):
        logger.info("Extracting predicted proteins...")
        files = [
            f
            for f in sorted(os.listdir(self.pred_genes_dir))
            if any(busco_id in f for busco_id in self.coords)
        ]
        for filename in files:
            self._extract_genes_from_augustus_output(filename)

        if not self.any_gene_found and self.run_number == 1:
            raise NoGenesError("Augustus")

        self.gene_details = dict(self.gene_details)

        self._merge_stderr_logs()
        self._remove_individual_err_logs()

        return

    def _count_jobs(self):
        n = 0
        for busco_group, contigs in self.coords.items():
            for _ in contigs:
                n += 1
        return n

    def sort_jobs(self):
        jobs_size_info = []
        for busco_group, contigs in self.coords.items():

            for contig_name, contig_info in contigs.items():
                contig_start = contig_info["contig_start"]
                contig_end = contig_info["contig_end"]
                pred_size = int(contig_end) - int(contig_start)
                jobs_size_info.append(
                    {
                        "busco_group": busco_group,
                        "contig_name": contig_name,
                        "contig_start": contig_start,
                        "contig_end": contig_end,
                        "pred_size": pred_size,
                    }
                )
        job_sizes = [item["pred_size"] for item in jobs_size_info]
        new_job_order = np.argsort(job_sizes)[::-1]
        ordered_jobs = [jobs_size_info[i] for i in new_job_order]
        return ordered_jobs

    def generate_job_args(self):
        contig_ordinal_inds = defaultdict(int)
        njobs = 0

        ordered_jobs = self.sort_jobs()

        for job_info in ordered_jobs:
            contig_name = job_info["contig_name"]
            busco_group = job_info["busco_group"]
            contig_start = job_info["contig_start"]
            contig_end = job_info["contig_end"]
            contig_tmp_file = "{}.temp".format(
                contig_name[:100]
            )  # Avoid very long filenames
            contig_ordinal_inds[busco_group] += 1
            output_index = contig_ordinal_inds[busco_group]
            out_filename = os.path.join(
                self.pred_genes_dir, "{}.out.{}".format(busco_group, output_index)
            )
            njobs += 1

            yield busco_group, contig_tmp_file, contig_start, contig_end, out_filename

    @log(
        "Additional parameters for Augustus are {}: ",
        logger,
        attr_name="_target_species",
    )
    def parse_parameters(self):
        accepted_keys = []
        accepted_values = []
        if self.extra_params:
            self.extra_params = self.extra_params.strip("\" '")
            try:
                if self.extra_params.startswith("--"):
                    key_val_pairs = self.extra_params.split(" --")
                    for kv in key_val_pairs:
                        key_vals = kv.strip("- ").split("=")
                        if len(key_vals) == 2:
                            key, val = key_vals
                            if key in type(self).ACCEPTED_PARAMETERS:
                                accepted_keys.append(key.strip())
                                accepted_values.append(val.strip())
                            else:
                                logger.warning(
                                    "{} is not an accepted parameter for Augustus.".format(
                                        key
                                    )
                                )
                        else:
                            raise AugustusParsingError
                else:
                    raise AugustusParsingError
            except AugustusParsingError:
                logger.warning(
                    "Augustus parameters are not correctly formatted. Please enter them as follows: "
                    '"--param1=value1 --param2=value2" etc. Proceeding without additional parameters.'
                )
                return [], []
        return accepted_keys, accepted_values

    def _merge_stderr_logs(self):
        with open(self.err_logfile, "a") as f:
            for err_logfile in self.err_logfiles:
                with open(err_logfile, "r") as g:
                    content = g.readlines()
                    f.writelines(content)
        return

    def _remove_individual_err_logs(self):
        shutil.rmtree(self.tmp_dir)
        return

    def get_version(self):  # todo: need to handle all possible exceptions
        augustus_help_output = subprocess.check_output(
            [self.cmd, "--version"], stderr=subprocess.STDOUT, shell=False
        )
        augustus_help_output = augustus_help_output.decode("utf-8")
        s = augustus_help_output.split("\n")[0]
        augustus_version = s[s.find("(") + 1 : s.find(")")]
        return augustus_version

    def configure_job(
        self, busco_group, contig_tmp_file, contig_start, contig_end, out_filename
    ):
        # Augustus does not provide an option to write to an output file, so have to change the pipe target from the
        # log file to the desired output file
        self.logfile_path_out = out_filename
        err_logfile = os.path.join(
            self.tmp_dir, os.path.basename(out_filename.rpartition(".out")[0] + ".err")
        )
        self.logfile_path_err = err_logfile
        self.err_logfiles.append(err_logfile)

        augustus_job = self.create_job()
        augustus_job.add_parameter("--codingseq=1")
        augustus_job.add_parameter(
            "--proteinprofile={}".format(
                os.path.join(
                    self.lineage_dataset, "prfl", "{}.prfl".format(busco_group)
                )
            )
        )
        augustus_job.add_parameter("--predictionStart={}".format(contig_start))
        augustus_job.add_parameter("--predictionEnd={}".format(contig_end))
        augustus_job.add_parameter("--species={}".format(self._target_species))
        for k, key in enumerate(self.param_keys):
            augustus_job.add_parameter("--{}={}".format(key, self.param_values[k]))
        augustus_job.add_parameter(os.path.join(self.seqs_path, contig_tmp_file))
        return augustus_job

    def _extract_genes_from_augustus_output(self, filename):
        # todo: consider parallelizing this and other parsing functions

        gene_id = None
        gene_info = []
        sequences_aa = []
        sequences_nt = []
        gene_found = False
        completed_record = False

        with open(
            os.path.join(self.pred_genes_dir, filename), "r", encoding="utf-8"
        ) as f:
            # utf-8 encoding needed to handle the umlaut in the third line of the file.
            gene_info_section = False
            nt_sequence_section = False
            aa_sequence_section = False
            nt_sequence_parts = []
            aa_sequence_parts = []

            for line in f:

                if aa_sequence_section:
                    if "]" in line:
                        line = line.strip().lstrip("# ").rstrip("]")
                        aa_sequence_parts.append(line)
                        aa_sequence_section = False
                        completed_record = True
                        if gene_id is not None:
                            aa_sequence = "".join(aa_sequence_parts)
                            nt_sequence = "".join(nt_sequence_parts)
                            seq_record_aa = SeqRecord(
                                Seq(aa_sequence.upper()), id=gene_id
                            )
                            seq_record_nt = SeqRecord(
                                Seq(nt_sequence.upper()), id=gene_id
                            )
                            sequences_aa.append(seq_record_aa)
                            sequences_nt.append(seq_record_nt)
                            aa_sequence_parts = []
                            nt_sequence_parts = []
                            gene_id = None
                        continue

                    else:
                        line = line.strip().lstrip("# ").rstrip("]")
                        aa_sequence_parts.append(line)
                        continue

                if line.startswith("# protein"):
                    nt_sequence_section = False
                    aa_sequence_section = True
                    if "]" in line:
                        line = line.strip().rstrip("]").split("[")
                        aa_sequence_parts.append(line[1])
                        aa_sequence_section = False
                        completed_record = True
                        if gene_id is not None:
                            aa_sequence = "".join(aa_sequence_parts)
                            nt_sequence = "".join(nt_sequence_parts)
                            seq_record_aa = SeqRecord(
                                Seq(aa_sequence.upper()), id=gene_id
                            )
                            seq_record_nt = SeqRecord(
                                Seq(nt_sequence.upper()), id=gene_id
                            )
                            sequences_aa.append(seq_record_aa)
                            sequences_nt.append(seq_record_nt)
                            aa_sequence_parts = []
                            nt_sequence_parts = []
                            gene_id = None
                    else:
                        line = line.strip().rstrip("]").split("[")
                        aa_sequence_parts.append(line[1])
                    continue

                if nt_sequence_section:
                    line = line.strip().lstrip("# ").rstrip("]")
                    nt_sequence_parts.append(line)
                    continue

                if line.startswith("# coding sequence"):
                    gene_info = []
                    gene_info_section = False
                    nt_sequence_section = True
                    line = (
                        line.strip().rstrip("]").split("[")
                    )  # Extract sequence part of line
                    nt_sequence_parts.append(line[1])
                    continue

                if gene_info_section:
                    line = line.strip().split()
                    seq_name = line[0].strip()
                    gene_start = line[3].strip()
                    gene_end = line[4].strip()
                    strand = line[6].strip()
                    if not gene_id:
                        gene_id = "{}:{}-{}".format(seq_name, gene_start, gene_end)
                        self.gene_details[gene_id].append(
                            {
                                "gene_start": gene_start,
                                "gene_end": gene_end,
                                "strand": strand,
                            }
                        )
                    gene_info.append("\t".join(line))
                    continue

                if line.startswith("# start gene"):
                    gene_found = True
                    self.any_gene_found = True
                    gene_info_section = True
                    completed_record = False
                    continue

            if gene_found and not completed_record:
                logger.warning("Augustus output file {} truncated".format(filename))

        self.sequences_aa.update({record.id: record for record in sequences_aa})
        self.sequences_nt.update({record.id: record for record in sequences_nt})
        if gene_found:
            self._write_sequences_to_file(filename, sequences_nt, sequences_aa)

        return

    def make_gff_files(self, single_copy_buscos):

        for b in single_copy_buscos:
            gene_info = []
            busco_files = [
                f for f in os.listdir(self.pred_genes_dir) if f.startswith(b)
            ]
            pred_genes_dir_current = self.pred_genes_dir
            if (
                len(busco_files) == 0
                and pred_genes_dir_current == self.pred_genes_dir_rerun
            ):
                pred_genes_dir_current = self.pred_genes_dir_initial
                busco_files = [
                    f for f in os.listdir(pred_genes_dir_current) if f.startswith(b)
                ]
            gff_filename = os.path.join(self.gff_dir, "{}.gff".format(b))
            single_copy_busco_gene = list(single_copy_buscos[b].keys())[0]
            gene_id_parts = single_copy_busco_gene.split(":")
            if (
                len(gene_id_parts) > 2
            ):  # if a ":" is present in the gene id, we don't want to break it up
                gene_id_parts = [":".join(gene_id_parts[:-1]), gene_id_parts[-1]]
            single_copy_busco_gene_id = gene_id_parts[0]
            (
                single_copy_busco_gene_start_coord,
                single_copy_busco_gene_end_coord,
            ) = gene_id_parts[1].split("-")
            gene_found = False
            for filename in busco_files:
                match_number = filename.split(".")[-1]
                with open(
                    os.path.join(pred_genes_dir_current, filename),
                    "r",
                    encoding="utf-8",
                ) as f:
                    gene_info_section = False
                    for line in f:
                        if gene_info_section and line.startswith("# coding sequence"):
                            with open(gff_filename, "w") as g:
                                g.write("\n".join(gene_info) + "\n")
                            gene_info = []
                            break

                        if line.startswith("# start gene"):
                            gene_info_section = True
                            continue

                        if gene_info_section:
                            line = line.strip().split()
                            seq_name = line[0]
                            gene_start = line[3]
                            gene_end = line[4]
                            if gene_found or (
                                seq_name == single_copy_busco_gene_id
                                and gene_start == single_copy_busco_gene_start_coord
                                and gene_end == single_copy_busco_gene_end_coord
                            ):
                                gene_found = True
                                gene_id_info = line[-1]
                                line[-1] = self.edit_gene_identifier(
                                    gene_id_info, match_number
                                )
                                if len(line) == 12:
                                    gene_id_info_2 = line[-3]
                                    line[-3] = self.edit_gene_identifier(
                                        gene_id_info_2, match_number
                                    )
                                gene_info.append("\t".join(line))
                            else:
                                gene_info_section = False
                            continue
                if gene_found:
                    break
            if not gene_found and self.run_number == 1:
                raise SystemExit(
                    "Unable to find single copy BUSCO gene in Augustus output."
                )

        return

    def edit_gene_identifier(self, orig_str, match_num):
        modified_str = re.sub(
            r"g([0-9])", r"r{}.m{}.g\1".format(self.run_number, match_num), orig_str
        )
        return modified_str

    def _write_sequences_to_file(self, filename, sequences_nt, sequences_aa):

        filename_parts = filename.rpartition(".out")
        output_fna = os.path.join(
            self.extracted_prot_dir, filename_parts[0] + ".fna" + filename_parts[-1]
        )
        output_faa = os.path.join(
            self.extracted_prot_dir, filename_parts[0] + ".faa" + filename_parts[-1]
        )
        self.output_sequences.append(output_faa)

        with open(output_fna, "w") as out_fna:
            SeqIO.write(sequences_nt, out_fna, "fasta")
        with open(output_faa, "w") as out_faa:
            SeqIO.write(sequences_aa, out_faa, "fasta")

        return

    def move_retraining_parameters(self):
        """
        This function moves retraining parameters from augustus species folder
        to the run folder
        """
        augustus_species_path = os.path.join(
            self._augustus_config_path, "species", self._target_species
        )
        if os.path.exists(augustus_species_path):
            new_path = os.path.join(
                self._output_folder, "retraining_parameters", self._target_species
            )
            shutil.move(augustus_species_path, new_path)
        elif self.config.getboolean("busco_run", "restart") and os.path.exists(
            os.path.join(
                self._output_folder, "retraining_parameters", self._target_species
            )
        ):
            pass
        else:
            logger.warning("Augustus did not produce a retrained species folder.")
        return


class GFF2GBRunner(BaseRunner):

    name = "gff2gbSmallDNA.pl"
    cmd = "gff2gbSmallDNA.pl"

    def __init__(self):
        super().__init__()
        self._output_folder = os.path.join(self.run_folder, "augustus_output")
        self.gff_folder = os.path.join(self._output_folder, "gff")
        self.gb_folder = os.path.join(self._output_folder, "gb")
        self.create_dirs([self.gff_folder, self.gb_folder])

        self.init_checkpoint_file()

    def configure_runner(self, single_copy_buscos):
        self.run_number += 1
        self.single_copy_buscos = single_copy_buscos

    def run(self):
        super().run()
        self.total = self._count_jobs()
        self.run_jobs()

    def _count_jobs(self):
        n = len(self.single_copy_buscos)
        return n

    def generate_job_args(self):
        for busco_id in self.single_copy_buscos:
            yield busco_id

    def configure_job(self, busco_id):
        gff2_gb_small_dna_pl_job = self.create_job()
        gff2_gb_small_dna_pl_job.add_parameter(
            os.path.join(self.gff_folder, "{}.gff".format(busco_id))
        )
        gff2_gb_small_dna_pl_job.add_parameter(self.input_file)
        gff2_gb_small_dna_pl_job.add_parameter("1000")
        gff2_gb_small_dna_pl_job.add_parameter(
            os.path.join(self.gb_folder, "{}.raw.gb".format(busco_id))
        )
        return gff2_gb_small_dna_pl_job

    def check_tool_dependencies(self):
        pass

    def get_version(self):
        return

    @property
    def output_folder(self):
        return self._output_folder


class NewSpeciesRunner(BaseRunner):

    name = "new_species.pl"
    cmd = "new_species.pl"

    def __init__(self):
        super().__init__()
        self._output_folder = os.path.join(self.run_folder, "augustus_output")
        self.new_species_name = "BUSCO_{}".format(os.path.basename(self.main_out))
        self.init_checkpoint_file()
        self.run_number += 1

    def run(self):
        super().run()
        self.total = 1
        self.run_jobs()

    def configure_job(self, *args):

        new_species_pl_job = self.create_job()
        # bacteria clade needs to be flagged as "prokaryotic"
        if self.domain == "prokaryota":
            new_species_pl_job.add_parameter("--prokaryotic")
        new_species_pl_job.add_parameter(
            "--species={}".format(os.path.basename(self.new_species_name))
        )
        return new_species_pl_job

    def check_tool_dependencies(self):
        pass

    def generate_job_args(self):
        yield

    def get_version(self):
        return

    @property
    def output_folder(self):
        return self._output_folder


class ETrainingRunner(BaseRunner):

    name = "etraining"
    cmd = "etraining"

    def __init__(self):
        super().__init__()
        self._output_folder = os.path.join(self.run_folder, "augustus_output")
        self._gb_folder = os.path.join(self._output_folder, "gb")
        self.augustus_config_path = self.config.get("busco_run", "augustus_config_path")
        self._training_file = os.path.join(self._output_folder, "training_set.db")

        self.init_checkpoint_file()

    def configure_runner(self, new_species_name):
        self.run_number += 1
        self.new_species_name = new_species_name
        self._merge_gb_files()

    def run(self):
        super().run()
        self.total = 1
        self.run_jobs()
        self._validate_run()

    def check_tool_dependencies(self):
        pass

    def generate_job_args(self):
        yield

    def _merge_gb_files(self):
        """Concatenate all GB files into one large file"""
        with open(self._training_file, "w") as outfile:
            for fname in os.listdir(self._gb_folder):
                with open(os.path.join(self._gb_folder, fname), "r") as infile:
                    outfile.writelines(infile.readlines())
        return

    def _validate_run(self):
        species_filepath = os.path.join(
            self.augustus_config_path, "species", self.new_species_name
        )
        if os.path.exists(species_filepath) and any(
            "exon_probs" in f for f in os.listdir(species_filepath)
        ):
            return
        else:
            raise SystemExit(
                "Retraining did not complete correctly. Check your Augustus config path environment variable."
            )

    def configure_job(self, *args):
        etraining_job = self.create_job()
        etraining_job.add_parameter("--species={}".format(self.new_species_name))
        etraining_job.add_parameter(
            os.path.join(self.run_folder, "augustus_output", "training_set.db")
        )
        return etraining_job

    def get_version(self):
        return

    @property
    def output_folder(self):
        return self._output_folder


class OptimizeAugustusRunner(BaseRunner):

    name = "optimize_augustus.pl"
    cmd = "optimize_augustus.pl"

    def __init__(self):
        super().__init__()
        self._output_folder = None
        self.training_set_db = None
        self.new_species_name = None

    def configure_runner(self, output_folder, new_species_name):
        self.run_number += 1
        self._output_folder = output_folder
        self.training_set_db = os.path.join(self._output_folder, "training_set.db")
        self.new_species_name = new_species_name

        self.init_checkpoint_file()

    def configure_job(self, *args):
        optimize_augustus_pl_job = self.create_job()
        optimize_augustus_pl_job.add_parameter("--cpus={}".format(self.cpus))
        optimize_augustus_pl_job.add_parameter(
            "--species={}".format(self.new_species_name)
        )
        optimize_augustus_pl_job.add_parameter(self.training_set_db)
        return optimize_augustus_pl_job

    def run(self):
        super().run()
        self.total = 1
        self.run_jobs()

    def generate_job_args(self):
        yield

    def check_tool_dependencies(self):
        pass

    def get_version(self):
        return

    @property
    def output_folder(self):
        return self._output_folder
