"""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.patches import Patch
from matplotlib.path import Path as PlotPath
from matplotlib.patheffects import AbstractPathEffect
from matplotlib.transforms import Affine2D, 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]
GRID_COLOR = "#E6E9EC"
LINE_COLOR = "#C9CED3"
BAND_COLOR = "#F7F9FA"
COMPILER_STYLES = {"raptor": PALETTE["sage"], "pimcomp": PALETTE["sand"]}
METRIC_STYLES = {"latency": PALETTE["blue"], "throughput": 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"}


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 "./\\"))
        renderer.draw_path(base_gc, tpath, affine, rgbFace)
        base_gc.restore()

        bounds = affine.transform_path(tpath).get_extents()
        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
            half_length = (abs(dx) * bounds.width + abs(dy) * 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
                vertices.extend(((x - dx * half_length, y - dy * half_length),
                                 (x + dx * half_length, y + dy * half_length)))
        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))
            centers = [(center_x + x * spacing, center_y + y * spacing)
                       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_markers(line_gc, PlotPath.unit_circle(),
                                  Affine2D().scale(renderer.points_to_pixels(0.9)),
                                  PlotPath(centers), IdentityTransform(), gc.get_hatch_color())
        line_gc.restore()


def configure_style() -> None:
    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",
    })


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],
                            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) -> 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))
    upper = max(0, *(upper for _, upper in limits))
    if lower == upper:
        upper = lower + 1
    for _ in range(12):
        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 save_figure(fig, basename: str, output_dir: Path) -> None:
    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)