"""fractal_fair_helpers.py

Small rendering helpers shared by the Fractal Fair examples.

The chapter files deliberately keep *mathematics* and *construction rules*
local, while this module hides repetitive Matplotlib plumbing: canvases,
axis cleanup, gradient paths, point clouds, images, and SVG export.

That division is pedagogical, not architectural theatre: after importing
these helpers, a chapter script should read almost like executable
pseudocode for the fractal itself.
"""
from __future__ import annotations

from pathlib import Path
import matplotlib
matplotlib.use("SVG")
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import LineCollection, PatchCollection
from matplotlib.patches import Polygon

FIGURES_DIR = Path(__file__).resolve().parent.parent / "figures"
FIGURES_DIR.mkdir(parents=True, exist_ok=True)

PALETTE_FIRE = mcolors.LinearSegmentedColormap.from_list(
    "fair_fire", ["#0b0033", "#3b0f70", "#8c2981", "#de4968", "#fe9f6d", "#fcfdbf"]
)
PALETTE_OCEAN = mcolors.LinearSegmentedColormap.from_list(
    "fair_ocean", ["#03045e", "#0077b6", "#00b4d8", "#90e0ef", "#caf0f8"]
)
PALETTE_PLANT = mcolors.LinearSegmentedColormap.from_list(
    "fair_plant", ["#6b4423", "#2e7d32", "#8bc34a"]
)


def save_svg(fig, name: str, **kwargs):
    """Save a Matplotlib figure using Fractal Fair's SVG defaults."""
    path = FIGURES_DIR / f"{name}.svg"
    options = dict(format="svg", bbox_inches="tight", pad_inches=0.05,
                   transparent=True, dpi=220)
    options.update(kwargs)
    fig.savefig(path, **options)
    plt.close(fig)
    print(f"wrote {path}")
    return path


def clean(ax, *, equal=True):
    """Remove chart furniture; optionally keep one unit equal in x and y."""
    if equal:
        ax.set_aspect("equal")
    ax.axis("off")
    return ax


def fit(ax, points, pad=0.05):
    """Fit axes around Nx2 points or Mx2x2 line segments."""
    p = np.asarray(points)
    xs, ys = p[..., 0], p[..., 1]
    dx, dy = np.ptp(xs), np.ptp(ys)
    ax.set_xlim(xs.min() - pad * max(dx, 1), xs.max() + pad * max(dx, 1))
    ax.set_ylim(ys.min() - pad * max(dy, 1), ys.max() + pad * max(dy, 1))
    return ax


def path_figure(points, *, palette=PALETTE_OCEAN, figsize=(7.5, 7.0),
                linewidth=1.5, pad=0.05, fill=None, rasterized=False):
    """Draw a polyline with a traversal-order gradient."""
    pts = np.asarray(points)
    segs = np.stack([pts[:-1], pts[1:]], axis=1)
    fig, ax = plt.subplots(figsize=figsize)
    colours = palette(np.linspace(0, 1, len(segs)))
    lines = LineCollection(segs, colors=colours, linewidths=linewidth)
    lines.set_rasterized(rasterized)
    ax.add_collection(lines)
    if fill is not None:
        colour, alpha = fill
        ax.fill(pts[:, 0], pts[:, 1], color=colour, alpha=alpha)
    fit(ax, pts, pad=pad)
    clean(ax)
    return fig


def segments_figure(segments, *, palette=PALETTE_OCEAN, figsize=(7.5, 7.5),
                    linewidth=1.2, pad=0.04, values=None, rasterized=False):
    """Draw independent segments with colours supplied by a scalar value."""
    segs = np.asarray(segments)
    if values is None:
        values = np.linspace(0, 1, len(segs))
    fig, ax = plt.subplots(figsize=figsize)
    lines = LineCollection(segs, colors=palette(np.asarray(values)), linewidths=linewidth)
    lines.set_rasterized(rasterized)
    ax.add_collection(lines)
    fit(ax, segs, pad=pad)
    clean(ax)
    return fig


def points_figure(points, *, figsize=(6.5, 9.0), values=None, cmap="Greens",
                  size=0.25, rasterized=True):
    """Draw a dense point cloud without exposing scatter boilerplate."""
    pts = np.asarray(points)
    fig, ax = plt.subplots(figsize=figsize)
    values = pts[:, 1] if values is None else values
    dots = ax.scatter(pts[:, 0], pts[:, 1], s=size, c=values,
                      cmap=cmap, linewidths=0)
    dots.set_rasterized(rasterized)
    clean(ax)
    return fig


def image_figure(image, *, extent, figsize=(8, 7), cmap=None, transform=None):
    """Display a computed scalar/RGB field as a clean mathematical image."""
    data = transform(image) if transform else image
    fig, ax = plt.subplots(figsize=figsize)
    ax.imshow(data, cmap=cmap, extent=extent, origin="lower")
    clean(ax, equal=False)
    return fig


def triangles_figure(triangles, *, palette=PALETTE_FIRE, figsize=(8, 7)):
    """Render surviving triangles, coloured by vertical position."""
    tris = np.asarray(triangles)
    heights = tris[:, :, 1].mean(axis=1)
    values = heights / max(heights.max(), 1e-12)
    patches = [Polygon(t, closed=True) for t in tris]
    fig, ax = plt.subplots(figsize=figsize)
    collection = PatchCollection(patches, facecolor=palette(values),
                                 edgecolor=palette(values), linewidths=0.15)
    ax.add_collection(collection)
    fit(ax, tris, pad=0.02)
    clean(ax)
    return fig


def staircase_figure(bars, depth, *, palette=PALETTE_OCEAN, figsize=(9, 3.6)):
    """Render (a,b,remaining_depth) triples as a Cantor-style staircase."""
    fig, ax = plt.subplots(figsize=figsize)
    for a, b, remaining in bars:
        row = depth - remaining
        ax.plot([a, b], [row, row], lw=4, solid_capstyle="butt",
                color=palette(row / max(depth, 1)))
    ax.set_xlim(-0.02, 1.02)
    ax.set_ylim(-0.7, depth + 0.5)
    ax.invert_yaxis()
    clean(ax, equal=False)
    return fig