"""Tests for Theorem Z: support periodicity of v_p(D(n)).

Theorem Z states that for any prime p, v_p(D(n)) > 0 if and only if
n+1 >= p and either:
  (a) ord_2(p) | (n+1)  [Mersenne channel], or
  (b) alpha_L(p) | (n+1) and (n+1)/alpha_L(p) is odd  [Lucas channel].

For p = 2: v_2(D(n)) > 0 iff 3 | (n+1) and n >= 1.

The support period tau_p = vp_support_period(p) is the minimal period
of the binary indicator of conditions (a) OR (b).
"""

import math
import pytest

from math import gcd, factorial
from watkins_nn.kaleidoscope import lucas, mersenne
from watkins_nn.analytic import (
    ord_2,
    alpha_lucas,
    vp_support_period,
    euler_factor_sequence,
)


# ── ord_2 ─────────────────────────────────────────────────────────────────────

class TestOrd2:

    def test_known_values(self):
        """Known multiplicative orders of 2 mod p."""
        expected = {
            3: 2, 5: 4, 7: 3, 11: 10, 13: 12,
            17: 8, 19: 18, 23: 11, 29: 28, 31: 5,
        }
        for p, o in expected.items():
            assert ord_2(p) == o, f"ord_2({p}) should be {o}, got {ord_2(p)}"

    def test_p2_returns_none(self):
        assert ord_2(2) is None

    def test_ord_divides_p_minus_1(self):
        """By Fermat's little theorem, ord_2(p) | (p-1)."""
        for p in [3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]:
            o = ord_2(p)
            assert (p - 1) % o == 0, f"ord_2({p})={o} does not divide {p-1}"

    def test_2_to_ord_is_1_mod_p(self):
        """Verify 2^ord_2(p) = 1 mod p."""
        for p in [3, 5, 7, 11, 13, 17, 19, 23, 29, 31]:
            o = ord_2(p)
            assert pow(2, o, p) == 1

    def test_minimality(self):
        """No smaller k has 2^k = 1 mod p."""
        for p in [3, 5, 7, 11, 13, 17, 19, 23, 29, 31]:
            o = ord_2(p)
            for k in range(1, o):
                assert pow(2, k, p) != 1, f"2^{k} = 1 mod {p} but ord_2 = {o}"

    def test_invalid_raises(self):
        with pytest.raises(ValueError):
            ord_2(1)


# ── alpha_lucas ───────────────────────────────────────────────────────────────

class TestAlphaLucas:

    def test_known_values(self):
        """Known ranks of apparition in Lucas sequence."""
        expected = {2: 3, 3: 2, 7: 4, 11: 5, 19: 9, 29: 7, 31: 15}
        for p, a in expected.items():
            assert alpha_lucas(p) == a, f"alpha_lucas({p}) should be {a}"

    def test_none_for_lucas_immune_primes(self):
        """Primes that never divide any Lucas number."""
        for p in [5, 13, 17]:
            assert alpha_lucas(p) is None, f"alpha_lucas({p}) should be None"

    def test_divisibility_at_alpha(self):
        """L(alpha_L(p)) must be divisible by p."""
        for p in [2, 3, 7, 11, 19, 29, 31]:
            a = alpha_lucas(p)
            assert lucas(a) % p == 0, f"L({a}) not divisible by {p}"

    def test_minimality(self):
        """No smaller k has p | L(k)."""
        for p in [2, 3, 7, 11, 19, 29, 31]:
            a = alpha_lucas(p)
            for k in range(1, a):
                assert lucas(k) % p != 0, f"L({k}) divisible by {p} but alpha = {a}"

    def test_lucas_divisibility_at_odd_multiples(self):
        """p | L(k) iff k/alpha_L(p) is a positive odd integer."""
        for p in [3, 7, 11, 19, 29]:
            a = alpha_lucas(p)
            for k in range(1, 200):
                divides = (lucas(k) % p == 0)
                predicted = (k % a == 0) and ((k // a) % 2 == 1)
                assert divides == predicted, (
                    f"p={p}, k={k}: L(k)%p==0 is {divides}, predicted {predicted}"
                )

    def test_invalid_raises(self):
        with pytest.raises(ValueError):
            alpha_lucas(1)


# ── vp_support_period ─────────────────────────────────────────────────────────

class TestVpSupportPeriod:

    def test_p2_period_3(self):
        assert vp_support_period(2) == 3

    def test_known_periods(self):
        """Support periods computed from the theoretical formula."""
        # These are verified computationally against euler_factor_sequence
        expected = {
            2: 3,    # alpha_L(2) = 3
            3: 2,    # ord_2=2, both contribute with combined period 2
            5: 4,    # ord_2=4, alpha_L=None, period = ord_2
            7: 24,   # lcm(3, 8) = 24, no internal reduction
            11: 5,   # lcm(10, 10) = 10, reduces to 5
            13: 12,  # ord_2=12, alpha_L=None, period = 12
            17: 8,   # ord_2=8, alpha_L=None, period = 8
            19: 9,   # lcm(18, 18) = 18, reduces to 9
            23: 264, # lcm(11, 24) = 264, no reduction
            29: 28,  # lcm(28, 14) = 28, no reduction
            31: 5,   # lcm(5, 30) = 30, reduces to 5
        }
        for p, tau in expected.items():
            assert vp_support_period(p) == tau, (
                f"vp_support_period({p}) should be {tau}, got {vp_support_period(p)}"
            )

    def test_invalid_raises(self):
        with pytest.raises(ValueError):
            vp_support_period(1)


# ── Theorem Z: full support verification ──────────────────────────────────────

def _v_p(n, p):
    """p-adic valuation of n."""
    if n == 0:
        return float('inf')
    v = 0
    while n % p == 0:
        v += 1
        n //= p
    return v


class TestTheoremZSupport:

    @pytest.mark.parametrize("p", [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31])
    def test_support_matches_prediction(self, p):
        """v_p(D(n)) > 0 iff the Mersenne/Lucas channel predicts it."""
        N = 300
        o2 = ord_2(p)
        aL = alpha_lucas(p)
        vals = euler_factor_sequence(p, N)

        for n in range(1, N + 1):
            actual_nz = vals[n - 1] > 0
            k = n + 1
            if p == 2:
                predicted_nz = (k % 3 == 0) and (k >= 2)
            else:
                mersenne_hit = (k % o2 == 0)
                lucas_hit = (aL is not None and k % aL == 0 and (k // aL) % 2 == 1)
                predicted_nz = (mersenne_hit or lucas_hit) and (k >= p)
            assert actual_nz == predicted_nz, (
                f"p={p}, n={n}: v_p(D(n))>0 is {actual_nz}, "
                f"predicted {predicted_nz} (k={k}, o2={o2}, aL={aL})"
            )

    @pytest.mark.parametrize("p", [2, 3, 5, 7, 11, 13, 17, 19, 29, 31])
    def test_support_period_matches_data(self, p):
        """The theoretical tau_p matches the empirical support period."""
        tau = vp_support_period(p)
        N = max(4 * tau + p, 200)
        vals = euler_factor_sequence(p, N)

        # Extract support pattern starting after transient (n+1 >= p)
        start = max(p - 1, 1)
        binary = [1 if v > 0 else 0 for v in vals[start:]]

        # Verify the support repeats with period tau
        for i in range(len(binary)):
            assert binary[i] == binary[i % tau], (
                f"p={p}: support at position {start + i + 1} does not match "
                f"period {tau} (binary[{i}]={binary[i]}, binary[{i % tau}]={binary[i % tau]})"
            )

    def test_mersenne_channel_period_equals_ord2(self):
        """The Mersenne part v_p(gcd(M(n+1),(n+1)!)) has support period ord_2(p)."""
        for p in [3, 5, 7, 11, 13, 17, 19]:
            o2 = ord_2(p)
            for n in range(p, 200):
                k = n + 1
                Mk = mersenne(k)
                fk = factorial(k)
                vpM = _v_p(gcd(Mk, fk), p)
                predicted = (k % o2 == 0)
                assert (vpM > 0) == predicted, (
                    f"p={p}, n={n}: Mersenne channel mismatch"
                )

    def test_lucas_channel_period(self):
        """The Lucas part v_p(gcd(L(n+1),(n+1)!)) has support at odd multiples of alpha_L."""
        for p in [3, 7, 11, 19, 29]:
            aL = alpha_lucas(p)
            for n in range(p, 200):
                k = n + 1
                Lk = lucas(k)
                fk = factorial(k)
                vpL = _v_p(gcd(Lk, fk), p)
                predicted = (k % aL == 0) and ((k // aL) % 2 == 1)
                assert (vpL > 0) == predicted, (
                    f"p={p}, n={n}: Lucas channel mismatch (k={k}, aL={aL})"
                )


# ── Additive decomposition ────────────────────────────────────────────────────

class TestVpDecomposition:

    @pytest.mark.parametrize("p", [2, 3, 5, 7, 11])
    def test_vp_D_equals_vp_gM_plus_vp_gL(self, p):
        """v_p(D(n)) = v_p(gcd(M,f!)) + v_p(gcd(L,f!)) since D = gM * gL."""
        for n in range(1, 100):
            Mn1 = mersenne(n + 1)
            Ln1 = lucas(n + 1)
            fn1 = factorial(n + 1)
            gM = gcd(Mn1, fn1)
            gL = gcd(Ln1, fn1)
            D = gM * gL
            vp_D = _v_p(D, p)
            vp_gM = _v_p(gM, p)
            vp_gL = _v_p(gL, p)
            assert vp_D == vp_gM + vp_gL, (
                f"p={p}, n={n}: v_p(D)={vp_D} != v_p(gM)+v_p(gL)={vp_gM}+{vp_gL}"
            )
