#!/usr/bin/env python3
# /// script
# requires-python = "==3.12.*"
# dependencies = ["matplotlib"]
# ///
"""
Witness/liar table for n=561 (smallest Carmichael number = 3 × 11 × 17).

Shows bases a=2..20 against Fermat and Miller-Rabin tests.
- Green: witness (correctly identifies 561 as composite)
- Red:   liar (incorrectly outputs "probably prime")

Carmichael property: every base coprime to 561 is a Fermat liar.
Miller-Rabin: computes the squaring chain for 561-1 = 560 = 2^4 × 35.
"""

import math
import sys
from pathlib import Path

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from figure_quality import configure_matplotlib, figure_output_path, validate_figure

configure_matplotlib(font_family="DejaVu Serif")

N = 561  # 3 × 11 × 17


# ---------------------------------------------------------------------------
# Fermat test: a^(n-1) ≡ 1 (mod n)?  liar if True (and gcd(a,n)==1)
# ---------------------------------------------------------------------------
def fermat_liar(a, n):
    """Return True if a is a Fermat liar for n (i.e., Fermat test says 'probably prime')."""
    if math.gcd(a, n) > 1:
        return False  # composite witness via gcd
    return pow(a, n - 1, n) == 1


# ---------------------------------------------------------------------------
# Miller-Rabin: write n-1 = 2^s * d, then check squaring chain
# liar if: a^d ≡ 1 (mod n)  OR  a^(2^r * d) ≡ -1 (mod n) for some r < s
# ---------------------------------------------------------------------------
def miller_rabin_liar(a, n):
    """Return True if a is a Miller-Rabin liar for n."""
    if math.gcd(a, n) > 1:
        return False  # composite witness via gcd
    # factor out powers of 2 from n-1
    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1
    # now n-1 = 2^s * d
    x = pow(a, d, n)
    if x == 1 or x == n - 1:
        return True  # liar
    for _ in range(s - 1):
        x = pow(x, 2, n)
        if x == n - 1:
            return True  # liar
    return False  # witness


# ---------------------------------------------------------------------------
# Build data
# ---------------------------------------------------------------------------
BASES = list(range(2, 21))  # 2 .. 20

fermat_results = [fermat_liar(a, N) for a in BASES]  # True = liar
mr_results = [miller_rabin_liar(a, N) for a in BASES]  # True = liar

# ---------------------------------------------------------------------------
# Plot
# ---------------------------------------------------------------------------
LIAR_COLOR = "#FEE2E2"
WITNESS_COLOR = "#D1FAE5"
TEXT_CELL = "#111827"
TEXT_DARK = "#222222"

CELL_W = 1.6  # width per column
CELL_H = 0.52  # height per row
LEFT_MARGIN = 0.7  # width for "a =" labels
TOP_MARGIN = 0.55
BOTTOM_MARGIN = 0.25

N_ROWS = len(BASES)
N_COLS = 2

fig_w_in = (
    (LEFT_MARGIN + N_COLS * CELL_W + 0.3)
    * (550 / (96 * (LEFT_MARGIN + N_COLS * CELL_W + 0.3)))
    * (550 / 96)
)
# Simpler: fix at 550px / 96 dpi in points -> inches
TARGET_PX = 550
DPI = 300
fig_w_inch = TARGET_PX / 96  # ~5.73 in at screen scale; we'll scale by DPI ratio
# Actually: produce at 300 DPI, keep pixel width ~550*300/96 ≈ 1718px
# Just set a sensible inch size
fig_w_inch = 5.5
fig_h_inch = TOP_MARGIN + N_ROWS * CELL_H + BOTTOM_MARGIN

fig, ax = plt.subplots(figsize=(fig_w_inch, fig_h_inch))
ax.set_xlim(0, LEFT_MARGIN + N_COLS * CELL_W)
ax.set_ylim(0, TOP_MARGIN + N_ROWS * CELL_H + BOTTOM_MARGIN)
ax.axis("off")
fig.patch.set_facecolor("white")
ax.set_facecolor("white")

col_centers = [
    LEFT_MARGIN + 0.5 * CELL_W,
    LEFT_MARGIN + 1.5 * CELL_W,
]
col_headers = ["Fermat test", "Miller-Rabin"]

# Column headers
for cx, header in zip(col_centers, col_headers):
    ax.text(
        cx,
        TOP_MARGIN + N_ROWS * CELL_H + 0.12,
        header,
        ha="center",
        va="bottom",
        fontsize=10.5,
        fontweight="bold",
        color="#222222",
    )

# Rows (top = base 2, bottom = base 20)
for row_idx, a in enumerate(BASES):
    # y increases upward; row 0 (a=2) at top
    y_center = TOP_MARGIN + (N_ROWS - 1 - row_idx) * CELL_H + 0.5 * CELL_H

    # Base label
    ax.text(
        LEFT_MARGIN - 0.12,
        y_center,
        f"$a = {a}$",
        ha="right",
        va="center",
        fontsize=10,
        color="#333333",
    )

    results = [fermat_results[row_idx], mr_results[row_idx]]

    for col_idx, is_liar in enumerate(results):
        cx = col_centers[col_idx]
        color = LIAR_COLOR if is_liar else WITNESS_COLOR
        label = "liar" if is_liar else "witness"
        text_color = TEXT_CELL

        # Cell background
        rect = mpatches.FancyBboxPatch(
            (cx - CELL_W * 0.46, y_center - CELL_H * 0.42),
            CELL_W * 0.92,
            CELL_H * 0.84,
            boxstyle="round,pad=0.02",
            facecolor=color,
            edgecolor="none",
            zorder=2,
        )
        ax.add_patch(rect)

        ax.text(
            cx,
            y_center,
            label,
            ha="center",
            va="center",
            fontsize=10,
            color=text_color,
            fontweight="normal",
            zorder=3,
        )

# Thin horizontal dividers between rows
for row_idx in range(N_ROWS + 1):
    y = TOP_MARGIN + row_idx * CELL_H
    ax.axhline(
        y,
        xmin=(LEFT_MARGIN) / ax.get_xlim()[1],
        xmax=1.0,
        color="#e8e8e8",
        linewidth=0.5,
        zorder=1,
    )

# Caption
caption = (
    "$n = 561 = 3 \\times 11 \\times 17$ (smallest Carmichael number). "
    "Every coprime base shown is a Fermat liar; bases sharing a factor reject "
    "561 immediately. Miller-Rabin finds a witness for every base shown."
)
ax.text(
    (LEFT_MARGIN + N_COLS * CELL_W) / 2,
    BOTTOM_MARGIN * 0.4,
    caption,
    ha="center",
    va="center",
    fontsize=10,
    color="#555555",
    wrap=True,
)

plt.tight_layout(pad=0)
validate_figure(fig)
out_path = figure_output_path(__file__, "witness_table.png")
fig.savefig(out_path, dpi=300, bbox_inches="tight", facecolor="white")
print(f"Saved: {out_path}")

# Print summary for verification
print(f"\nn = {N} = 3 × 11 × 17  (n-1 = {N - 1} = 2^4 × 35)\n")
print(f"{'a':>4}  {'gcd':>6}  {'Fermat':>10}  {'Miller-Rabin':>14}")
print("-" * 42)
for a, fr, mr in zip(BASES, fermat_results, mr_results):
    g = math.gcd(a, N)
    print(
        f"{a:>4}  {g:>6}  {'LIAR' if fr else 'witness':>10}  {'LIAR' if mr else 'witness':>14}"
    )
