"""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,
)


# ── 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)
