# SPDX-FileCopyrightText: 2025 Theresa Pollinger
#
# SPDX-License-Identifier: GPL-3.0-or-later

import warnings

from pathlib import Path
from typing import Sequence, Union, Mapping, Optional
import subprocess

from dyada.coordinates import (
    CoordinateInterval,
)
from dyada.descriptor import branch_generator, RefinementDescriptor
from dyada.discretization import (
    Discretization,
    discretization_to_location_stack_strings,
)
from dyada.drawing_util import (
    side_corners_generator,
    get_colors_byte,
    letter_counter,
)


def latex_write_and_compile(latex_string: str, filename: str) -> None:
    dirname = Path.cwd()
    with open(dirname / filename, "w") as f:
        f.write(latex_string)

    # if in environment, run pdflatex to generate pdf
    # this needs `latexmk`, `pdflatex`, and `tikz` to be installed, e.g.
    # through `apt install texlive-latex-extra latexmk`
    try:
        subprocess.check_output(
            ["latexmk", "-interaction=nonstopmode", "-pdf", filename],
            cwd=dirname,
        )
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        warnings.warn(f"Error while running latexmk on {filename}: {e}")
    try:  # try to clean up regardless of success
        subprocess.run(
            ["latexmk", "-c"],
            check=False,
            cwd=dirname,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.STDOUT,
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass


def latex_add_color_defs(
    tikz_string: str, num_colors: int, colormap_name="CET_R3"
) -> str:
    for leaf, color in enumerate(get_colors_byte(num_colors, colormap_name)):
        color_str = "color_%d" % leaf
        tikz_string += "\\definecolor{%s}{RGB}{%d,%d,%d}\n" % (
            color_str,
            color[0],
            color[1],
            color[2],
        )
    return tikz_string


def plot_boxes_2d_tikz(
    intervals: Union[Sequence[CoordinateInterval], Mapping[CoordinateInterval, str]],
    labels: Optional[Sequence[str]],
    projection: Sequence[int],
    filename: Optional[str] = None,
    connect_centers=False,
    **kwargs,
) -> None:
    tikz_string = R"""\documentclass{standalone}
% autogenerated with dyada : https://github.com/freifrauvonbleifrei/DyAda
\usepackage{xcolor}
\usepackage{relsize}
\usepackage{tikz}
\begin{document}
\begin{tikzpicture}
"""
    # add one white rectangle to the background
    tikz_string += "\\draw[very thin, white] (0,0) rectangle (1,1);\n"
    letter_counter_first = letter_counter(len(intervals))

    def tikz_rectangle(interval: CoordinateInterval, option_string="", label_string=""):
        lower = interval[0][projection]
        upper = interval[1][projection]
        rectangle_string = "\\draw[%s] (%f,%f) rectangle (%f,%f);\n" % (
            option_string,
            *lower,
            *upper,
        )
        middle = (lower + upper) / 2.0
        min_extent = min(upper - lower)  # type: ignore
        if label_string != "":
            if min_extent <= 0.125:
                label_string = "\\tiny \\relsize{-1} " + label_string
            elif min_extent <= 0.25:
                label_string = "\\tiny \\relsize{-0.5}" + label_string
            elif min_extent <= 0.5:
                label_string = "\\footnotesize " + label_string
        label_tex = f"\\node[inner sep=0] ({next(letter_counter_first)}) at ({middle[0]},{middle[1]}) {{{label_string}}};\n"

        return rectangle_string + label_tex

    if labels is None:

        def none_iter():
            while True:
                yield ""

        label_iter = none_iter()
    else:
        label_iter = iter(labels)
    tikz_string = latex_add_color_defs(tikz_string, len(intervals))
    for grid_idx, interval in enumerate(intervals):
        color_str = "color_%d" % grid_idx
        option_string = "very thin, gray, fill=%s,fill opacity=0.3" % (color_str)
        tikz_string += tikz_rectangle(interval, option_string, next(label_iter))

    if connect_centers:
        # connect the named nodes with lines
        letter_counter_again = letter_counter(len(intervals))
        connect_string = "\\draw[thick, densely dotted] "
        for node_name in letter_counter_again:
            connect_string += f"({node_name}) -- "
        tikz_string += connect_string[:-4] + ";\n"

    tikz_string += "\\end{tikzpicture}\n"
    tikz_string += "\\end{document}\n"
    if filename is None:
        filename = "tikz_rectangles"
    filename += ".tex"
    latex_write_and_compile(tikz_string, filename)


# inspired by @griegler : https://github.com/griegler/octnet/issues/28
def plot_boxes_3d_tikz(
    intervals: Union[Sequence[CoordinateInterval], Mapping[CoordinateInterval, str]],
    labels: Optional[Sequence[str]],
    projection: Sequence[int],
    wireframe: bool = False,
    filename: Optional[str] = None,
    **kwargs,
) -> None:

    def tikz_cuboid(
        interval: CoordinateInterval, option_string="", label_string=""
    ) -> str:
        tikz_string = ""
        line_string = "\\draw[%s] (%f,%f,%f) -- (%f,%f,%f) -- (%f,%f,%f) -- (%f,%f,%f) -- cycle;\n"
        lower = interval[0][projection]
        upper = interval[1][projection]
        for side_corners in side_corners_generator(lower, upper):
            tikz_string += line_string % (
                option_string,
                *side_corners[0],
                *side_corners[1],
                *side_corners[2],
                *side_corners[3],
            )
        middle = (lower + upper) / 2.0
        min_extent = min(upper - lower)  # type: ignore
        if min_extent < 0.125:
            label_string = "\\tiny \\relsize{-1} " + label_string
        elif min_extent < 0.25:
            label_string = "\\tiny \\relsize{-0.5}" + label_string
        elif min_extent < 0.5:
            label_string = "\\footnotesize " + label_string
        tikz_string += (
            f"\\node at ({middle[0]},{middle[1]},{middle[2]}) {{{label_string}}};\n"
        )
        return tikz_string

    def tikz_grid(intervals, wireframe: bool):
        tikz_string = R"""\documentclass{standalone}
% autogenerated with dyada : https://github.com/freifrauvonbleifrei/DyAda
% inspired by @griegler : https://github.com/griegler/octnet/issues/28
\usepackage{xcolor}
\usepackage{relsize}
\usepackage{tikz,tikz-3dplot}
\begin{document}
\tdplotsetmaincoords{50}{130}
\begin{tikzpicture}[scale=1.0, tdplot_main_coords]"""
        if labels is None:

            def none_iter():
                while True:
                    yield ""

            label_iter = none_iter()
        else:
            label_iter = iter(labels)
        tikz_string = latex_add_color_defs(tikz_string, len(intervals))
        for grid_idx, interval in enumerate(intervals):
            color_str = "color_%d" % grid_idx
            if wireframe:
                option_string = "very thin, %s" % (color_str)
            else:
                option_string = "very thin, gray, fill=%s,fill opacity=0.3" % (
                    color_str
                )
            tikz_string += tikz_cuboid(interval, option_string, next(label_iter))

        tikz_string += "\\end{tikzpicture}\n"
        tikz_string += "\\end{document}\n"
        return tikz_string

    latex_string = tikz_grid(intervals, wireframe)
    if filename is None:
        filename = "tikz_cuboids"
    if wireframe:
        filename += "_wireframe.tex"
    else:
        filename += "_solid.tex"
    latex_write_and_compile(latex_string, filename)


def plot_tree_tikz(
    refinement_descriptor: RefinementDescriptor, labels=None, filename="omnitree"
):
    tikz_string = R"""\documentclass{standalone}
% autogenerated with dyada : https://github.com/freifrauvonbleifrei/DyAda
\usepackage{forest}
\begin{document}
% cf. https://tex.stackexchange.com/questions/332300/draw-lines-on-top-of-tikz-forest
\forestset{%
    declare keylist register={through},
    through={},
    tracing tree/.style={%
    delay={%
        for #1={%
        if phantom={}{through+/.option=name},
        }
    },
    before drawing tree={%
        tikz+/.wrap pgfmath arg={%
        \foreach \i [count=\j, remember=\i as \k] in {##1} \ifnum\j>1 \draw [densely dashed, ->] (\k.west) -- (\i.west)\fi;
        }{(through)}
    },
    }
}
\tikzset{every label/.style={xshift=-2ex,yshift=-2ex , font=\footnotesize, text=red}}
"""
    n_leaves = refinement_descriptor.get_num_boxes()
    tikz_string = latex_add_color_defs(tikz_string, n_leaves)

    assert labels is None or len(labels) == len(refinement_descriptor)

    tikz_string += "\\begin{forest}\n     tracing tree=tree,\n"
    d_zeros = "0" * refinement_descriptor.get_num_dimensions()
    leaf_string = (
        "[" + d_zeros + ",circle,draw,fill=color_%d,fill opacity=0.4,label={%s}]\n"
    )
    last_indent = 0
    tab = "    "
    num_box = 0
    num_patch = 0
    try:
        for current_branch, refinement in branch_generator(refinement_descriptor):
            current_indent = len(current_branch)
            while current_indent < last_indent:
                tikz_string += tab * current_indent + "]\n"
                last_indent -= 1
            if labels is not None:
                label_string = labels[num_patch]
            else:
                label_string = str(num_patch)
            if refinement.count() == 0:
                tikz_string += tab * current_indent + leaf_string % (
                    num_box,
                    label_string,
                )
                num_box += 1
            else:
                tikz_string += (
                    tab * current_indent
                    + "["
                    + refinement.to01()
                    + ",label={"
                    + label_string
                    + "}\n"
                )
            last_indent = current_indent
            num_patch += 1
    except IndexError:
        # if this happens, the descriptor is not valid but we can plot part of it
        warnings.warn("Descriptor is not valid, plotting as much as possible.")
        pass

    tikz_string += "]" * (last_indent)
    tikz_string += "\n\\end{forest}\n\\end{document}"

    if not filename.endswith(".tex"):
        filename += ".tex"
    latex_write_and_compile(tikz_string, filename)


def plot_descriptor_tikz(
    refinement_descriptor: RefinementDescriptor, filename="descriptor"
):
    tikz_string = R"""\documentclass{standalone}
% autogenerated with dyada : https://github.com/freifrauvonbleifrei/DyAda
\usepackage{tikz}
\usetikzlibrary{matrix}
\begin{document}"""
    num_colors = refinement_descriptor.get_num_boxes()
    tikz_string = latex_add_color_defs(tikz_string, num_colors)
    tikz_string += R"""\begin{tikzpicture}[every node/.style={draw,align=center,text height=2ex,minimum width=2ex, inner sep=0.2ex, fill opacity=0.4, text opacity=1}]
    \matrix [draw=none, matrix of nodes,nodes in empty cells]
    {"""
    tab = "    "
    box_counter = 0
    for refinement in refinement_descriptor:
        tikz_string += tab
        if refinement.count() == 0:
            tikz_string += f"|[fill=color_{box_counter}]| "
            box_counter += 1
        tikz_string += refinement.to01() + "&\n"

    # remove last newline and ampersand
    tikz_string = tikz_string[:-2]
    tikz_string += R"""\\
    };   
\end{tikzpicture}
\end{document}
"""
    if not filename.endswith(".tex"):
        filename += ".tex"
    latex_write_and_compile(tikz_string, filename)


def plot_location_stack_tikz(discretization: Discretization, filename="location_stack"):
    descriptor = discretization.descriptor
    location_stack = discretization_to_location_stack_strings(discretization)

    tikz_string = R"""\documentclass[tikz]{standalone}
% autogenerated with dyada : https://github.com/freifrauvonbleifrei/DyAda
\usepackage{tikz}
\usetikzlibrary{matrix}
\begin{document}"""
    num_colors = descriptor.get_num_boxes()
    tikz_string = latex_add_color_defs(tikz_string, num_colors)
    tikz_string += R"""\begin{tikzpicture}[
       every node/.style={align=left,anchor=west,
       text height=2ex,minimum width=2ex,
       inner sep=0.2ex,
       fill opacity=0.4, text opacity=1}]
    \matrix [
       draw=none, matrix of nodes] (n)
    {"""
    tab = "    "
    box_counter = 0
    for location_code_tuple, refinement in zip(location_stack, descriptor):
        tikz_string += tab
        for dimension in range(descriptor.get_num_dimensions()):
            location_code_string = "\\texttt{" + location_code_tuple[dimension] + "}"
            if dimension < descriptor.get_num_dimensions() - 1:
                location_code_string = location_code_string + "&"
            location_code_string = location_code_string.replace(
                "∩", "\\scalebox{0.79}{$\\cap$}"
            )
            location_code_string = location_code_string.replace(
                "λ", "\\scalebox{0.85}{$\\lambda$}"
            )

            if refinement.count() == 0 and len(location_code_tuple[dimension]) > 0:
                tikz_string += f"|[fill=color_{box_counter}]| "
            tikz_string += location_code_string
        tikz_string += "\\\\ \n"
        if refinement.count() == 0:
            box_counter += 1
    tikz_string = tikz_string[:-2]
    tikz_string += "\n};\n"
    tikz_string += "\\draw "
    for column in range(descriptor.get_num_dimensions()):
        tikz_string += f"(n-1-{column+1}.north west) -- (n-{len(location_stack)}-{column+1}.south west)"
    tikz_string += ";"
    tikz_string += R"""
\end{tikzpicture}
\end{document}
"""
    if not filename.endswith(".tex"):
        filename += ".tex"
    latex_write_and_compile(tikz_string, filename)
