#!/usr/bin/env python3
"""RCL-guided search for Erdős-Straus representations.

The search uses the Lean-certified reduction in
`IndisputableMonolith.NumberTheory.ErdosStrausRCL`:

    c = 4*x - n
    N = n*x
    (c*y - N) * (c*z - N) = N^2

For fixed `n`, start at the least denominator with positive residual,
`x = floor(n/4) + 1`, and increase by a small offset.  A factor pair of
`N^2` with both factors congruent to `-N mod c` gives `y,z` immediately.
"""

from __future__ import annotations

from dataclasses import dataclass
from math import isqrt


@dataclass(frozen=True)
class ErdosStrausSolution:
    n: int
    x: int
    y: int
    z: int
    c: int
    d: int
    e: int
    offset: int

    def check(self) -> bool:
        return 4 * self.x * self.y * self.z == self.n * (
            self.x * self.y + self.x * self.z + self.y * self.z
        )


def factor(n: int) -> list[tuple[int, int]]:
    """Trial-divide `n` into prime powers."""
    factors: list[tuple[int, int]] = []
    d = 2
    while d * d <= n:
        if n % d == 0:
            exponent = 0
            while n % d == 0:
                n //= d
                exponent += 1
            factors.append((d, exponent))
        d += 1 if d == 2 else 2
    if n > 1:
        factors.append((n, 1))
    return factors


def divisors_from_factorization(factors: list[tuple[int, int]]) -> list[int]:
    divisors = [1]
    for prime, exponent in factors:
        powers = [prime ** k for k in range(exponent + 1)]
        divisors = [d * p for d in divisors for p in powers]
    return divisors


def factor_pair_for_residual_square(N: int, c: int) -> tuple[int, int] | None:
    """Return d,e with d*e=N^2 and d=e=-N mod c, if found."""
    target = (-N) % c
    square = N * N
    square_factors = [(prime, 2 * exponent) for prime, exponent in factor(N)]
    for d in divisors_from_factorization(square_factors):
        if d > isqrt(square):
            continue
        e = square // d
        if d % c == target and e % c == target:
            return d, e
    return None


def solve(n: int, max_offset: int = 64) -> ErdosStrausSolution | None:
    if n < 2:
        raise ValueError("Erdős-Straus starts at n >= 2")
    base = n // 4 + 1
    for offset in range(max_offset + 1):
        x = base + offset
        c = 4 * x - n
        N = n * x
        pair = factor_pair_for_residual_square(N, c)
        if pair is None:
            continue
        d, e = pair
        y = (N + d) // c
        z = (N + e) // c
        out = ErdosStrausSolution(n, x, y, z, c, d, e, offset)
        if out.check():
            return out
    return None


def profile(limit: int, max_offset: int = 64, hard_class_only: bool = False) -> None:
    c_hist: dict[int, int] = {}
    max_seen_offset = 0
    misses: list[int] = []
    ns = range(5, limit + 1, 4) if hard_class_only else range(2, limit + 1)
    for n in ns:
        sol = solve(n, max_offset=max_offset)
        if sol is None:
            misses.append(n)
            continue
        c_hist[sol.c] = c_hist.get(sol.c, 0) + 1
        max_seen_offset = max(max_seen_offset, sol.offset)

    scope = "hard class n = 1 mod 4" if hard_class_only else "all n"
    print(f"checked {scope}, n <= {limit}")
    print(f"misses: {len(misses)}")
    if misses:
        print("first misses:", misses[:20])
    print(f"max offset used: {max_seen_offset}")
    print("c histogram:")
    for c, count in sorted(c_hist.items()):
        print(f"  c={c}: {count}")


if __name__ == "__main__":
    profile(250_000, hard_class_only=True)