"""
fta.arithmetic

Prime factorization and coordinate arithmetic.

Every positive integer maps uniquely to a coordinate vector in N^k where each
axis corresponds to a prime and each coordinate is the exponent of that prime:

    n = p1^a1 * p2^a2 * ... * pk^ak  <-->  (a1, a2, ..., ak)

Under this embedding:
    multiplication  <-->  coordinate addition
    divisibility    <-->  coordinate-wise <=
    gcd             <-->  coordinate-wise min
    lcm             <-->  coordinate-wise max

Original implementation by Shun Richard Honda (upfpy, 2002–2013).
"""

from __future__ import annotations

import math
import time
from typing import Callable, Dict, Tuple

# A Coordinate maps prime → exponent.
# e.g. 12 = 2^2 * 3^1  →  {2: 2, 3: 1}
Coordinate = Dict[int, int]


# ---------------------------------------------------------------------------
# Factorization backends
# ---------------------------------------------------------------------------

_TIMING_THRESHOLD: float = 0.01  # 10ms — suggest sympy if exceeded


def _wheel_factorize(n: int) -> Coordinate:
    """
    Trial division with 6k±1 wheel.

    Every prime > 3 is of the form 6k−1 or 6k+1. Checks 2 and 3 first,
    then 6k±1 candidates up to √n. Correct and fast for all integers this
    system typically encounters (up to ~10^8).
    """
    if n < 1:
        raise ValueError(f"factorize requires a positive integer, got {n}")
    if n == 1:
        return {}

    factors: Coordinate = {}
    remaining = n

    while remaining % 2 == 0:
        factors[2] = factors.get(2, 0) + 1
        remaining //= 2

    while remaining % 3 == 0:
        factors[3] = factors.get(3, 0) + 1
        remaining //= 3

    k = 1
    while True:
        for candidate in (6 * k - 1, 6 * k + 1):
            if candidate * candidate > remaining:
                break
            while remaining % candidate == 0:
                factors[candidate] = factors.get(candidate, 0) + 1
                remaining //= candidate
        else:
            k += 1
            continue
        break

    if remaining > 1:
        factors[remaining] = factors.get(remaining, 0) + 1

    return factors


def _sympy_factorize(n: int) -> Coordinate:
    """Sympy backend. Requires sympy to be installed."""
    try:
        from sympy import factorint
    except ImportError:
        raise ImportError(
            "sympy is not installed. "
            "pip install sympy  or use the default wheel backend."
        )
    if n < 1:
        raise ValueError(f"factorize requires a positive integer, got {n}")
    if n == 1:
        return {}
    return dict(factorint(n))


_backend: Callable[[int], Coordinate] = _wheel_factorize
_cache: Dict[int, Coordinate] = {}


def set_factorization_backend(backend) -> None:
    """
    Set the factorization backend.

    Parameters
    ----------
    backend : "wheel" | "sympy" | callable
        "wheel"  — internal 6k±1 trial division (default, no dependencies)
        "sympy"  — sympy.factorint (requires sympy)
        callable — any fn(n: int) -> dict[int, int]

    The cache is NOT cleared automatically. Call clear_factorization_cache()
    after switching if you need recomputed values.
    """
    global _backend
    if backend == "wheel":
        _backend = _wheel_factorize
    elif backend == "sympy":
        _backend = _sympy_factorize
    elif callable(backend):
        _backend = backend
    else:
        raise ValueError(
            f"Unknown backend {backend!r}. "
            f"Use 'wheel', 'sympy', or a callable fn(n: int) -> dict[int, int]."
        )


def clear_factorization_cache() -> None:
    """Clear the factorization cache."""
    _cache.clear()


def factorization_cache_size() -> int:
    """Return the number of integers currently cached."""
    return len(_cache)


# ---------------------------------------------------------------------------
# factorize — public entry point
# ---------------------------------------------------------------------------


def factorize(n: int) -> Coordinate:
    """
    Return the prime exponent coordinate of integer n.

    factorize(1)    -> {}
    factorize(12)   -> {2: 2, 3: 1}
    factorize(360)  -> {2: 3, 3: 2, 5: 1}

    Results are cached indefinitely. The empty dict is the coordinate of 1
    — the bottom of every divisibility lattice.
    """
    if n in _cache:
        return _cache[n]

    t0 = time.monotonic()
    result = _backend(n)
    elapsed = time.monotonic() - t0
    _cache[n] = result

    if elapsed > _TIMING_THRESHOLD and _backend is _wheel_factorize:
        print(
            f"[fta] factorize({n}) took {elapsed*1000:.1f}ms. "
            f"For large integers: set_factorization_backend('sympy')"
        )

    return result


# ---------------------------------------------------------------------------
# Coordinate arithmetic
# ---------------------------------------------------------------------------


def to_int(coord: Coordinate) -> int:
    """
    Reconstruct the integer from its prime exponent coordinate.

    to_int({})           -> 1
    to_int({2: 1})       -> 2
    to_int({2: 2, 3: 1}) -> 12
    """
    result = 1
    for prime, exp in coord.items():
        result *= prime ** exp
    return result


def _divisors(n: int) -> Tuple[int, ...]:
    """All positive divisors of n in ascending order."""
    coord = factorize(n)
    result = [1]
    for prime, exp in coord.items():
        result = [d * prime ** e for d in result for e in range(exp + 1)]
    return tuple(sorted(result))


def smallest_divisor_gte(n: int, floor: int) -> int:
    """
    Return the smallest divisor of n that is >= floor.

    smallest_divisor_gte(3600, 250) -> 300
    smallest_divisor_gte(3600, 300) -> 300
    smallest_divisor_gte(3600, 301) -> 400

    Useful for finding the coarsest admissible bin size given a minimum grain.
    """
    for d in _divisors(n):
        if d >= floor:
            return d
    raise ValueError(f"No divisor of {n} is >= {floor}")
