"""Tests for the analytic module: prime decomposition of D(n) and Dirichlet series.

Tests cover:
- prime_factorization_D consistency with arithmetic_denominator
- D_sequence batch vs individual computation
- divisor_count for known values
- D_growth_normalized basic properties
- dirichlet_partial convergence and monotonicity
- euler_factor_sequence periodicity patterns
- abscissa_estimate bounds
"""

import math
import pytest

from watkins_nn.coupling import arithmetic_denominator
from watkins_nn.analytic import (
    prime_factorization_D,
    D_sequence,
    divisor_count,
    D_growth_normalized,
    dirichlet_partial,
    euler_factor_sequence,
    abscissa_estimate,
    ord_2,
    alpha_lucas,
    vp_support_period,
    vp_support_period_fast,
    support_pattern,
    prime_classification_table,
    D_mersenne_component,
    D_lucas_component,
    dirichlet_channel_partial,
)


# ── prime_factorization_D ────────────────────────────────────────────────────

class TestPrimeFactorizationD:

    def test_product_equals_D_for_n_1_to_50(self):
        """Product of p^v_p must reconstruct D(n) exactly."""
        for n in range(1, 51):
            D_expected = arithmetic_denominator(n)[0]
            pf = prime_factorization_D(n)
            product = 1
            for p, v in pf.items():
                product *= p ** v
            assert product == D_expected, f"n={n}: product={product}, D={D_expected}"

    def test_D_1_returns_empty(self):
        """D(n) = 1 should give empty factorization."""
        # n=1, 4, 6, 10, 12, 16 all have D=1
        for n in [1, 4, 6, 10, 12, 16]:
            assert prime_factorization_D(n) == {}, f"n={n} should give empty dict"

    def test_known_factorization_n5(self):
        """D(5) = 162 = 2 * 3^4."""
        pf = prime_factorization_D(5)
        assert pf == {2: 1, 3: 4}

    def test_known_factorization_n8(self):
        """D(8) = 28 = 2^2 * 7."""
        pf = prime_factorization_D(8)
        assert pf == {2: 2, 7: 1}

    def test_known_factorization_n11(self):
        """D(11) = 4410 = 2 * 3^2 * 5 * 7^2."""
        pf = prime_factorization_D(11)
        assert pf == {2: 1, 3: 2, 5: 1, 7: 2}

    def test_all_factors_are_prime(self):
        """Every key in the factorization must be prime."""
        def is_prime(n):
            if n < 2:
                return False
            for d in range(2, int(n**0.5) + 1):
                if n % d == 0:
                    return False
            return True

        for n in range(1, 51):
            pf = prime_factorization_D(n)
            for p in pf:
                assert is_prime(p), f"n={n}: {p} is not prime"

    def test_all_exponents_positive(self):
        """Every exponent must be a positive integer."""
        for n in range(1, 51):
            pf = prime_factorization_D(n)
            for p, v in pf.items():
                assert isinstance(v, int) and v > 0, f"n={n}, p={p}: v={v}"

    def test_invalid_n_raises(self):
        with pytest.raises(ValueError):
            prime_factorization_D(0)


# ── D_sequence ───────────────────────────────────────────────────────────────

class TestDSequence:

    def test_matches_individual_calls(self):
        """D_sequence(N) must match individual arithmetic_denominator calls."""
        N = 50
        seq = D_sequence(N)
        assert len(seq) == N
        for i, D in enumerate(seq):
            assert D == arithmetic_denominator(i + 1)[0], f"mismatch at n={i+1}"

    def test_length(self):
        assert len(D_sequence(1)) == 1
        assert len(D_sequence(25)) == 25

    def test_first_values(self):
        seq = D_sequence(10)
        expected = [1, 2, 3, 1, 162, 1, 15, 28, 9, 1]
        assert seq == expected

    def test_invalid_N_raises(self):
        with pytest.raises(ValueError):
            D_sequence(0)


# ── divisor_count ────────────────────────────────────────────────────────────

class TestDivisorCount:

    def test_prime_has_2_divisors(self):
        for p in [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]:
            assert divisor_count(p) == 2, f"prime {p} should have 2 divisors"

    def test_prime_power(self):
        """sigma_0(p^k) = k + 1."""
        assert divisor_count(4) == 3    # 2^2
        assert divisor_count(8) == 4    # 2^3
        assert divisor_count(16) == 5   # 2^4
        assert divisor_count(27) == 4   # 3^3
        assert divisor_count(9) == 3    # 3^2

    def test_known_composite(self):
        assert divisor_count(12) == 6   # 1,2,3,4,6,12
        assert divisor_count(24) == 8   # 1,2,3,4,6,8,12,24
        assert divisor_count(36) == 9   # perfect square

    def test_one(self):
        assert divisor_count(1) == 1

    def test_multiplicative_property(self):
        """sigma_0 is multiplicative: sigma_0(mn) = sigma_0(m)*sigma_0(n) for gcd(m,n)=1."""
        # gcd(3,4) = 1, so sigma_0(12) = sigma_0(3) * sigma_0(4)
        assert divisor_count(12) == divisor_count(3) * divisor_count(4)
        # gcd(5,8) = 1
        assert divisor_count(40) == divisor_count(5) * divisor_count(8)

    def test_invalid_raises(self):
        with pytest.raises(ValueError):
            divisor_count(0)


# ── D_growth_normalized ──────────────────────────────────────────────────────

class TestDGrowthNormalized:

    def test_length(self):
        ratios = D_growth_normalized(10)
        assert len(ratios) == 9   # n=2..10

    def test_all_nonnegative(self):
        ratios = D_growth_normalized(50)
        for r in ratios:
            assert r >= 0.0

    def test_contains_zeros_for_D_equals_1(self):
        """When D(n)=1, ratio is 1/(sigma_0(n+1)*ln(n)) which is small but positive."""
        ratios = D_growth_normalized(20)
        # D(4)=1, n=4 is index 2 (0-based), ratio should be small
        assert ratios[2] < 1.0  # D(4)=1, ratio = 1/(sigma_0(5)*ln(4)) ~ 0.36

    def test_invalid_N_raises(self):
        with pytest.raises(ValueError):
            D_growth_normalized(1)


# ── dirichlet_partial ────────────────────────────────────────────────────────

class TestDirichletPartial:

    def test_single_term(self):
        """sum at N=1 should be D(1)/1^s = 1."""
        assert dirichlet_partial(1, 2.0) == 1.0

    def test_two_terms(self):
        """N=2: D(1)/1^s + D(2)/2^s = 1 + 2/4 = 1.5 for s=2."""
        assert abs(dirichlet_partial(2, 2.0) - 1.5) < 1e-12

    def test_monotone_in_N(self):
        """Partial sums are monotonically increasing (all D(n) >= 1)."""
        prev = 0.0
        for N in range(1, 51):
            val = dirichlet_partial(N, 2.0)
            assert val >= prev, f"N={N}: {val} < {prev}"
            prev = val

    def test_convergence_above_abscissa(self):
        """Above the abscissa (~6.2), partial sums should stabilize."""
        s = 8.0  # well above abscissa of convergence
        s50 = dirichlet_partial(50, s)
        s100 = dirichlet_partial(100, s)
        relative_change = abs(s100 - s50) / s50
        assert relative_change < 0.01, f"relative change {relative_change} too large"

    def test_s_zero_gives_sum_of_D(self):
        """At s=0, sum = D(1) + D(2) + ... + D(N)."""
        N = 10
        expected = sum(D_sequence(N))
        assert abs(dirichlet_partial(N, 0.0) - expected) < 1e-10

    def test_invalid_N_raises(self):
        with pytest.raises(ValueError):
            dirichlet_partial(0, 2.0)


# ── euler_factor_sequence ────────────────────────────────────────────────────

class TestEulerFactorSequence:

    def test_v2_period_3_pattern(self):
        """v_2(D(n)) has a periodic pattern with period 3 in chunks of 3."""
        vals = euler_factor_sequence(2, 50)
        # Check: v_2 is 0 whenever n mod 3 == 0 (1-indexed: n=3,6,9,...)
        for i in range(len(vals)):
            n = i + 1
            if n % 3 == 0:
                assert vals[i] == 0, f"v_2(D({n})) should be 0"

    def test_v2_known_values(self):
        vals = euler_factor_sequence(2, 10)
        expected = [0, 1, 0, 0, 1, 0, 0, 2, 0, 0]
        assert vals == expected

    def test_v3_known_values(self):
        vals = euler_factor_sequence(3, 10)
        expected = [0, 0, 1, 0, 4, 0, 1, 0, 2, 0]
        assert vals == expected

    def test_v5_period_4(self):
        """v_5 is nonzero only at positions n = 4k+3 (n=7,11,15,19,...)."""
        vals = euler_factor_sequence(5, 50)
        for i in range(len(vals)):
            n = i + 1
            if n % 4 == 3:
                # These are candidates for nonzero v_5
                pass
            else:
                assert vals[i] == 0, f"v_5(D({n}))={vals[i]} should be 0 (n mod 4 = {n%4})"

    def test_v11_nonzero_positions(self):
        """v_11(D(n)) is nonzero roughly every 5 odd positions."""
        vals = euler_factor_sequence(11, 50)
        nonzero_positions = [i + 1 for i, v in enumerate(vals) if v > 0]
        # From data: [14, 19, 24, 29, 34, 39, 44, 49]
        # Differences are all 5
        if len(nonzero_positions) >= 3:
            diffs = [nonzero_positions[i+1] - nonzero_positions[i]
                     for i in range(len(nonzero_positions) - 1)]
            assert all(d == 5 for d in diffs), f"v_11 diffs: {diffs}"

    def test_consistency_with_prime_factorization(self):
        """v_p from euler_factor_sequence matches prime_factorization_D."""
        for p in [2, 3, 5, 7]:
            vals = euler_factor_sequence(p, 30)
            for n in range(1, 31):
                pf = prime_factorization_D(n)
                expected = pf.get(p, 0)
                assert vals[n - 1] == expected, \
                    f"p={p}, n={n}: euler={vals[n-1]}, pf={expected}"

    def test_invalid_p_raises(self):
        with pytest.raises(ValueError):
            euler_factor_sequence(1, 10)

    def test_invalid_N_raises(self):
        with pytest.raises(ValueError):
            euler_factor_sequence(2, 0)


# ── abscissa_estimate ────────────────────────────────────────────────────────

class TestAbscissaEstimate:

    def test_positive(self):
        """Abscissa should be positive since D(n) can be >> 1."""
        a = abscissa_estimate(50)
        assert a > 0.0

    def test_finite(self):
        """Abscissa should be finite."""
        a = abscissa_estimate(100)
        assert math.isfinite(a)

    def test_bounded_above(self):
        """D(n) grows at most polynomially, so abscissa should be moderate."""
        a = abscissa_estimate(100)
        assert a < 10.0, f"abscissa {a} seems too large"

    def test_monotone_nondecreasing(self):
        """With more terms, the limsup can only stay or increase."""
        a50 = abscissa_estimate(50)
        a100 = abscissa_estimate(100)
        assert a100 >= a50 - 1e-12

    def test_invalid_N_raises(self):
        with pytest.raises(ValueError):
            abscissa_estimate(1)


# ── vp_support_period_fast ──────────────────────────────────────────────────

def _primes_up_to(n):
    """Simple sieve for test use."""
    sieve = [True] * (n + 1)
    sieve[0] = sieve[1] = False
    for i in range(2, int(n**0.5) + 1):
        if sieve[i]:
            for j in range(i*i, n + 1, i):
                sieve[j] = False
    return [i for i in range(2, n + 1) if sieve[i]]


class TestVpSupportPeriodFast:

    def test_agrees_with_original_for_small_primes(self):
        """vp_support_period_fast must agree with vp_support_period for p < 200."""
        for p in _primes_up_to(200):
            expected = vp_support_period(p)
            fast = vp_support_period_fast(p)
            assert fast == expected, (
                f"p={p}: fast={fast}, expected={expected}, "
                f"ord_2={ord_2(p)}, alpha_lucas={alpha_lucas(p)}"
            )

    def test_p2(self):
        assert vp_support_period_fast(2) == 3

    def test_p5_no_lucas(self):
        """p=5: alpha_lucas=None, so tau = ord_2(5) = 4."""
        assert alpha_lucas(5) is None
        assert vp_support_period_fast(5) == 4

    def test_p11(self):
        assert vp_support_period_fast(11) == 5

    def test_known_values(self):
        """Check a few known support periods."""
        known = {2: 3, 3: 2, 5: 4, 7: 24, 11: 5, 13: 12}
        for p, tau in known.items():
            assert vp_support_period_fast(p) == tau, f"p={p}"

    def test_large_prime_1009(self):
        """Fast version handles p=1009 easily."""
        tau = vp_support_period_fast(1009)
        assert isinstance(tau, int) and tau >= 1

    def test_large_prime_10007(self):
        """Fast version handles p=10007."""
        tau = vp_support_period_fast(10007)
        assert isinstance(tau, int) and tau >= 1

    def test_large_prime_100003(self):
        """Fast version handles p=100003."""
        tau = vp_support_period_fast(100003)
        assert isinstance(tau, int) and tau >= 1

    def test_period_divides_lcm(self):
        """tau_p must divide lcm(ord_2(p), 2*alpha_lucas(p)) when both exist."""
        for p in _primes_up_to(100):
            if p == 2:
                continue
            a = ord_2(p)
            b = alpha_lucas(p)
            tau = vp_support_period_fast(p)
            if b is not None:
                full = math.lcm(a, 2 * b)
                assert full % tau == 0, f"p={p}: tau={tau} does not divide lcm={full}"
            else:
                assert tau == a, f"p={p}: tau={tau} != ord_2={a}"

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


# ── support_pattern ─────────────────────────────────────────────────────────

class TestSupportPattern:

    def test_p2_first_10(self):
        """support_pattern(2, 10) should match v_2 > 0."""
        import numpy as np
        pat = support_pattern(2, 10)
        # v_2(D(n)) for n=1..10: [0,1,0,0,1,0,0,2,0,0]
        expected = [0, 1, 0, 0, 1, 0, 0, 1, 0, 0]
        assert list(pat) == expected

    def test_p5_matches_euler(self):
        """support_pattern(5, N) matches euler_factor_sequence support."""
        import numpy as np
        N = 50
        pat = support_pattern(5, N)
        v5 = euler_factor_sequence(5, N)
        expected = [1 if x > 0 else 0 for x in v5]
        assert list(pat) == expected, f"mismatch: {list(pat)} vs {expected}"

    def test_p3_matches_euler(self):
        import numpy as np
        N = 50
        pat = support_pattern(3, N)
        v3 = euler_factor_sequence(3, N)
        expected = [1 if x > 0 else 0 for x in v3]
        assert list(pat) == expected

    def test_p7_matches_euler(self):
        import numpy as np
        N = 50
        pat = support_pattern(7, N)
        v7 = euler_factor_sequence(7, N)
        expected = [1 if x > 0 else 0 for x in v7]
        assert list(pat) == expected

    def test_p11_matches_euler(self):
        import numpy as np
        N = 50
        pat = support_pattern(11, N)
        v11 = euler_factor_sequence(11, N)
        expected = [1 if x > 0 else 0 for x in v11]
        assert list(pat) == expected

    def test_p13_matches_euler(self):
        import numpy as np
        N = 60
        pat = support_pattern(13, N)
        v13 = euler_factor_sequence(13, N)
        expected = [1 if x > 0 else 0 for x in v13]
        assert list(pat) == expected

    def test_all_primes_to_50_match_euler(self):
        """Comprehensive check: support_pattern matches brute-force for all primes < 50."""
        import numpy as np
        N = 80
        for p in _primes_up_to(50):
            pat = support_pattern(p, N)
            vp = euler_factor_sequence(p, N)
            expected = [1 if x > 0 else 0 for x in vp]
            assert list(pat) == expected, f"p={p}: mismatch"

    def test_dtype_and_shape(self):
        import numpy as np
        pat = support_pattern(7, 100)
        assert pat.shape == (100,)
        assert pat.dtype == np.uint8

    def test_zeros_before_threshold(self):
        """For n+1 < p, support must be 0."""
        import numpy as np
        p = 37
        pat = support_pattern(p, 100)
        # n+1 < p means n < p-1 = 36, i.e. indices 0..34 (n=1..35)
        for i in range(p - 2):
            assert pat[i] == 0, f"p={p}, n={i+1}: should be 0 before threshold"

    def test_invalid_p_raises(self):
        with pytest.raises(ValueError):
            support_pattern(1, 10)

    def test_invalid_N_raises(self):
        with pytest.raises(ValueError):
            support_pattern(2, 0)


# ── prime_classification_table ─────────────────────────────────────────────

class TestPrimeClassificationTable:

    def test_keys_are_primes(self):
        """All keys must be primes less than N."""
        table = prime_classification_table(50)
        expected_primes = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47}
        assert set(table.keys()) == expected_primes

    def test_fields_present(self):
        """Each entry must have all five fields."""
        table = prime_classification_table(20)
        for p, info in table.items():
            assert 'ord2' in info
            assert 'alpha_L' in info
            assert 'chi4' in info
            assert 'tau' in info
            assert 'mod4' in info

    def test_chi4_values(self):
        """chi_minus4 values must be correct."""
        table = prime_classification_table(30)
        assert table[2]['chi4'] == 0
        assert table[5]['chi4'] == 1    # 5 ≡ 1 (mod 4)
        assert table[13]['chi4'] == 1   # 13 ≡ 1 (mod 4)
        assert table[3]['chi4'] == -1   # 3 ≡ 3 (mod 4)
        assert table[7]['chi4'] == -1   # 7 ≡ 3 (mod 4)

    def test_mod4_values(self):
        table = prime_classification_table(20)
        for p, info in table.items():
            assert info['mod4'] == p % 4

    def test_consistency_with_individual_functions(self):
        """Table values must match individual function calls."""
        table = prime_classification_table(50)
        for p, info in table.items():
            assert info['ord2'] == ord_2(p)
            assert info['alpha_L'] == alpha_lucas(p)
            assert info['tau'] == vp_support_period(p)

    def test_p1_mod4_mostly_no_alpha_L(self):
        """For p ≡ 1 (mod 4), alpha_L is None for most primes."""
        table = prime_classification_table(100)
        class1 = {p: info for p, info in table.items()
                   if info['mod4'] == 1}
        no_alpha = sum(1 for info in class1.values() if info['alpha_L'] is None)
        # From data: 9 out of 11 primes p≡1(4) < 100 have alpha_L = None
        assert no_alpha >= len(class1) // 2, (
            f"Expected most p≡1(4) to have alpha_L=None, got {no_alpha}/{len(class1)}"
        )


# ── D_mersenne_component and D_lucas_component ────────────────────────────

class TestChannelComponents:

    def test_product_equals_D_for_n_1_to_50(self):
        """D_M(n) * D_L(n) must equal D(n) for all n."""
        for n in range(1, 51):
            D = arithmetic_denominator(n)[0]
            dm = D_mersenne_component(n)
            dl = D_lucas_component(n)
            assert dm * dl == D, f"n={n}: D_M={dm}, D_L={dl}, product={dm*dl}, D={D}"

    def test_both_positive(self):
        """Both channels must produce positive integers."""
        for n in range(1, 51):
            assert D_mersenne_component(n) >= 1
            assert D_lucas_component(n) >= 1

    def test_known_mersenne_values(self):
        """Check specific D_M values from the research output."""
        assert D_mersenne_component(7) == 15
        assert D_mersenne_component(8) == 7
        assert D_mersenne_component(3) == 3
        assert D_mersenne_component(1) == 1

    def test_known_lucas_values(self):
        """Check specific D_L values from the research output."""
        assert D_lucas_component(7) == 1
        assert D_lucas_component(8) == 4
        assert D_lucas_component(2) == 2
        assert D_lucas_component(1) == 1

    def test_not_multiplicative_mersenne(self):
        """D_M is NOT multiplicative: D_M(6) != D_M(2)*D_M(3)."""
        # gcd(2,3)=1 but D_M(6) != D_M(2)*D_M(3)
        dm6 = D_mersenne_component(6)
        dm2_dm3 = D_mersenne_component(2) * D_mersenne_component(3)
        assert dm6 != dm2_dm3, "D_M should NOT be multiplicative"

    def test_not_multiplicative_lucas(self):
        """D_L is NOT multiplicative: D_L(6) != D_L(2)*D_L(3)."""
        dl6 = D_lucas_component(6)
        dl2_dl3 = D_lucas_component(2) * D_lucas_component(3)
        assert dl6 != dl2_dl3, "D_L should NOT be multiplicative"

    def test_invalid_n_raises_mersenne(self):
        with pytest.raises(ValueError):
            D_mersenne_component(0)

    def test_invalid_n_raises_lucas(self):
        with pytest.raises(ValueError):
            D_lucas_component(0)


# ── dirichlet_channel_partial ──────────────────────────────────────────────

class TestDirichletChannelPartial:

    def test_single_term_both_channels(self):
        """At N=1, both channels give D_X(1)/1^s = 1."""
        assert dirichlet_channel_partial('mersenne', 1, 2.0) == 1.0
        assert dirichlet_channel_partial('lucas', 1, 2.0) == 1.0

    def test_monotone_in_N(self):
        """Partial sums are monotonically nondecreasing."""
        for channel in ('mersenne', 'lucas'):
            prev = 0.0
            for N in range(1, 31):
                val = dirichlet_channel_partial(channel, N, 4.0)
                assert val >= prev - 1e-12, f"{channel}, N={N}: {val} < {prev}"
                prev = val

    def test_mersenne_dominates_lucas(self):
        """Mersenne channel has larger partial sums than Lucas channel."""
        for s in [2.0, 3.0, 4.0]:
            sm = dirichlet_channel_partial('mersenne', 50, s)
            sl = dirichlet_channel_partial('lucas', 50, s)
            assert sm > sl, f"s={s}: Mersenne {sm} should dominate Lucas {sl}"

    def test_invalid_channel_raises(self):
        with pytest.raises(ValueError):
            dirichlet_channel_partial('invalid', 10, 2.0)

    def test_invalid_N_raises(self):
        with pytest.raises(ValueError):
            dirichlet_channel_partial('mersenne', 0, 2.0)

    def test_sum_product_vs_full(self):
        """The product of channel sums is NOT equal to the full sum (non-multiplicativity)."""
        s = 4.0
        N = 30
        sm = dirichlet_channel_partial('mersenne', N, s)
        sl = dirichlet_channel_partial('lucas', N, s)
        full = dirichlet_partial(N, s)
        # Product of channel sums != full Dirichlet sum (they're sums not products)
        # But the ratio should be finite and positive
        assert sm > 0 and sl > 0 and full > 0
