"""Analytic module: prime decomposition of D(n) and its Dirichlet series.

The arithmetic denominator D(n) = gcd(M(n+1), (n+1)!) * gcd(L(n+1), (n+1)!)
admits a prime factorization D(n) = prod p^{v_p(D(n))} whose p-adic
valuations v_p exhibit striking periodicity.

This module provides:
- Prime factorization of D(n) and the full D-sequence
- Divisor count function sigma_0
- Normalized growth analysis D(n) / (sigma_0(n+1) * ln(n))
- Dirichlet partial sums sum_{n=1}^{N} D(n) / n^s
- Euler factor sequences v_p(D(n)) for fixed primes p
- Abscissa of convergence estimation for the Dirichlet series
- Theorem Z: support periodicity of v_p(D(n))

Key observations:
    v_2(D(n)): period 3 pattern [0, 1, 0, 0, 1, 0, 0, 2, 0, ...]
    v_p(D(n)): for prime p, nonzero only when (p-1) | n (approximately)
    D(n) is NOT multiplicative, but log D decomposes additively over primes

Theorem Z (Support Periodicity of p-adic valuations):
    For odd prime p, v_p(D(n)) > 0 if and only if n+1 >= p and either:
      (a) ord_2(p) | (n+1), i.e. p divides the Mersenne number M(n+1), or
      (b) alpha_L(p) | (n+1) and (n+1)/alpha_L(p) is odd, i.e. p divides L(n+1).
    For p = 2, v_2(D(n)) > 0 iff 3 | (n+1) and n >= 1.

    The support period tau_p is the minimal period of the indicator function
    I(k) = [ord_2(p)|k OR (alpha_L(p)|k AND k/alpha_L(p) is odd)].
    In general tau_p | lcm(ord_2(p), 2*alpha_L(p)).

    Proof sketch: D(n) = gcd(M(n+1),(n+1)!) * gcd(L(n+1),(n+1)!).
    The p-adic valuation decomposes as v_p(D(n)) = min(v_p(M(n+1)), v_p((n+1)!))
    + min(v_p(L(n+1)), v_p((n+1)!)). The first term is nonzero iff p | M(n+1)
    AND v_p((n+1)!) >= 1 (i.e. n+1 >= p). Since M(n+1) = 2^{n+1}-1, we have
    p | M(n+1) iff ord_2(p) | (n+1). The second term is nonzero iff p | L(n+1)
    AND n+1 >= p. By the Lucas divisibility law, for odd prime p, p | L(k) iff
    k/alpha_L(p) is a positive odd integer.
"""

import math
from collections import Counter
from typing import Dict, List, Optional

from .coupling import arithmetic_denominator
from .kaleidoscope import lucas, prime_factors


def prime_factorization_D(n: int) -> Dict[int, int]:
    """Return the prime factorization of D(n) as {p: v_p}.

    Parameters:
        n: positive integer

    Returns:
        Dictionary mapping prime p to its exponent v_p(D(n)).
        Returns {} if D(n) = 1.

    Examples:
        >>> prime_factorization_D(5)
        {2: 1, 3: 4}
        >>> prime_factorization_D(8)
        {2: 2, 7: 1}
    """
    if n < 1:
        raise ValueError(f"n must be >= 1, got {n}")
    D = arithmetic_denominator(n)[0]
    if D == 1:
        return {}
    factors = prime_factors(D)
    return dict(Counter(factors))


def D_sequence(N: int) -> List[int]:
    """Return [D(1), D(2), ..., D(N)].

    Parameters:
        N: positive integer, length of sequence

    Returns:
        List of D(n) values for n = 1, 2, ..., N.
    """
    if N < 1:
        raise ValueError(f"N must be >= 1, got {N}")
    return [arithmetic_denominator(n)[0] for n in range(1, N + 1)]


def divisor_count(n: int) -> int:
    """sigma_0(n): count of positive divisors of n.

    Parameters:
        n: positive integer

    Returns:
        Number of divisors of n.

    Examples:
        >>> divisor_count(12)
        6
        >>> divisor_count(7)
        2
    """
    if n < 1:
        raise ValueError(f"n must be >= 1, got {n}")
    count = 0
    # Only iterate up to sqrt(n)
    isqrt = int(math.isqrt(n))
    for d in range(1, isqrt + 1):
        if n % d == 0:
            count += 1
            if d != n // d:
                count += 1
    return count


def D_growth_normalized(N: int) -> List[float]:
    """Normalized growth ratio D(n) / (sigma_0(n+1) * ln(n)) for n=2..N.

    The conjectured asymptotic D(n) ~ sigma_0(n+1) * ln(n) would make
    this ratio converge to a constant.  In practice the ratio is highly
    variable because D(n) = 1 for many n (when M(n+1) and L(n+1) are
    coprime to (n+1)!).

    Parameters:
        N: upper bound (must be >= 2)

    Returns:
        List of ratios for n = 2, 3, ..., N (length N-1).
    """
    if N < 2:
        raise ValueError(f"N must be >= 2, got {N}")
    ratios = []
    for n in range(2, N + 1):
        D = arithmetic_denominator(n)[0]
        s0 = divisor_count(n + 1)
        denom = s0 * math.log(n)
        ratios.append(D / denom)
    return ratios


def dirichlet_partial(N: int, s: float) -> float:
    """Partial sum of the Dirichlet series: sum_{n=1}^{N} D(n) / n^s.

    Parameters:
        N: number of terms (>= 1)
        s: real exponent

    Returns:
        The partial sum as a float.

    Notes:
        For s large enough (s > abscissa of convergence), this converges
        as N -> infinity.  The abscissa depends on the growth rate of D(n).
    """
    if N < 1:
        raise ValueError(f"N must be >= 1, got {N}")
    total = 0.0
    for n in range(1, N + 1):
        D = arithmetic_denominator(n)[0]
        total += D / (n ** s)
    return total


def euler_factor_sequence(p: int, N: int) -> List[int]:
    """Return [v_p(D(1)), v_p(D(2)), ..., v_p(D(N))] for prime p.

    The p-adic valuation v_p(D(n)) measures how many times prime p
    divides D(n).  These sequences exhibit periodic or quasi-periodic
    structure:
        v_2: period 3 in the nonzero positions
        v_3: complex quasi-periodic structure
        v_5: period 4 pattern
        v_p (p >= 7): nonzero roughly every (p-1) terms

    Parameters:
        p: a prime number
        N: sequence length (>= 1)

    Returns:
        List of p-adic valuations of D(n) for n = 1..N.
    """
    if N < 1:
        raise ValueError(f"N must be >= 1, got {N}")
    if p < 2:
        raise ValueError(f"p must be a prime >= 2, got {p}")
    result = []
    for n in range(1, N + 1):
        D = arithmetic_denominator(n)[0]
        v = 0
        if D > 0:
            temp = D
            while temp % p == 0:
                v += 1
                temp //= p
        result.append(v)
    return result


def abscissa_estimate(N: int) -> float:
    """Estimate the abscissa of convergence of sum D(n)/n^s.

    Uses the Cauchy-Hadamard criterion: the abscissa sigma_c satisfies
        sigma_c = limsup_{n->inf} ln(D(n)) / ln(n).

    Since D(n) can be 1 for many n, and occasionally very large, we
    compute the running maximum of ln(D(n))/ln(n) over n = 2..N.

    Parameters:
        N: number of terms to use (>= 2)

    Returns:
        Estimated abscissa of convergence (float).
    """
    if N < 2:
        raise ValueError(f"N must be >= 2, got {N}")
    max_ratio = 0.0
    for n in range(2, N + 1):
        D = arithmetic_denominator(n)[0]
        if D > 1:
            ratio = math.log(D) / math.log(n)
            if ratio > max_ratio:
                max_ratio = ratio
        # D=1 contributes ratio=0, doesn't affect limsup
    return max_ratio


# ── Theorem Z: support periodicity of v_p(D(n)) ──────────────────────────────

def ord_2(p: int) -> Optional[int]:
    """Multiplicative order of 2 modulo p.

    The smallest positive integer k such that 2^k = 1 (mod p).
    Returns None for p = 2 (since 2 mod 2 = 0, order is undefined).

    Parameters:
        p: an odd prime

    Returns:
        The multiplicative order, or None if p = 2.

    Examples:
        >>> ord_2(3)
        2
        >>> ord_2(7)
        3
        >>> ord_2(11)
        10
    """
    if p < 2:
        raise ValueError(f"p must be a prime >= 2, got {p}")
    if p == 2:
        return None
    k = 1
    val = 2 % p
    while val != 1:
        k += 1
        val = (val * 2) % p
    return k


def alpha_lucas(p: int) -> Optional[int]:
    """Rank of apparition of prime p in the Lucas sequence.

    The smallest positive integer k such that p divides L(k).
    Returns None if p never divides any Lucas number (this occurs
    for primes p = 5 (mod 10) that are congruent to +/-1 mod 5,
    specifically when the Legendre symbol (5/p) = 1 and the
    Wall-Sun-Sun conditions exclude Lucas divisibility).

    Parameters:
        p: a prime >= 2

    Returns:
        The rank of apparition, or None if p never divides L(k).

    Examples:
        >>> alpha_lucas(2)
        3
        >>> alpha_lucas(3)
        2
        >>> alpha_lucas(7)
        4
        >>> alpha_lucas(5) is None
        True
    """
    if p < 2:
        raise ValueError(f"p must be a prime >= 2, got {p}")
    # Search up to a generous bound; for primes p, alpha_L(p) <= 2(p+1)
    bound = max(4 * (p + 1), 200)
    for k in range(1, bound):
        if lucas(k) % p == 0:
            return k
    return None


def vp_support_period(p: int) -> int:
    """Support period tau_p of the p-adic valuation v_p(D(n)).

    Theorem Z: the binary sequence [v_p(D(n)) > 0] is eventually
    periodic with minimal period tau_p.  For p = 2, tau_2 = 3.
    For odd p with alpha_L(p) undefined, tau_p = ord_2(p).
    For odd p with both ord_2(p) and alpha_L(p) defined, tau_p is
    the minimal period of the indicator
        I(k) = [ord_2(p) | k] OR [alpha_L(p) | k AND k/alpha_L(p) odd].

    Parameters:
        p: a prime >= 2

    Returns:
        The support period tau_p.

    Examples:
        >>> vp_support_period(2)
        3
        >>> vp_support_period(5)
        4
        >>> vp_support_period(11)
        5
    """
    if p < 2:
        raise ValueError(f"p must be a prime >= 2, got {p}")
    if p == 2:
        return 3  # alpha_L(2) = 3; M(n+1) always odd

    o2 = ord_2(p)
    aL = alpha_lucas(p)

    if aL is None:
        # Only Mersenne contributes
        return o2

    # Full period is lcm(ord_2(p), 2*alpha_L(p))
    full_period = math.lcm(o2, 2 * aL)

    # Build one full period of the indicator and find minimal sub-period
    indicator = []
    for j in range(1, full_period + 1):
        mersenne_hit = (j % o2 == 0)
        lucas_hit = (j % aL == 0) and ((j // aL) % 2 == 1)
        indicator.append(1 if (mersenne_hit or lucas_hit) else 0)

    # Find minimal period dividing full_period
    for T in range(1, full_period + 1):
        if full_period % T != 0:
            continue
        valid = True
        for i in range(full_period):
            if indicator[i] != indicator[i % T]:
                valid = False
                break
        if valid:
            return T

    return full_period  # fallback (should not reach here)


# ── Fast CRT-based support period and pattern ─────────────────────────────────

def _alpha_lucas_fast(p: int) -> Optional[int]:
    """Rank of apparition of prime p in Lucas sequence, computed mod p.

    Uses the recurrence L(k) = L(k-1) + L(k-2) mod p to avoid
    computing enormous Lucas numbers.  O(p) time, O(1) space.

    Parameters:
        p: a prime >= 2

    Returns:
        Smallest k >= 1 with p | L(k), or None if no such k exists
        within the search bound 2(p+1).
    """
    if p == 2:
        return 3
    if p == 5:
        return None  # 5 never divides a Lucas number
    # L(0)=2, L(1)=1, L(k) = L(k-1) + L(k-2)
    bound = max(2 * (p + 1), 200)
    prev, curr = 2 % p, 1 % p  # L(0) mod p, L(1) mod p
    if curr == 0:
        return 1
    for k in range(2, bound):
        prev, curr = curr, (prev + curr) % p
        if curr == 0:
            return k
    return None


def vp_support_period_fast(p: int) -> int:
    """Fast CRT-based computation of the support period tau_p.

    The support indicator I(k) = [a|k] OR [b|k AND (k/b) odd] where
    a = ord_2(p) and b = alpha_lucas(p).  This is a union of two
    arithmetic progressions:
        M = {k : a|k}       — period a
        L = {b, 3b, 5b, ...} — period 2b

    The combined period divides lcm(a, 2b).  We compute the indicator
    over exactly one period of lcm(a, 2b) and find the minimal divisor
    T such that the indicator is T-periodic.

    Special-case optimizations:
    - p = 2: tau_2 = 3 (only Lucas contributes; M(n+1) is always odd)
    - alpha_lucas(p) = None: tau_p = ord_2(p) (only Mersenne contributes)
    - alpha_lucas(p) is not None and a | b: if b/a is odd, then every
      Lucas position is also a Mersenne position, so L subset M and
      tau_p = a.  Otherwise tau_p = lcm(a, 2b) / gcd(...) — use general
      algorithm.

    Parameters:
        p: a prime >= 2

    Returns:
        The support period tau_p.

    Examples:
        >>> vp_support_period_fast(2)
        3
        >>> vp_support_period_fast(5)
        4
        >>> vp_support_period_fast(11)
        5
    """
    if p < 2:
        raise ValueError(f"p must be a prime >= 2, got {p}")
    if p == 2:
        return 3

    a = ord_2(p)
    b = _alpha_lucas_fast(p)

    if b is None:
        # Only Mersenne contributes: I(k) = [a|k], period = a
        return a

    # --- Special case: if a divides b and b/a is odd, then L ⊆ M ---
    # L = {b, 3b, 5b, ...}.  Each element mb (m odd) is a multiple of a
    # iff a|b.  When a|b, every Lucas position is automatically a Mersenne
    # position, so I(k) = [a|k] with period a.
    if b % a == 0:
        # a | b.  Check: is L truly ⊆ M?  L positions are multiples of b
        # with odd quotient.  Since a|b, a divides every multiple of b,
        # so every L-position is an M-position.  Period = a.
        return a

    # --- General case ---
    L = math.lcm(a, 2 * b)

    # I(j) = [j ≡ 0 mod a] ∨ [j ≡ b mod 2b]  (1-indexed)
    #
    # Find minimal T | L such that I is T-periodic.
    # Method: start with T = L, then for each prime factor q of L,
    # check if we can divide T by q (i.e. T/q is still a valid period).
    # Repeat until no further reduction is possible.
    return _minimal_period_mersenne_lucas(a, b, L)


def _minimal_period_mersenne_lucas(a: int, b: int, L: int) -> int:
    """Find minimal period of I(j) = [a|j] ∨ [j ≡ b mod 2b] within L.

    Uses prime-factor reduction: start with T = L, then try dividing
    by each prime factor q.  For candidate = T/q, verify periodicity
    by checking I(j) == I(j + candidate) for j = 1..T-candidate.
    Since I(j) uses only modular arithmetic, each evaluation is O(1).
    """
    def _ind(j):
        return (j % a == 0) or (j % (2 * b) == b)

    # Factorize L
    prime_factors_of_L = []
    temp = L
    d = 2
    while d * d <= temp:
        while temp % d == 0:
            if not prime_factors_of_L or prime_factors_of_L[-1] != d:
                prime_factors_of_L.append(d)
            temp //= d
        d += 1
    if temp > 1:
        prime_factors_of_L.append(temp)

    T = L
    changed = True
    while changed:
        changed = False
        for q in prime_factors_of_L:
            if T % q != 0:
                continue
            candidate = T // q
            # Check: is I periodic with period candidate?
            # Since I has period T, check I(j) == I(j + candidate)
            # for all j in [1..T - candidate].  (Equivalently, check
            # all residue classes mod candidate within [1..T].)
            valid = True
            for j in range(1, T - candidate + 1):
                if _ind(j) != _ind(j + candidate):
                    valid = False
                    break
            if valid:
                T = candidate
                changed = True
                break  # restart with smaller T

    return T


def _sorted_divisors(n: int) -> List[int]:
    """Return sorted list of all positive divisors of n."""
    divs = []
    isqrt = int(math.isqrt(n))
    for d in range(1, isqrt + 1):
        if n % d == 0:
            divs.append(d)
            if d != n // d:
                divs.append(n // d)
    divs.sort()
    return divs


def support_pattern(p: int, N: int) -> "numpy.ndarray":
    """Binary support pattern of v_p(D(n)) for n = 1, ..., N.

    Returns a numpy array of 0/1 indicators where entry i (0-indexed)
    is 1 iff v_p(D(i+1)) > 0.  Uses the analytic criterion from
    Theorem Z rather than computing D(n) directly, so this is O(N)
    with tiny constant factor.

    For n+1 < p, the indicator is always 0 (since p cannot divide
    (n+1)! when n+1 < p, so v_p(D(n)) = 0 regardless of Mersenne/Lucas).

    Parameters:
        p: a prime >= 2
        N: length of pattern (>= 1)

    Returns:
        numpy ndarray of shape (N,) with dtype uint8, entries in {0, 1}.

    Examples:
        >>> import numpy as np
        >>> pat = support_pattern(2, 10)
        >>> list(pat)
        [0, 1, 0, 0, 1, 0, 0, 1, 0, 0]
    """
    import numpy as np

    if p < 2:
        raise ValueError(f"p must be a prime >= 2, got {p}")
    if N < 1:
        raise ValueError(f"N must be >= 1, got {N}")

    result = np.zeros(N, dtype=np.uint8)

    if p == 2:
        # v_2(D(n)) > 0 iff 3 | (n+1) and n >= 1.
        # M(k) = 2^k - 1 is always odd so only Lucas contributes.
        # L(k) mod 2: [0,1,1,0,1,1,...] period 3, zero at k ≡ 0 mod 3.
        # So 2 | L(k) iff 3 | k, giving support at n = 2, 5, 8, ...
        for i in range(N):
            k = i + 2  # k = n+1
            if k % 3 == 0:
                result[i] = 1
        return result

    a = ord_2(p)
    b = _alpha_lucas_fast(p)

    for i in range(N):
        n = i + 1
        k = n + 1  # conditions are on k = n+1
        if k < p:
            continue  # v_p((n+1)!) = 0, so v_p(D(n)) = 0
        # Mersenne hit: a | k
        if k % a == 0:
            result[i] = 1
        # Lucas hit: b | k and k/b odd
        elif b is not None and k % b == 0 and (k // b) % 2 == 1:
            result[i] = 1

    return result


# ── Hecke–Dirichlet bridge: channel decomposition ────────────────────────────

def _sieve_primes(N: int) -> List[int]:
    """Return list of primes less than N via sieve of Eratosthenes."""
    if N < 3:
        return [2] if N > 2 else []
    sieve = [True] * N
    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, i):
                sieve[j] = False
    return [i for i in range(2, N) if sieve[i]]


def prime_classification_table(N: int) -> Dict[int, Dict]:
    """Build a classification table for all primes p < N.

    For each prime p, compute:
    - ord2: multiplicative order of 2 mod p (None for p=2)
    - alpha_L: rank of apparition in Lucas sequence (None if p never divides L(k))
    - chi4: Dirichlet character chi_{-4}(p)
    - tau: support period tau_p of v_p(D(n))
    - mod4: p mod 4

    The Hecke eigenvalue of theta_3^2 on Gamma_0(4) is lambda_p = 1 + chi_{-4}(p):
    - lambda_p = 2 for p ≡ 1 (mod 4)
    - lambda_p = 0 for p ≡ 3 (mod 4)
    - lambda_p = 1 for p = 2

    Key correlations:
    - For p ≡ 1 (mod 4): alpha_L is almost always None (only Mersenne channel
      active), so tau_p = ord_2(p).  Average tau_p is small.
    - For p ≡ 3 (mod 4): both Mersenne and Lucas channels are active,
      tau_p involves lcm(ord_2(p), 2*alpha_L(p)).  Average tau_p is
      dramatically larger due to the lcm interaction.

    Parameters:
        N: upper bound (exclusive) for primes.

    Returns:
        Dictionary {p: {'ord2': ..., 'alpha_L': ..., 'chi4': ..., 'tau': ..., 'mod4': ...}}.

    Examples:
        >>> t = prime_classification_table(10)
        >>> sorted(t.keys())
        [2, 3, 5, 7]
        >>> t[5]['chi4']
        1
        >>> t[3]['mod4']
        3
    """
    from .theta import chi_minus4

    primes = _sieve_primes(N)
    table = {}
    for p in primes:
        o2 = ord_2(p)
        aL = alpha_lucas(p)
        chi = chi_minus4(p)
        tau = vp_support_period(p)
        table[p] = {
            'ord2': o2,
            'alpha_L': aL,
            'chi4': chi,
            'tau': tau,
            'mod4': p % 4,
        }
    return table


def D_mersenne_component(n: int) -> int:
    """Mersenne-channel contribution to D(n).

    D_M(n) = gcd(M(n+1), (n+1)!) where M(k) = 2^k - 1.

    The full arithmetic denominator factorizes as D(n) = D_M(n) * D_L(n),
    separating contributions from primes dividing Mersenne numbers from those
    dividing Lucas numbers.  This factorization holds because M(k) and L(k)
    are coprime for all k (they share no common prime factor).

    Parameters:
        n: positive integer

    Returns:
        The Mersenne-channel component (a positive integer).

    Examples:
        >>> D_mersenne_component(7)
        15
        >>> D_mersenne_component(8)
        7
    """
    if n < 1:
        raise ValueError(f"n must be >= 1, got {n}")
    from .kaleidoscope import mersenne
    k = n + 1
    M = mersenne(k)
    fact_k = math.factorial(k)
    return math.gcd(M, fact_k)


def D_lucas_component(n: int) -> int:
    """Lucas-channel contribution to D(n).

    D_L(n) = gcd(L(n+1), (n+1)!) where L(k) is the k-th Lucas number.

    The full arithmetic denominator factorizes as D(n) = D_M(n) * D_L(n).

    Parameters:
        n: positive integer

    Returns:
        The Lucas-channel component (a positive integer).

    Examples:
        >>> D_lucas_component(7)
        1
        >>> D_lucas_component(8)
        4
    """
    if n < 1:
        raise ValueError(f"n must be >= 1, got {n}")
    k = n + 1
    L = lucas(k)
    fact_k = math.factorial(k)
    return math.gcd(L, fact_k)


def dirichlet_channel_partial(channel: str, N: int, s: float) -> float:
    """Partial sum of a channel Dirichlet series.

    Computes sum_{n=1}^{N} D_X(n) / n^s where X is either the Mersenne
    channel ('mersenne') or the Lucas channel ('lucas').

    Neither channel is multiplicative (the factorial GCD destroys it),
    so these do NOT admit Euler products.  However, the channel separation
    reveals that the Hecke eigenvalue structure (chi_{-4} splitting) governs
    which primes contribute to which channel:
    - Mersenne channel: all odd primes contribute (via ord_2(p) | (n+1))
    - Lucas channel: primarily p ≡ 3 (mod 4) contribute (via alpha_L)
      since alpha_L is None for most p ≡ 1 (mod 4)

    Parameters:
        channel: 'mersenne' or 'lucas'
        N: number of terms (>= 1)
        s: real exponent

    Returns:
        The partial sum as a float.

    Raises:
        ValueError: if channel is not 'mersenne' or 'lucas', or N < 1.

    Examples:
        >>> dirichlet_channel_partial('mersenne', 1, 2.0)
        1.0
        >>> dirichlet_channel_partial('lucas', 1, 2.0)
        1.0
    """
    if channel not in ('mersenne', 'lucas'):
        raise ValueError(f"channel must be 'mersenne' or 'lucas', got {channel!r}")
    if N < 1:
        raise ValueError(f"N must be >= 1, got {N}")

    func = D_mersenne_component if channel == 'mersenne' else D_lucas_component
    total = 0.0
    for n in range(1, N + 1):
        total += func(n) / (n ** s)
    return total
