import argparse
import math
import sys

PHI = 1.6180339887498948482
LN_PHI = math.log(PHI)
SQRT5 = math.sqrt(5.0)
GOLDEN_ANG = 3.8832220774530017295
NUM_DN = 16
FIB_TABLE = [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987]
PRIME_TABLE = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]

ALL_MERSENNE_P = [
    2, 3, 5, 7, 13, 17, 19, 31, 61, 89, 107, 127,
    521, 607, 1279, 2203, 2281, 3217, 4253, 4423,
    9689, 9941, 11213, 19937, 21701, 23209, 44497,
    86243, 110503, 132049, 216091, 756839, 859433,
    1257787, 1398269, 2976221, 3021377, 6972593,
    13466917, 20996011, 24036583, 25964951, 30402457,
    32582657, 37156667, 42643801, 43112609,
    57885161, 74207281, 77232917, 82589933, 136279841
]


def gen_fib(x):
    return (PHI ** x - math.cos(math.pi * x) * PHI ** (-x)) / SQRT5


def prismatic_omega_dim(omega_p, n):
    return 0.5 + 0.5 * math.sin(omega_p + n * GOLDEN_ANG)


def compute_Dn_r(n, r, omega_p):
    if n < 1 or n > NUM_DN:
        return 0.0
    beta = omega_p - math.floor(omega_p)
    x = n + beta
    k = (n + 1) / 8.0
    F_x = abs(gen_fib(x))
    Omega_n = prismatic_omega_dim(omega_p, n)
    P_n = PRIME_TABLE[n - 1]
    two_n = 2.0 ** n
    phi_kx = PHI ** (k * x)
    coeff = math.sqrt(PHI * F_x * two_n * P_n * Omega_n * phi_kx)
    delta = abs(math.cos(math.pi * beta * PHI)) * math.log(P_n) / (PHI ** x)
    one_eff = 1.0 + delta
    scale = coeff * abs(r) ** k
    phase_re = one_eff * math.cos(math.pi * beta)
    phase_im = one_eff * math.sin(math.pi * beta)
    re = scale + phase_re
    return math.sqrt(re * re + phase_im * phase_im)


def compute_pln(p):
    log2x = float(p)
    lnx = log2x * math.log(2.0)
    arg = lnx / LN_PHI
    omega = math.log(arg) / LN_PHI if arg > 1.0 else 0.0
    D = int(math.floor(omega) % 8) + 1
    theta = omega % 1.0
    frac = omega - math.floor(omega)
    theta_rad = frac * 2.0 * math.pi
    D_cont = (omega % 8.0) + 1.0
    lat = PHI * PHI * (1.0 - PHI ** (-D_cont)) + 1.0 / D_cont
    lat = max(1.0, min(lat, PHI * PHI))
    Rre = math.cos(theta_rad) - 1.0 + lat
    Rim = math.sin(theta_rad)
    r = math.hypot(Rre, Rim)
    return {'p': p, 'omega': omega, 'D': D, 'theta': theta, 'r': r}


def compute_fingerprint(p):
    pln = compute_pln(p)
    features = []
    for n in range(1, NUM_DN + 1):
        features.append(compute_Dn_r(n, pln['r'], pln['omega']))
    features.extend([
        pln['omega'], pln['theta'], pln['r'], float(pln['D']),
        pln['omega'] * 170.6180339887 % 1.0,
        pln['omega'] * 1.6180339887 % 1.0,
        pln['r'] * 8.8541019662 % 1.0,
        (pln['theta'] + 0.6180339887) % 1.0,
        (pln['omega'] * 0.2360679775) % 1.0,
        abs(math.sin(pln['omega'] * 0.6180339887 * math.pi))
    ])
    for i in range(4):
        features.append((pln['omega'] % [8.3141592654, 7.8541019662, 6.4721359549, 5.6180339887][i]) /
                        [8.3141592654, 7.8541019662, 6.4721359549, 5.6180339887][i])
    base_seeds = [
        0.6180339887, 1.6180339887, 2.6180339887, 3.6180339887,
        4.8541019662, 5.6180339887, 6.4721359549, 7.8541019662,
        8.3141592654, 0.0901699437, 0.1458980338, 0.2360679775,
        0.3090169944, 0.3819660113, 0.4721359549, 0.6545084972,
        0.8729833462, 1.0000000000, 1.2360679775, 1.6180339887,
        2.2360679775, 2.6180339887, 3.1415926535, 3.6180339887,
        4.2360679775, 4.8541019662, 5.6180339887, 6.4721359549,
        7.2360679775, 7.8541019662, 8.6180339887, 9.2360679775,
        9.8541019662, 10.6180339887, 11.0901699437, 11.9442719100,
        12.6180339887, 13.6180339887, 14.2360679775, 14.8541019662,
        15.6180339887, 16.4721359549, 17.2360679775, 17.9442719100,
        18.6180339887, 19.2360679775, 19.8541019662, 20.6180339887,
        21.0901699437, 21.9442719100, 22.6180339887, 23.6180339887,
        24.2360679775, 24.8541019662, 25.6180339887, 26.4721359549,
        27.2360679775, 27.9442719100, 28.6180339887, 29.0344465435,
        29.6180339887, 30.2360679775, 30.8541019662, 31.6180339887
    ]
    omega_mod32 = pln['omega'] % 32.0
    distances = sorted(abs(omega_mod32 - seed) for seed in base_seeds)
    features.extend(distances[:4])
    d1 = compute_Dn_r(1, pln['r'], pln['omega'])
    d4 = compute_Dn_r(4, pln['r'], pln['omega'])
    d5 = compute_Dn_r(5, pln['r'], pln['omega'])
    d8 = compute_Dn_r(8, pln['r'], pln['omega'])
    ratio = (d1 * d8) / (d4 * d5) if abs(d4 * d5) > 1e-15 else 0.0
    features.append(ratio)
    golden_dist = abs(pln['r'] - 1.0 / PHI)
    features.append(math.exp(-golden_dist * golden_dist * 10.0))
    features.extend([
        (p * PHI) % 1.0,
        (p * PHI * PHI) % 1.0,
        (p * 1.41421356237) % 1.0,
        (p * math.pi) % 1.0,
        (p % 6) / 6.0,
        (p % 8) / 8.0,
        math.sqrt(p) - math.floor(math.sqrt(p)),
        math.log2(p) - math.floor(math.log2(p))
    ])
    return {'p': p, 'features': features}


def fit_model():
    fps = [compute_fingerprint(p) for p in ALL_MERSENNE_P]
    feature_count = len(fps[0]['features'])
    feat_std = [0.0] * feature_count
    feat_weight = [0.0] * feature_count
    for f in range(feature_count):
        vals = [fp['features'][f] for fp in fps]
        mean = sum(vals) / len(vals)
        var = sum((v - mean) ** 2 for v in vals) / len(vals)
        feat_std[f] = math.sqrt(var) if var > 1e-20 else 1.0
        feat_weight[f] = 0.0 if var < 1e-10 else 1.0 / feat_std[f]
    total = sum(feat_weight)
    if total > 0:
        feat_weight = [w / total for w in feat_weight]
    return {'fps': fps, 'feat_std': feat_std, 'feat_weight': feat_weight}


def score_candidate(model, fp):
    max_sim = 0.0
    for known in model['fps']:
        wdist2 = 0.0
        for f, w in enumerate(model['feat_weight']):
            if w < 1e-15:
                continue
            diff = fp['features'][f] - known['features'][f]
            z = diff / model['feat_std'][f]
            wdist2 += w * z * z
        sim = math.exp(-wdist2 / 2.0)
        max_sim = max(max_sim, sim)
    return max_sim


def is_prime_u64(n):
    if n < 2:
        return False
    small = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
    for p in small:
        if n == p:
            return True
        if n % p == 0:
            return False
    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1
    bases = [2, 325, 9375, 28178, 450775, 9780504, 1795265022]
    for a in bases:
        if a % n == 0:
            return True
        x = pow(a, d, n)
        if x == 1 or x == n - 1:
            continue
        composite = True
        for _ in range(s - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                composite = False
                break
        if composite:
            return False
    return True


def trial_factor(p, trial_limit, progress_callback=None):
    max_k = int((trial_limit - 1) // (2 * p)) + 1
    report_every = max(1, max_k // 50)
    if progress_callback is not None:
        progress_callback('trial', 0.0, p=p)
    for k in range(1, max_k):
        if progress_callback is not None and (k % report_every == 0 or k == 1 or k == max_k - 1):
            progress_callback('trial', float(k) / float(max_k), p=p)
        q = 2 * k * p + 1
        if q > trial_limit:
            break
        if q % 8 not in (1, 7):
            continue
        if any(q % d == 0 for d in [3, 5, 7, 11, 13] if q > d):
            continue
        if pow(2, p, q) == 1:
            if progress_callback is not None:
                progress_callback('trial', 1.0, p=p)
            return q
    if progress_callback is not None:
        progress_callback('trial', 1.0, p=p)
    return 0


def mersenne_reduce(x, p):
    if p == 0:
        return 0
    M = (1 << p) - 1
    while x.bit_length() > p:
        x = (x & M) + (x >> p)
    if x >= M:
        x -= M
    return x


def lucas_lehmer(p, progress_callback=None):
    if p < 2:
        return False
    if p == 2:
        return True
    s = 4
    if progress_callback is not None:
        progress_callback('ll', 0.0, p=p, s=s)
    M = (1 << p) - 1
    report_every = max(1, (p - 2) // 50)
    for i in range(p - 2):
        if progress_callback is not None and ((i + 1) % report_every == 0 or i == 0 or i == p - 3):
            progress_callback('ll', float(i + 1) / float(p - 2), p=p, s=s)
        s = mersenne_reduce(s * s - 2, p)
    if progress_callback is not None:
        progress_callback('ll', 1.0, p=p, s=s)
    return s == 0


def verify_candidate(p, trial_bits=24, max_ll_p=100000, progress_callback=None):
    result = {'p': p}
    result['prime_exponent'] = is_prime_u64(p)
    result['trial_factor'] = trial_factor(p, 1 << trial_bits, progress_callback=progress_callback)
    model = fit_model()
    fp = compute_fingerprint(p)
    result['hdgl_score'] = score_candidate(model, fp)
    if p <= max_ll_p:
        result['lucas_lehmer'] = lucas_lehmer(p, progress_callback=progress_callback)
    else:
        result['lucas_lehmer'] = None
    return result


def main():
    parser = argparse.ArgumentParser(description='Mersenne HDGL sieve and verifier')
    parser.add_argument('--score', nargs='+', type=int, help='Score candidate exponents')
    parser.add_argument('--verify', type=int, help='Verify a candidate exponent')
    parser.add_argument('--rank', nargs=3, metavar=('CENTER', 'DELTA', 'TOPN'), help='Rank nearby prime exponents')
    parser.add_argument('--ll', type=int, help='Run exact Lucas-Lehmer for small p')
    args = parser.parse_args()

    if args.score:
        model = fit_model()
        for p in args.score:
            fp = compute_fingerprint(p)
            score = score_candidate(model, fp)
            pln = compute_pln(p)
            print(f'p={p} score={score:.6f} omega={pln["omega"]:.10f} D={pln["D"]}')
        return

    if args.verify:
        p = args.verify
        print(f'Verifying p={p}')
        print(' exponent primality:', 'prime' if is_prime_u64(p) else 'composite')
        factor = trial_factor(p, 1 << 24)
        print(' small factor:', factor if factor else 'none below 2^24')
        model = fit_model()
        fp = compute_fingerprint(p)
        print(f' HDGL score={score_candidate(model, fp):.6f}')
        if p <= 100000:
            print(' Lucas-Lehmer: starting')
            result = lucas_lehmer(p)
            print(' Lucas-Lehmer:', 'prime' if result else 'composite')
        else:
            print(' Lucas-Lehmer: skipped for p > 100000 in Python fallback')
        return

    if args.rank:
        center = int(args.rank[0])
        delta = float(args.rank[1])
        top_n = int(args.rank[2])
        lo = max(3, int(center * (1 - delta / 100.0)))
        hi = int(center * (1 + delta / 100.0))
        primes = [p for p in range(lo, hi + 1) if is_prime_u64(p)]
        model = fit_model()
        ranked = []
        for p in primes:
            fp = compute_fingerprint(p)
            ranked.append((score_candidate(model, fp), p))
        ranked.sort(reverse=True)
        for score, p in ranked[:top_n]:
            print(f'p={p} score={score:.6f}')
        return

    if args.ll:
        p = args.ll
        print('Lucas-Lehmer for p=', p)
        print('result:', 'prime' if lucas_lehmer(p) else 'composite')
        return

    parser.print_help()


if __name__ == '__main__':
    main()
