"""Reproducible examples for the first Stochastic Simulation/NASODE note."""

from __future__ import annotations

from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFont


SEED = 20240918


def lcg(n: int, *, seed: int = 0, a: int = 5, c: int = 1, m: int = 32) -> np.ndarray:
    """Return n values from a deliberately small full-period LCG."""
    state = seed
    values = np.empty(n)
    for i in range(n):
        state = (a * state + c) % m
        values[i] = state / m
    return values


def inverse_exponential(
    n: int, rate: float, rng: np.random.Generator
) -> np.ndarray:
    """Generate Exp(rate) samples by inverse transform."""
    u = rng.random(n)
    return -np.log1p(-u) / rate


def rejection_beta_2_5(
    n: int, rng: np.random.Generator
) -> tuple[np.ndarray, float]:
    """Generate Beta(2, 5) samples using Uniform(0, 1) proposals."""
    # f(x) = 30*x*(1-x)^4, maximized at x=1/5.
    envelope = 30.0 * (1.0 / 5.0) * (4.0 / 5.0) ** 4
    accepted: list[np.ndarray] = []
    accepted_count = 0
    proposed_count = 0

    while accepted_count < n:
        batch_size = max(1024, 2 * (n - accepted_count))
        y = rng.random(batch_size)
        u = rng.random(batch_size)
        density = 30.0 * y * (1.0 - y) ** 4
        batch = y[u <= density / envelope]
        accepted.append(batch)
        accepted_count += batch.size
        proposed_count += batch_size

    sample = np.concatenate(accepted)[:n]
    # The batched implementation slightly overshoots n, so report the fraction of
    # all generated proposals that were accepted in the completed batches.
    observed_rate = accepted_count / proposed_count
    return sample, observed_rate


def box_muller(n: int, rng: np.random.Generator) -> np.ndarray:
    """Generate independent standard normal samples with Box-Muller."""
    pairs = (n + 1) // 2
    # 1 - U is uniform on (0, 1], which keeps log(0) out of the formula.
    u1 = 1.0 - rng.random(pairs)
    u2 = rng.random(pairs)
    radius = np.sqrt(-2.0 * np.log(u1))
    angle = 2.0 * np.pi * u2
    z = np.column_stack((radius * np.cos(angle), radius * np.sin(angle)))
    return z.ravel()[:n]


def load_font(size: int, *, bold: bool = False) -> ImageFont.FreeTypeFont:
    """Load a font bundled with Pillow, with a minimal fallback."""
    name = "DejaVuSans-Bold.ttf" if bold else "DejaVuSans.ttf"
    try:
        return ImageFont.truetype(name, size)
    except OSError:
        return ImageFont.load_default()


def draw_panel_axes(
    draw: ImageDraw.ImageDraw,
    box: tuple[int, int, int, int],
    title: str,
    x_label: str,
    y_label: str,
) -> tuple[int, int, int, int]:
    """Draw one panel and return its plotting rectangle."""
    left, top, right, bottom = box
    plot = (left + 82, top + 58, right - 24, bottom - 62)
    x0, y0, x1, y1 = plot
    draw.text((left + 18, top + 12), title, font=load_font(25, bold=True), fill="#172033")
    draw.rectangle(plot, outline="#5b6475", width=2)
    for fraction in (0.25, 0.5, 0.75):
        x = x0 + fraction * (x1 - x0)
        y = y0 + fraction * (y1 - y0)
        draw.line((x, y0, x, y1), fill="#e2e6ee", width=1)
        draw.line((x0, y, x1, y), fill="#e2e6ee", width=1)
    draw.text(
        ((x0 + x1) / 2 - 12, bottom - 44),
        x_label,
        font=load_font(21),
        fill="#30394b",
    )
    draw.text((left + 15, (y0 + y1) / 2 - 10), y_label, font=load_font(18), fill="#30394b")
    return plot


def draw_scatter(
    draw: ImageDraw.ImageDraw,
    plot: tuple[int, int, int, int],
    x: np.ndarray,
    y: np.ndarray,
    *,
    radius: int,
    color: tuple[int, int, int, int],
) -> None:
    x0, y0, x1, y1 = plot
    px = x0 + np.clip(x, 0, 1) * (x1 - x0)
    py = y1 - np.clip(y, 0, 1) * (y1 - y0)
    for xx, yy in zip(px, py):
        draw.ellipse((xx - radius, yy - radius, xx + radius, yy + radius), fill=color)


def draw_histogram(
    draw: ImageDraw.ImageDraw,
    plot: tuple[int, int, int, int],
    sample: np.ndarray,
    domain: tuple[float, float],
    density,
) -> None:
    x0, y0, x1, y1 = plot
    counts, edges = np.histogram(sample, bins=54, range=domain, density=True)
    curve_x = np.linspace(domain[0], domain[1], 500)
    curve_y = density(curve_x)
    ymax = 1.08 * max(float(counts.max()), float(curve_y.max()))

    def sx(value):
        return x0 + (value - domain[0]) / (domain[1] - domain[0]) * (x1 - x0)

    def sy(value):
        return y1 - value / ymax * (y1 - y0)

    for count, start, end in zip(counts, edges[:-1], edges[1:]):
        draw.rectangle((sx(start), sy(count), sx(end), y1), fill=(103, 150, 230, 105))
    points = [(sx(xx), sy(yy)) for xx, yy in zip(curve_x, curve_y)]
    draw.line(points, fill="#c23b53", width=4, joint="curve")


def save_figure(
    lcg_values: np.ndarray,
    modern_values: np.ndarray,
    exp_sample: np.ndarray,
    beta_sample: np.ndarray,
    figure_path: Path,
) -> None:
    """Create a dependency-light four-panel PNG with Pillow."""
    image = Image.new("RGBA", (1600, 1160), "white")
    draw = ImageDraw.Draw(image, "RGBA")
    boxes = (
        (35, 30, 790, 555),
        (810, 30, 1565, 555),
        (35, 585, 790, 1125),
        (810, 585, 1565, 1125),
    )

    plot = draw_panel_axes(draw, boxes[0], "Small full-period LCG", "u[n]", "u[n+1]")
    draw_scatter(
        draw, plot, lcg_values[:-1], lcg_values[1:], radius=5, color=(34, 98, 190, 220)
    )

    plot = draw_panel_axes(
        draw, boxes[1], "NumPy default generator", "u[n]", "u[n+1]"
    )
    draw_scatter(
        draw,
        plot,
        modern_values[:-1],
        modern_values[1:],
        radius=2,
        color=(34, 98, 190, 70),
    )

    plot = draw_panel_axes(draw, boxes[2], "Inverse transform: Exp(2)", "x", "density")
    draw_histogram(draw, plot, exp_sample, (0.0, 4.0), lambda x: 2.0 * np.exp(-2.0 * x))

    plot = draw_panel_axes(
        draw, boxes[3], "Rejection sampling: Beta(2, 5)", "x", "density"
    )
    draw_histogram(
        draw, plot, beta_sample, (0.0, 1.0), lambda x: 30.0 * x * (1.0 - x) ** 4
    )

    figure_path.parent.mkdir(parents=True, exist_ok=True)
    image.convert("RGB").save(figure_path, quality=94)


def main() -> None:
    rng = np.random.default_rng(SEED)
    n = 50_000

    exp_sample = inverse_exponential(n, rate=2.0, rng=rng)
    beta_sample, beta_acceptance = rejection_beta_2_5(n, rng=rng)
    normal_sample = box_muller(n, rng=rng)

    covariance = np.array([[1.0, 0.8], [0.8, 2.0]])
    factor = np.linalg.cholesky(covariance)
    standard = box_muller(2 * n, rng=rng).reshape(n, 2)
    correlated = standard @ factor.T

    print(f"Exp(2): mean={exp_sample.mean():.6f}, var={exp_sample.var():.6f}")
    print(
        "Beta(2,5): "
        f"mean={beta_sample.mean():.6f}, var={beta_sample.var():.6f}, "
        f"acceptance={beta_acceptance:.6f}"
    )
    print(
        "Box-Muller N(0,1): "
        f"mean={normal_sample.mean():.6f}, var={normal_sample.var():.6f}"
    )
    print("Correlated normal empirical covariance:")
    print(np.cov(correlated, rowvar=False, ddof=0))

    lcg_values = lcg(33)
    modern_values = rng.random(2_000)

    output = Path(__file__).resolve().parents[1] / "media" / "stochsim-nasode"
    figure_path = output / "01-random-number-generation-and-sampling.png"
    save_figure(lcg_values, modern_values, exp_sample, beta_sample, figure_path)
    print(f"Figure written to {figure_path}")


if __name__ == "__main__":
    main()