#!/usr/bin/env python3
"""乌岚素数螺旋 (Ulam Spiral) 实验

1963 年,数学家斯坦尼斯瓦夫·乌岚在听一个冗长报告时百无聊赖,
在一张纸条上把整数写成螺旋,然后圈出素数——结果素数的分布
呈现出神秘的斜向条纹,至今没有完全解释清楚。

本脚本:
1. 生成 size×size 的乌岚螺旋(1 在中心,顺时针向外);
2. 用 ASCII 渲染素数分布;
3. 统计两条主对角线的素数密度,与整体密度对比;
4. 验证欧拉素数多项式 n^2 + n + 41 的 40 个连续素数
   是否恰好落在螺旋的一条直线上。
"""

import sys


def sieve(n: int) -> list[bool]:
    """埃氏筛,返回 is_prime[0..n]"""
    is_prime = [True] * (n + 1)
    is_prime[0] = is_prime[1] = False
    for i in range(2, int(n ** 0.5) + 1):
        if is_prime[i]:
            for j in range(i * i, n + 1, i):
                is_prime[j] = False
    return is_prime


def build_spiral(size: int):
    """顺时针乌岚螺旋:1 在中心,2 在右边,3 在上方……
    返回 (grid, pos_of):grid[(x,y)] = 数值,pos_of[数值] = (x,y)"""
    assert size % 2 == 1, "size 必须为奇数"
    grid = {(0, 0): 1}
    pos_of = {1: (0, 0)}
    dirs = [(1, 0), (0, 1), (-1, 0), (0, -1)]  # 右、上、左、下
    x = y = 0
    v = 1
    d = 0
    step = 1
    taken = 0
    while v < size * size:
        dx, dy = dirs[d]
        x += dx
        y += dy
        v += 1
        grid[(x, y)] = v
        pos_of[v] = (x, y)
        taken += 1
        if taken == step:
            taken = 0
            d = (d + 1) % 4
            if d % 2 == 0:
                step += 1  # 步长序列:1,1,2,2,3,3,...
    return grid, pos_of


def render(grid: dict, size: int, is_prime: list[bool]) -> str:
    """渲染 ASCII 图:素数 #,合数空白,外框 ·"""
    half = size // 2
    lines = []
    for y in range(half, -half - 1, -1):
        row = ["·"] + ["#" if is_prime[grid[(x, y)]] else " " for x in range(-half, half + 1)] + ["·"]
        lines.append("".join(row))
    border = "·" * (size + 2)
    return "\n".join([border] + lines + [border])


def main() -> None:
    size = int(sys.argv[1]) if len(sys.argv) > 1 else 61
    n_max = size * size
    is_prime = sieve(n_max)
    grid, pos_of = build_spiral(size)

    art = render(grid, size, is_prime)
    print(art)
    print()

    # ---- 整体素数密度 ----
    primes = [v for v in range(2, n_max + 1) if is_prime[v]]
    total_density = len(primes) / n_max
    print(f"螺旋 {size}×{size} = {n_max} 个数,素数 {len(primes)} 个,"
          f"整体密度 {total_density:.3%}")

    # ---- 两条主对角线(穿过中心的 x==y 与 x==-y)的素数密度 ----
    half = size // 2
    diag_cells = [grid[(i, i)] for i in range(-half, half + 1)]
    anti_cells = [grid[(i, -i)] for i in range(-half, half + 1)]
    for name, cells in (("主对角线 x=y", diag_cells), ("副对角线 x=-y", anti_cells)):
        p = sum(1 for v in cells if is_prime[v])
        print(f"{name}:{len(cells)} 格,素数 {p} 个,"
              f"密度 {p / len(cells):.3%}"
              f"{'  ↑ 高于整体' if p / len(cells) > total_density else ''}")

    # ---- 扫描所有对角线,找出素数最密的几条 ----
    print("\n--- 全螺旋对角线扫描(素数密度 TOP 3,格子数 ≥ 40) ---")
    lines = []
    for c in range(-2 * half, 2 * half + 1):
        cells = [grid[(i, c - i)] for i in range(-half, half + 1) if (i, c - i) in grid]
        if len(cells) >= 40:
            p = sum(1 for v in cells if is_prime[v])
            lines.append((p / len(cells), p, len(cells), f"x+y={c}"))
        cells = [grid[(i, i - c)] for i in range(-half, half + 1) if (i, i - c) in grid]
        if len(cells) >= 40:
            p = sum(1 for v in cells if is_prime[v])
            lines.append((p / len(cells), p, len(cells), f"x-y={c}"))
    lines.sort(reverse=True)
    for density, p, n, name in lines[:3]:
        print(f"  {name}:{n} 格,素数 {p} 个,密度 {density:.3%}")

    # ---- 欧拉多项式 n^2 + n + 41 ----
    print("\n--- 欧拉素数多项式 n²+n+41 (n=0..39) ---")
    pts = []
    ok_prime = True
    for n in range(40):
        p = n * n + n + 41
        if not is_prime[p]:
            ok_prime = False
            print(f"  n={n}: {p} 不是素数!")
        pts.append(pos_of[p])
    print(f"  40 个值全部为素数: {ok_prime}")
    xs = [x for x, _ in pts]
    ys = [y for _, y in pts]
    if len(set(x - y for x, y in pts)) == 1:
        line = f"x − y = {xs[0] - ys[0]}"
    elif len(set(x + y for x, y in pts)) == 1:
        line = f"x + y = {xs[0] + ys[0]}"
    else:
        line = "并不在一条直线上(而是沿一条弧线爬升)"
    print(f"  这 40 个点在螺旋中的位置:{line}")
    print(f"  坐标范围 x∈[{min(xs)},{max(xs)}], y∈[{min(ys)},{max(ys)}]")
    # 这 40 个点里,落在某一条对角线上的最多有多少个?
    from collections import Counter
    c1 = Counter(x - y for x, y in pts).most_common(1)[0]
    c2 = Counter(x + y for x, y in pts).most_common(1)[0]
    best = c1 if c1[1] >= c2[1] else c2
    print(f"  落在同一条对角线上的最多有 {best[1]} 个(x{'−y' if best is c1 else '+y'}={best[0]})")
    # 对比:这段弧线所在环带(半径 20~21)内的整体素数密度
    ring = [v for (x, y), v in grid.items() if max(abs(x), abs(y)) >= 19]
    ring_density = sum(1 for v in ring if is_prime[v]) / len(ring)
    print(f"  对照:螺旋外环带(半径 19~30)整体素数密度 {ring_density:.3%}")


if __name__ == "__main__":
    main()