"""Shared appearance and PNG/SVG/PDF export for PIM result plots."""

from pathlib import Path
from math import ceil

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import to_rgba
from matplotlib.legend_handler import HandlerBase
from matplotlib.lines import Line2D
from matplotlib.patches import Patch
from matplotlib.path import Path as PlotPath
from matplotlib.patheffects import AbstractPathEffect
from matplotlib.text import Text
from matplotlib.transforms import Affine2D, Bbox, IdentityTransform, TransformedPath


PALETTE = {
    "sage": ("#B3D5CA", "#4B7C6A"),
    "sand": ("#E1BC8B", "#856A44"),
    "pale_blue": ("#EDF3FA", "#456B90"),
    "blue": ("#B3CCE3", "#456B90"),
    "pale_green": ("#EAF5F1", "#3E8271"),
    "pale_sand": ("#FCF1E5", "#A87940"),
    "red": ("#FF9494", "#C52020"),
    "purple": ("#D1C0E0", "#77558F"),
}
TEXT_COLOR = "#222222"
MATRIX_TEXT_COLOR = "#1D2530"
ERROR_COLOR = PALETTE["red"][1]
DAGGER_COLOR = "#AA1C1C"
GRID_COLOR = "#E6E9EC"
LINE_COLOR = "#C9CED3"
BAND_COLOR = "#F7F9FA"
COMPILER_STYLES = {"raptor": PALETTE["sage"], "pimcomp": PALETTE["sand"]}
ARCH_STYLES = {"arch-a": (*PALETTE["blue"], ""), "arch-b": (*PALETTE["sage"], "//")}
NETWORK_STYLES = {"vgg8": PALETTE["blue"], "resnet18": PALETTE["sage"],
                  "resnet34": PALETTE["sand"], "googlenet": PALETTE["purple"],
                  "yolo11n": PALETTE["red"]}
COMPILER_NAMES = {"raptor": "Raptor", "pimcomp": "PIMCOMP"}
ARCH_NAMES = {"arch-a": "Arch-A", "arch-b": "Arch-B"}
MODEL_NAMES = {"vgg8": "VGG-8", "resnet18": "ResNet-18", "resnet34": "ResNet-34",
               "googlenet": "GoogLeNet", "yolo11n": "YOLO11n"}
MODE_NAMES = {"latency": "Single-inference", "throughput": "Throughput"}
COMPACT_MODE_NAMES = {"latency": "Latency-mode", "throughput": "Throughput-mode"}
_USE_COLORS = True


class CenteredHatch(AbstractPathEffect):
    """Anchor dot and diagonal hatch patterns to each rectangle's center."""

    def draw_path(self, renderer, gc, tpath, affine, rgbFace=None):
        hatch = gc.get_hatch() or ""
        base_gc = renderer.new_gc()
        base_gc.copy_properties(gc)
        base_gc.set_hatch("".join(symbol for symbol in hatch if symbol not in "./\\") or None)
        renderer.draw_path(base_gc, tpath, affine, rgbFace)
        base_gc.restore()

        bounds = affine.transform_path(tpath).get_extents()
        clip_bounds = bounds if gc.get_clip_rectangle() is None else Bbox.intersection(bounds, gc.get_clip_rectangle())
        if clip_bounds is None:
            return
        center_x, center_y = (bounds.x0 + bounds.x1) / 2, (bounds.y0 + bounds.y1) / 2
        vertices = []
        diagonal = 2 ** -0.5
        for symbols, dx, dy in (("/", diagonal, diagonal), ("\\", diagonal, -diagonal)):
            count = sum(hatch.count(symbol) for symbol in symbols)
            if count == 0:
                continue
            spacing = renderer.points_to_pixels(12 / count)
            half_span = (abs(dy) * bounds.width + abs(dx) * bounds.height) / 2
            extent = ceil(half_span / spacing)
            for step in range(-extent, extent + 1):
                x = center_x - dy * step * spacing
                y = center_y + dx * step * spacing
                y_limits = sorted(((clip_bounds.y0 - y) / dy, (clip_bounds.y1 - y) / dy))
                start = max((clip_bounds.x0 - x) / dx, y_limits[0])
                end = min((clip_bounds.x1 - x) / dx, y_limits[1])
                if start < end:
                    vertices.extend(((x + dx * start, y + dy * start),
                                     (x + dx * end, y + dy * end)))
        line_gc = renderer.new_gc()
        line_gc.copy_properties(gc)
        line_gc.set_hatch(None)
        line_gc.set_foreground(gc.get_hatch_color())
        line_gc.set_linewidth(gc.get_hatch_linewidth())
        line_gc.set_clip_path(TransformedPath(tpath, affine))
        if vertices:
            renderer.draw_path(line_gc, PlotPath(vertices, [PlotPath.MOVETO, PlotPath.LINETO] * (len(vertices) // 2)),
                               IdentityTransform())
        if "." in hatch:
            spacing = renderer.points_to_pixels(12 / hatch.count("."))
            x_extent = ceil(bounds.width / (2 * spacing))
            y_extent = ceil(bounds.height / (2 * spacing))
            # Trim geometry too: SVG importers can discard clip paths.
            dots = [PlotPath.unit_circle().transformed(
                        Affine2D().scale(renderer.points_to_pixels(0.9)).translate(
                            center_x + x * spacing, center_y + y * spacing)).clip_to_bbox(clip_bounds)
                    for x in range(-x_extent, x_extent + 1)
                    for y in range(-y_extent, y_extent + 1)]
            line_gc.set_linewidth(0)
            renderer.draw_path(line_gc, PlotPath.make_compound_path(*dots),
                               IdentityTransform(), gc.get_hatch_color())
        line_gc.restore()


def configure_style(*, colors: bool = True, anonymize: bool = False) -> None:
    global _USE_COLORS
    _USE_COLORS = colors
    COMPILER_NAMES["raptor"] = "Anonimized" if anonymize else "Raptor"
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.labelcolor": TEXT_COLOR,
        "text.color": TEXT_COLOR,
        "xtick.color": TEXT_COLOR,
        "ytick.color": TEXT_COLOR,
        "axes.edgecolor": LINE_COLOR,
        "grid.color": GRID_COLOR,
        "grid.linewidth": 0.8,
        "hatch.color": TEXT_COLOR,
        "savefig.facecolor": "white",
        "figure.facecolor": "white",
    })


class TextLegendHandler(HandlerBase):
    def create_artists(self, legend, original, xdescent, ydescent, width, height, fontsize, transform):
        return [Text(width / 2 - xdescent, height / 2 - ydescent, original.get_text(),
                     color=original.get_color(), fontproperties=original.get_fontproperties(),
                     math_fontfamily=original.get_math_fontfamily(), ha="center", va="center",
                     transform=transform)]


def add_legend(fig, handles) -> None:
    """Use the fewest balanced legend rows that fit the figure width."""
    renderer = fig.canvas.get_renderer()
    for rows in range(1, len(handles) + 1):
        columns = ceil(len(handles) / rows)
        grid = []
        start = 0
        for row in range(rows):
            count = len(handles) // rows + (row < len(handles) % rows)
            grid.append(handles[start:start + count] + [None] * (columns - count))
            start += count
        ordered = [handle for column in zip(*grid) for handle in column]
        legend = fig.legend(handles=[handle if handle is not None else Patch(visible=False) for handle in ordered],
                            labels=[handle.get_label() if handle is not None else "" for handle in ordered],
                            handler_map={Text: TextLegendHandler()},
                            loc="outside upper center", ncols=columns, frameon=False, fontsize=9.5,
                            handlelength=1.7, columnspacing=1.5, handletextpad=0.6, labelspacing=0.4)
        for handle, patch, text in zip(ordered, legend.legend_handles, legend.get_texts()):
            if handle is None:
                patch.set_visible(False)
                text.set_visible(False)
        if legend.get_window_extent(renderer).width <= fig.bbox.width - 20:
            return
        legend.remove()


def value_label(ax, x, y, label: str, status: str = "PASS", direction: int = 1,
                fontsize: float = 8, invalid_label: str = "INVALID") -> None:
    alignment = "bottom" if direction > 0 else "top"
    rotation = 90 if direction > 0 else -90
    offset = direction * fontsize * 0.65 if status != "PASS" else 0
    ax.annotate(label, (x, y), xytext=(-offset, 0), textcoords="offset points",
                rotation=rotation, ha="center", va=alignment, fontsize=fontsize,
                color=TEXT_COLOR)
    if status != "PASS":
        ax.annotate(invalid_label if status == "FAIL" else status, (x, y),
                    xytext=(offset, 0), textcoords="offset points",
                    rotation=rotation, ha="center", va=alignment, fontsize=fontsize,
                    color=ERROR_COLOR if status == "FAIL" else TEXT_COLOR)


def fit_x_tick_labels(fig) -> None:
    """Rotate only overlapping labels, clockwise by 30 degrees."""
    fig.canvas.draw()
    renderer = fig.canvas.get_renderer()
    for minor in (False, True):
        labels = [(text, text.get_window_extent(renderer))
                  for ax in fig.axes for text in ax.get_xticklabels(minor=minor)
                  if text.get_visible() and text.get_text()]
        crowded = set()
        for index, (text, bounds) in enumerate(labels):
            for other_text, other_bounds in labels[:index]:
                if bounds.overlaps(other_bounds):
                    crowded.update((text, other_text))
        for text in crowded:
            text.set_rotation(-30)
            text.set_rotation_mode("anchor")
            text.set_ha("right" if text.get_va() == "bottom" else "left")
    fig.canvas.draw()
    renderer = fig.canvas.get_renderer()
    for ax in fig.axes:
        bottom_minor = [text for text in ax.get_xticklabels(minor=True)
                        if text.get_va() == "top" and text.get_visible() and text.get_text()]
        if bottom_minor:
            height = max(text.get_window_extent(renderer).height for text in bottom_minor)
            ax.tick_params(axis="x", which="major", pad=height * 72 / fig.dpi + 6)
    fig.canvas.draw()


def fit_bar_widths(fig) -> None:
    """Reserve enough horizontal space for each bar's rendered annotations."""
    fig.canvas.draw()
    renderer = fig.canvas.get_renderer()
    scale = 1.0
    for ax in fig.axes:
        labels = {}
        for text in ax.texts:
            if hasattr(text, "xy"):
                labels.setdefault(text.xy[0], []).append(text.get_window_extent(renderer))
        for container in ax.containers:
            for bar in container.patches:
                bounds = labels.get(bar.get_x() + bar.get_width() / 2, [])
                if bounds:
                    width = max(box.x1 for box in bounds) - min(box.x0 for box in bounds)
                    scale = max(scale, (width + renderer.points_to_pixels(2)) / bar.get_window_extent(renderer).width)
    if scale > 1:
        fig.set_size_inches(fig.get_figwidth() * scale, fig.get_figheight())


def fit_value_labels(fig, *, lower_limit: float | None = None) -> None:
    """Fit common vertical limits to the data and rendered value labels."""
    for ax in fig.axes:
        for text in ax.texts:
            text.set_in_layout(False)
    fig.canvas.draw()
    renderer = fig.canvas.get_renderer()
    limits = []
    labels = []
    for ax in fig.axes:
        bounds = ax.get_window_extent(renderer)
        for text in ax.texts:
            x, y = text.xy if hasattr(text, "xy") else text.get_position()
            anchor = ax.transData.transform((x, y))[1]
            extent = text.get_window_extent(renderer)
            labels.append((y, max(0, anchor - extent.y0 + 3) / bounds.height,
                           max(0, extent.y1 - anchor + 3) / bounds.height))
        limits.append(ax.dataLim.intervaly if ax.has_data() else ax.get_ylim())
    lower = min(0, *(lower for lower, _ in limits)) if lower_limit is None else lower_limit
    upper = max(0, *(upper for _, upper in limits))
    if lower == upper:
        upper = lower + 1
    for _ in range(12):
        if lower_limit is None:
            lower = min([lower, *((y - below * upper) / (1 - below) for y, below, _ in labels)])
        upper = max([upper, *((y - above * lower) / (1 - above) for y, _, above in labels)])
    for ax in fig.axes:
        ax.set_ylim(lower, upper)


def grayscale_figure(fig) -> None:
    def gray(color):
        red, green, blue, alpha = to_rgba(color)
        level = 0.2126 * red + 0.7152 * green + 0.0722 * blue
        return level, level, level, alpha

    for artist in fig.findobj():
        if isinstance(artist, Patch):
            artist.set_facecolor(gray(artist.get_facecolor()))
            artist.set_edgecolor(gray(artist.get_edgecolor()))
        elif isinstance(artist, Line2D):
            artist.set_color(gray(artist.get_color()))
            artist.set_markerfacecolor(gray(artist.get_markerfacecolor()))
            artist.set_markeredgecolor(gray(artist.get_markeredgecolor()))
        elif isinstance(artist, Text):
            artist.set_color(gray(artist.get_color()))


def save_figure(fig, basename: str, output_dir: Path) -> None:
    if not _USE_COLORS:
        grayscale_figure(fig)
    for patch in fig.findobj(Patch):
        hatch = patch.get_hatch() or ""
        if any(symbol in hatch for symbol in "./\\"):
            patch.set_path_effects([CenteredHatch()])
    for extension in ("png", "svg", "pdf"):
        target = output_dir / extension / f"{basename}.{extension}"
        target.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(target, dpi=220, bbox_inches="tight", pad_inches=0.04)
        print(target)
    plt.close(fig)