"""Exact finite-orbit computations for second-order recurrences.

The recurrence is

    x[n + 2] = a*x[n + 1] + b*x[n] (mod m).

For b coprime to m, the state map on pairs is invertible, so every orbit is
purely periodic.  Fibonacci, Lucas, and Pell all have b = 1.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
from fractions import Fraction
from math import gcd
from typing import Iterable


@dataclass(frozen=True)
class SequenceSpec:
    name: str
    x0: int
    x1: int
    a: int
    b: int
    definition: str


@dataclass(frozen=True)
class OrbitStats:
    modulus: int
    period: int
    preperiod: int
    residue_count: int
    missing_count: int
    residue_fraction: str
    residue_proportion: float
    returned_to_start: bool
    residues: tuple[int, ...] | None = None


FIBONACCI = SequenceSpec(
    name="Fibonacci",
    x0=0,
    x1=1,
    a=1,
    b=1,
    definition="F(0)=0, F(1)=1, F(n+2)=F(n+1)+F(n)",
)

LUCAS = SequenceSpec(
    name="Lucas",
    x0=2,
    x1=1,
    a=1,
    b=1,
    definition="L(0)=2, L(1)=1, L(n+2)=L(n+1)+L(n)",
)

PELL = SequenceSpec(
    name="Pell",
    x0=0,
    x1=1,
    a=2,
    b=1,
    definition="P(0)=0, P(1)=1, P(n+2)=2P(n+1)+P(n)",
)

SEQUENCES = (FIBONACCI, LUCAS, PELL)

# Larger exponents are used for the small primes; the largest state orbit in
# this plan remains comfortably enumerable on a laptop.
EXPERIMENT_PLAN: dict[int, int] = {
    2: 10,
    3: 8,
    5: 7,
    7: 6,
    11: 5,
    13: 5,
    19: 4,
    31: 4,
}


def next_state(state: tuple[int, int], spec: SequenceSpec, modulus: int) -> tuple[int, int]:
    """Advance one pair state modulo ``modulus``."""

    x, y = state
    return y, (spec.a * y + spec.b * x) % modulus


def orbit_statistics(
    spec: SequenceSpec,
    modulus: int,
    *,
    residue_list_limit: int = 400,
) -> OrbitStats:
    """Enumerate the exact state orbit and count its distinct first coordinates.

    When gcd(b, modulus) = 1, the state map is a permutation and the starting
    pair must return with no preperiod.  This fast path needs no set of states.
    For a noninvertible altered recurrence, a state dictionary is used so the
    preperiod and eventual period are reported correctly.
    """

    if modulus < 2:
        raise ValueError("modulus must be at least 2")

    start = (spec.x0 % modulus, spec.x1 % modulus)
    observed = bytearray(modulus)
    residue_count = 0

    def observe(value: int) -> None:
        nonlocal residue_count
        if observed[value] == 0:
            observed[value] = 1
            residue_count += 1

    if gcd(spec.b, modulus) == 1:
        state = start
        period = 0
        while True:
            observe(state[0])
            state = next_state(state, spec, modulus)
            period += 1
            if state == start:
                break
            if period > modulus * modulus:
                raise RuntimeError("invertible orbit exceeded the finite state space")
        preperiod = 0
        returned_to_start = True
    else:
        state = start
        first_seen: dict[tuple[int, int], int] = {}
        step = 0
        while state not in first_seen:
            first_seen[state] = step
            observe(state[0])
            state = next_state(state, spec, modulus)
            step += 1
        preperiod = first_seen[state]
        period = step - preperiod
        returned_to_start = state == start

    fraction = Fraction(residue_count, modulus)
    residues = None
    if modulus <= residue_list_limit:
        residues = tuple(index for index, flag in enumerate(observed) if flag)

    return OrbitStats(
        modulus=modulus,
        period=period,
        preperiod=preperiod,
        residue_count=residue_count,
        missing_count=modulus - residue_count,
        residue_fraction=str(fraction),
        residue_proportion=float(fraction),
        returned_to_start=returned_to_start,
        residues=residues,
    )


def run_experiments(
    sequences: Iterable[SequenceSpec] = SEQUENCES,
    plan: dict[int, int] = EXPERIMENT_PLAN,
    *,
    residue_list_limit: int = 400,
) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
    """Run the full experiment plan and return records plus small residue sets."""

    records: list[dict[str, object]] = []
    small_sets: list[dict[str, object]] = []

    for spec in sequences:
        for prime, max_k in plan.items():
            prior_period: int | None = None
            prior_count: int | None = None
            for k in range(1, max_k + 1):
                modulus = prime**k
                stats = orbit_statistics(
                    spec,
                    modulus,
                    residue_list_limit=residue_list_limit,
                )
                record = {
                    "sequence": spec.name,
                    "definition": spec.definition,
                    "prime": prime,
                    "k": k,
                    **{key: value for key, value in asdict(stats).items() if key != "residues"},
                    "period_growth": None if prior_period is None else stats.period / prior_period,
                    "residue_growth": None if prior_count is None else stats.residue_count / prior_count,
                }
                records.append(record)
                if stats.residues is not None:
                    small_sets.append(
                        {
                            "sequence": spec.name,
                            "prime": prime,
                            "k": k,
                            "modulus": modulus,
                            "residue_count": stats.residue_count,
                            "residues": list(stats.residues),
                        }
                    )
                prior_period = stats.period
                prior_count = stats.residue_count

    return records, small_sets


def direct_terms(spec: SequenceSpec, modulus: int, length: int) -> list[int]:
    """Return a fixed number of terms, used only for independent small tests."""

    state = (spec.x0 % modulus, spec.x1 % modulus)
    values: list[int] = []
    for _ in range(length):
        values.append(state[0])
        state = next_state(state, spec, modulus)
    return values