# -*- coding: utf-8 -*-
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#    A copy of the GNU General Public License is available at
#    http://www.gnu.org/licenses/gpl-3.0.html

"""Process msp and compute tree"""
from pathlib import Path
from collections import defaultdict
from meteor.session import Session, Component
from meteor.phylogeny import Phylogeny
from dataclasses import dataclass
from tempfile import mkdtemp
import ete3  # type: ignore[import]
from ete3 import Tree  # , TreeStyle
import logging
import sys
import pandas as pd
import lzma
from shutil import rmtree
from typing import ClassVar


@dataclass
class TreeBuilder(Session):
    """Counter session map and count"""

    DEFAULT_MAX_GAP: ClassVar[float] = 0.5
    DEFAULT_GAP_CHAR: ClassVar[str] = "-"
    OUTPUT_FORMATS: ClassVar[list[str | None]] = [None, "png", "svg", "pdf", "txt"]
    DEFAULT_OUTPUT_FORMAT: ClassVar[str | None] = None
    DEFAULT_WIDTH: ClassVar[int] = 500
    DEFAULT_HEIGHT: ClassVar[int] = 500
    DEFAULT_NUM_THREADS: ClassVar[int] = 1

    meteor: type[Component]
    max_gap: float
    width: int
    height: int
    format: str | None
    gap_char: str

    def __post_init__(self) -> None:
        if self.format not in TreeBuilder.OUTPUT_FORMATS:
            raise ValueError(f"{self.format} is not a valid output format")

        self.meteor.tmp_dir = Path(mkdtemp(dir=self.meteor.tmp_path))
        self.meteor.tree_dir.mkdir(exist_ok=True, parents=True)

    def concatenate(self, msp_file_dict: dict[str, list[Path]]) -> list[Path]:
        """Concatenate fasta file from distinct samples
        :param msp_file_dict: (dict) A dictionnary of each msp version in the different samples
        :return: (list) A list of all concatenated fasta
        """
        msp_list = []
        # Concatenate files that occur in more than one directory
        for filename, paths in msp_file_dict.items():
            if len(paths) > 1:
                res = self.meteor.tree_dir / f"{filename}".replace(
                    ".fasta.xz", ".fasta"
                )
                with res.open("wt", encoding="UTF-8") as outfile:
                    for path in paths:
                        with lzma.open(path, "rt") as infile:
                            outfile.write(infile.read())
                msp_list += [res]
        logging.info("%d MSPs are available for tree analysis.", len(msp_list))
        return msp_list

    def get_msp_distance(self, tree: ete3.TreeNode) -> pd.DataFrame:
        samples = list(tree)
        distance_matrix = pd.DataFrame(
            index=[n.name for n in samples], columns=[n.name for n in samples]
        )
        distance_matrix.values[[range(len(samples))] * 2] = 0.0
        for i in samples[:-1]:
            for j in samples[1:]:
                distance = tree.get_distance(i, j)
                distance_matrix.loc[i.name, j.name] = distance
                distance_matrix.loc[j.name, i.name] = distance
        return distance_matrix

    def execute(self) -> None:
        "Merge all files generated by either mapper or profiler."
        # Fetch all census ini files
        all_census = list(Path(self.meteor.strain_dir).glob("**/*census_stage_3.json"))
        if len(all_census) == 0:
            logging.error("No census stage found in the specified repository.")
            sys.exit(1)
        else:
            logging.info("%d samples have been detected.", len(all_census))
        msp_file_dict = defaultdict(list)
        for filepath in self.meteor.strain_dir.glob("**/*.fasta.xz"):
            if not filepath.name.endswith("_consensus.fasta.xz"):
                msp_file_dict[filepath.name].append(filepath)
        # Concatenate msp files
        msp_file_list = self.concatenate(msp_file_dict)
        # Compute phylogenies
        phylogeny_process = Phylogeny(
            self.meteor, msp_file_list, self.max_gap, self.gap_char
        )
        phylogeny_process.execute()
        # Analyze tree data
        for msp_file in msp_file_list:
            # Check for a .raxml.bestTree file
            raxml_tree_file = self.meteor.tree_dir / f"{msp_file.stem}.raxml.bestTree"
            if raxml_tree_file.exists():
                tree_file = raxml_tree_file
            else:
                # If no .raxml.bestTree file exists, check for a .tree file
                tree_file = self.meteor.tree_dir / f"{msp_file.stem}.tree"
            img_file = self.meteor.tree_dir / f"{msp_file.stem}.{self.format}"
            # try:
            msp_tree = Tree(str(tree_file.resolve()))
            # Generate a distance msp by msp
            matrix = self.get_msp_distance(msp_tree)
            matrix.to_csv(self.meteor.tree_dir / f"{msp_file.stem}.tsv", sep="\t")
            # Draw trees
            if not self.format:
                pass
            elif self.format == "txt":
                with img_file.open("wt", encoding="UTF-8") as outfile:
                    outfile.write(msp_tree.get_ascii(show_internal=True))
            else:
                # ts = TreeStyle()
                # ts.show_leaf_name = True
                # ts.show_branch_length = True
                msp_tree.render(
                    str(img_file.resolve()),
                    w=self.width,
                    h=self.height,
                    units="px",
                    dpi=300,
                )
            # except ete3.parser.newick.NewickError:
            #     logging.info("Not sufficient info in %s.", str(tree_file.resolve()))
        rmtree(self.meteor.tmp_dir, ignore_errors=True)
