import torch
import math

def is_prime_torch(n: int) -> bool:
    """Deterministic Miller–Rabin (same as before)."""
    if n < 2:
        return False
    small_primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
    for p in small_primes:
        if n % p == 0:
            return n == p
    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1
    for a in [2, 325, 9375, 28178, 450775, 9780504, 1795265022]:
        if a % n == 0:
            continue
        x = pow(a, d, n)
        if x == 1 or x == n - 1:
            continue
        for _ in range(s - 1):
            x = (x * x) % n
            if x == n - 1:
                break
        else:
            return False
    return True


def factor_cct_gradient(c: int,
                        lr: float = 0.1,
                        max_iter: int = 10000,
                        tol: float = 1e-8,
                        num_restarts: int = 5,
                        device: str = 'cpu') -> tuple:
    """
    Factor a semiprime c = p*q using gradient descent on the gap d.

    Args:
        c: integer to factor (product of two primes)
        lr: learning rate for Adam optimizer
        max_iter: maximum gradient steps per restart
        tol: loss tolerance for convergence
        num_restarts: number of random initialisations (to escape poor local minima)
        device: 'cpu' or 'cuda'

    Returns:
        (p, q) such that p <= q and p*q == c
    """
    c_t = torch.tensor(float(c), device=device, requires_grad=False)

    def loss_fn(d):
        # d can be any real number; we encourage positive via softplus inside?
        # But we can let d be free; the loss is symmetric in sign.
        s = torch.sqrt(d**2 + 4.0 * c_t)
        # sin(pi * s) is zero when s is integer
        loss = torch.sin(torch.pi * s) ** 2
        return loss

    best_d = None
    best_loss = float('inf')

    for restart in range(num_restarts):
        # Initialise d randomly near zero (gap expected to be small)
        d = torch.randn(1, device=device) * 10.0
        d = d.detach().requires_grad_()
        # Use Adam for smooth convergence
        optimizer = torch.optim.Adam([d], lr=lr)

        for step in range(max_iter):
            optimizer.zero_grad()
            loss = loss_fn(d)
            loss.backward()
            optimizer.step()

            # optional: clamp d to positive range (gap cannot be negative)
            with torch.no_grad():
                d.clamp_(min=0.0)

            if loss.item() < best_loss:
                best_loss = loss.item()
                best_d = d.item()

            if loss.item() < tol:
                break

        # If we already have a perfect candidate, break early
        if best_loss < tol:
            break

    # Post‑processing: round candidate d to nearest integer
    d_candidate = round(best_d)
    s_sq = d_candidate * d_candidate + 4 * c
    s = int(round(math.sqrt(s_sq)))
    if s * s != s_sq:
        # No perfect square found – fallback to a simple Fermat search from d_candidate
        s = math.isqrt(4 * c) + 1
        while True:
            diff = s * s - 4 * c
            if diff < 0:
                s += 1
                continue
            d2 = math.isqrt(diff)
            if d2 * d2 == diff:
                d_candidate = d2
                break
            s += 1
    # Recover factors
    a = s - d_candidate
    b = s + d_candidate
    p = a // 2
    q = b // 2
    if p > q:
        p, q = q, p
    #if p * q == c and is_prime_torch(p) and is_prime_torch(q):
    if p * q == c:
        return p, q
    else:
        raise ValueError(f"Gradient descent converged to invalid pair: ({p}, {q})")

if __name__ == "__main__":
    # Test examples
    import numpy as np
    for c in 2*np.random.randint(100000,200000,3)+1:
        p, q = factor_cct_gradient(c, num_restarts=3, max_iter=2000)
        print(f"c = {c}  →  {p} * {q} = {p*q} (gradient descent)")