"""Generate deterministic light/dark figures for the Learning fractal lessons.

Inside this repository the assets are written to their canonical public path.
A downloaded copy instead writes to ``learning-figures/dinamica-caos`` below
the current working directory, so it remains executable as a standalone file.
"""

from __future__ import annotations

from itertools import product
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


SCRIPT_ROOT = Path(__file__).resolve().parents[1]
IN_REPOSITORY = (SCRIPT_ROOT / "src" / "content" / "learning").is_dir()
OUT = (
    SCRIPT_ROOT / "public" / "images" / "learning" / "dinamica-caos"
    if IN_REPOSITORY
    else Path.cwd() / "learning-figures" / "dinamica-caos"
)


def palette(dark: bool) -> dict[str, str]:
    return {
        "bg": "#07111f" if dark else "#f8fafc",
        "panel": "#0c1a2e" if dark else "#ffffff",
        "text": "#e5edf7" if dark else "#172033",
        "muted": "#9fb0c5" if dark else "#526174",
        "grid": "#30445f" if dark else "#d7e0ea",
        "rose": "#fb7185" if dark else "#be123c",
        "cyan": "#67e8f9" if dark else "#0e7490",
        "amber": "#fbbf24" if dark else "#b45309",
    }


def destination(name: str, dark: bool) -> Path:
    folder = OUT if dark else OUT / "light"
    folder.mkdir(parents=True, exist_ok=True)
    return folder / name


def style_axes(ax: plt.Axes, colors: dict[str, str]) -> None:
    ax.set_facecolor(colors["panel"])
    ax.tick_params(colors=colors["muted"], labelsize=8)
    for spine in ax.spines.values():
        spine.set_color(colors["grid"])
    ax.xaxis.label.set_color(colors["muted"])
    ax.yaxis.label.set_color(colors["muted"])


def cantor_intervals(depth: int) -> list[tuple[float, float]]:
    """Return the closed intervals retained after ``depth`` Cantor steps."""
    intervals = [(0.0, 1.0)]
    for _ in range(depth):
        next_intervals = []
        for left, right in intervals:
            third = (right - left) / 3.0
            next_intervals.extend([(left, left + third), (right - third, right)])
        intervals = next_intervals
    return intervals


def make_ifs_julia_mandelbrot_svg(dark: bool) -> None:
    """Generate the three-panel IFS, Julia and Mandelbrot comparison."""
    colors = palette(dark)
    plt.rcParams["svg.hashsalt"] = "fyskode-learning-fractals"
    figure, axes = plt.subplots(1, 3, figsize=(12.0, 5.0), constrained_layout=True)
    figure.patch.set_facecolor(colors["bg"])

    left, center, right = axes
    style_axes(left, colors)
    for depth in range(5):
        for start, end in cantor_intervals(depth):
            left.plot([start, end], [-depth, -depth], color=colors["cyan"], lw=7.0, solid_capstyle="butt")
    left.set(xlabel="posición", ylabel="iteración", xlim=(-0.04, 1.04), ylim=(-4.55, 0.55))
    left.set_yticks([0, -1, -2, -3, -4], ["0", "1", "2", "3", "4"])
    left.grid(True, axis="x", alpha=0.22)

    julia = escape_counts((-1.65, 1.65), (-1.35, 1.35), (420, 480), 80, -0.8 + 0.156j)
    mandelbrot = escape_counts((-2.1, 0.8), (-1.3, 1.3), (420, 480), 80)
    for ax, values, extent, x_label in [
        (center, julia, (-1.65, 1.65, -1.35, 1.35), r"$\operatorname{Re}(z_0)$"),
        (right, mandelbrot, (-2.1, 0.8, -1.3, 1.3), r"$\operatorname{Re}(c)$"),
    ]:
        style_axes(ax, colors)
        ax.imshow(
            values,
            extent=extent,
            origin="lower",
            cmap="magma",
            interpolation="bilinear",
            vmin=0,
            vmax=55,
            aspect="equal",
        )
        ax.set_xlabel(x_label)
        ax.set_ylabel(r"parte imaginaria")

    svg_path = destination("26_ifs_julia_mandelbrot.svg", dark)
    figure.savefig(
        svg_path,
        format="svg",
        facecolor=colors["bg"],
        metadata={"Date": None},
    )
    plt.close(figure)
    svg_text = svg_path.read_text(encoding="utf-8")
    svg_path.write_text(
        "\n".join(line.rstrip() for line in svg_text.splitlines()) + "\n",
        encoding="utf-8",
    )


def escape_counts(
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    shape: tuple[int, int],
    max_iter: int,
    julia_c: complex | None = None,
) -> np.ndarray:
    xs = np.linspace(*xlim, shape[1])
    ys = np.linspace(*ylim, shape[0])
    plane = xs[None, :] + 1j * ys[:, None]
    z = np.zeros_like(plane) if julia_c is None else plane.copy()
    c = plane if julia_c is None else np.full_like(plane, julia_c)
    escaped = np.zeros(shape, dtype=bool)
    counts = np.full(shape, max_iter, dtype=float)
    for iteration in range(max_iter):
        z[~escaped] = z[~escaped] * z[~escaped] + c[~escaped]
        newly = (~escaped) & (np.abs(z) > 2.0)
        if np.any(newly):
            counts[newly] = iteration + 1 - np.log2(np.log2(np.abs(z[newly])))
        escaped |= newly
    return counts


def make_escape_figure(dark: bool) -> None:
    colors = palette(dark)
    figure, axes = plt.subplots(1, 2, figsize=(11.2, 5.25), constrained_layout=True)
    figure.patch.set_facecolor(colors["bg"])

    datasets = [
        (
            escape_counts((-2.1, 0.8), (-1.3, 1.3), (520, 580), 90),
            (-2.1, 0.8, -1.3, 1.3),
        ),
        (
            escape_counts((-1.65, 1.65), (-1.35, 1.35), (520, 580), 90, -0.8 + 0.156j),
            (-1.65, 1.65, -1.35, 1.35),
        ),
    ]
    image = None
    for ax, (values, extent) in zip(axes, datasets):
        style_axes(ax, colors)
        image = ax.imshow(
            values,
            extent=extent,
            origin="lower",
            cmap="magma",
            interpolation="bilinear",
            vmin=0,
            vmax=55,
            aspect="equal",
        )
        ax.set_xlabel(r"parte real")
        ax.set_ylabel(r"parte imaginaria")

    colorbar = figure.colorbar(image, ax=axes, shrink=0.83, pad=0.025)
    colorbar.set_label("iteraciones antes de |z| > 2", color=colors["muted"], fontsize=8)
    colorbar.ax.tick_params(colors=colors["muted"], labelsize=7)
    colorbar.outline.set_edgecolor(colors["grid"])
    figure.savefig(destination("27_mandelbrot_julia_escape.png", dark), dpi=170, facecolor=colors["bg"])
    plt.close(figure)


def sierpinski_iterate(points: np.ndarray, rounds: int) -> np.ndarray:
    vertices = np.array([[0.0, 0.0], [1.0, 0.0], [0.5, np.sqrt(3.0) / 2.0]])
    current = points
    for _ in range(rounds):
        current = np.concatenate([(current + vertex) / 2.0 for vertex in vertices])
    return current


def make_ifs_figure(dark: bool) -> None:
    colors = palette(dark)
    t = np.linspace(0.0, 1.0, 90)
    boundary = np.vstack(
        [
            np.column_stack((t, np.zeros_like(t))),
            np.column_stack((np.ones_like(t), t)),
            np.column_stack((t[::-1], np.ones_like(t))),
            np.column_stack((np.zeros_like(t), t[::-1])),
        ]
    )
    rounds = [0, 1, 2, 5]
    figure, axes = plt.subplots(1, 4, figsize=(11.2, 3.55), constrained_layout=True)
    figure.patch.set_facecolor(colors["bg"])
    for ax, count in zip(axes, rounds):
        style_axes(ax, colors)
        points = sierpinski_iterate(boundary, count)
        ax.scatter(points[:, 0], points[:, 1], s=1.2 if count < 4 else 0.35, c=colors["cyan"], alpha=0.82)
        ax.set_xlim(-0.05, 1.05)
        ax.set_ylim(-0.05, 1.02)
        ax.set_aspect("equal")
        ax.set_xticks([])
        ax.set_yticks([])
    figure.savefig(destination("28_ifs_compact_attractor.png", dark), dpi=170, facecolor=colors["bg"])
    plt.close(figure)


def cantor_points(depth: int = 8) -> tuple[np.ndarray, np.ndarray]:
    addresses = np.array(list(product([0, 1], repeat=depth)), dtype=int)
    weights = 2.0 / (3.0 ** np.arange(1, depth + 1))
    x = addresses @ weights
    return x, addresses


def make_fractal_homeomorphism_figure(dark: bool) -> None:
    colors = palette(dark)
    x, addresses = cantor_points()
    code = addresses[:, :3] @ np.array([4, 2, 1])
    y_warped = 0.24 * np.sin(2 * np.pi * x)
    figure, axes = plt.subplots(1, 2, figsize=(11.2, 4.0), constrained_layout=True)
    figure.patch.set_facecolor(colors["bg"])
    cmap = plt.get_cmap("viridis", 8)

    for ax in axes:
        style_axes(ax, colors)
        ax.set_xlim(-0.04, 1.04)
        ax.set_ylim(-0.38, 0.38)
        ax.set_yticks([])
        ax.set_xlabel("posición codificada por la dirección")
    axes[0].scatter(x, np.zeros_like(x), c=code, cmap=cmap, s=9, alpha=0.9)
    axes[1].scatter(x, y_warped, c=code, cmap=cmap, s=9, alpha=0.9)
    axes[1].plot(np.linspace(0, 1, 400), 0.24 * np.sin(2 * np.pi * np.linspace(0, 1, 400)), color=colors["grid"], lw=0.7, alpha=0.55)

    for index in [0, 36, 109, 182, 255]:
        label = "".join(str(bit) for bit in addresses[index, :5]) + "…"
        axes[0].annotate(label, (x[index], 0), xytext=(0, 11), textcoords="offset points", ha="center", fontsize=6.5, color=colors["amber"])
        axes[1].annotate(label, (x[index], y_warped[index]), xytext=(0, 11), textcoords="offset points", ha="center", fontsize=6.5, color=colors["amber"])

    figure.savefig(destination("29_fractal_homeomorphism_codes.png", dark), dpi=170, facecolor=colors["bg"])
    plt.close(figure)


def main() -> None:
    for dark in (True, False):
        make_ifs_julia_mandelbrot_svg(dark)
        make_escape_figure(dark)
        make_ifs_figure(dark)
        make_fractal_homeomorphism_figure(dark)
    print("Generated 8 Learning fractal figure variants.")


if __name__ == "__main__":
    main()