"""
Dimensional PRNG Tomography v4
==============================

Research prototype for identifying LOW-COMPLEXITY, NON-CRYPTOGRAPHIC
pseudorandom generators from observed outputs.

The framework uses relational transforms to:
  1. annihilate nuisance parameters,
  2. lift nonlinear dynamics into linear feature coordinates,
  3. measure recurrence complexity,
  4. validate models on held-out outputs.

This is not intended to attack cryptographically secure generators.
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Sequence, Iterable
from math import gcd
import secrets


# ---------------------------------------------------------------------------
# Prime-field arithmetic and linear algebra
# ---------------------------------------------------------------------------

def invmod(a: int, p: int) -> int:
    a %= p
    if a == 0:
        raise ZeroDivisionError("zero has no inverse")
    return pow(a, p - 2, p)


def solve_square_mod(A: list[list[int]], b: list[int], p: int):
    """Solve square A x=b over prime field F_p. Returns None if singular/inconsistent."""
    n = len(A)
    M = [[x % p for x in row] + [rhs % p] for row, rhs in zip(A,b)]
    r = 0
    pivots = []
    for c in range(n):
        pivot = next((i for i in range(r,n) if M[i][c] % p), None)
        if pivot is None:
            continue
        M[r], M[pivot] = M[pivot], M[r]
        iv = invmod(M[r][c], p)
        M[r] = [(x*iv) % p for x in M[r]]
        for i in range(n):
            if i != r and M[i][c]:
                f = M[i][c] % p
                M[i] = [(x - f*y) % p for x,y in zip(M[i],M[r])]
        pivots.append(c)
        r += 1
    # detect inconsistency
    for i in range(n):
        if all(M[i][c] == 0 for c in range(n)) and M[i][n] != 0:
            return None
    if len(pivots) != n:
        return None
    x = [0]*n
    for i,c in enumerate(pivots):
        x[c] = M[i][n]
    return x


# ---------------------------------------------------------------------------
# Polynomial recurrence family
# x[n+1] = c0 + c1*x[n] + ... + cd*x[n]^d mod p
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class PolynomialRecovery:
    degree: int
    coefficients: tuple[int, ...] | None
    exact_train: bool
    exact_holdout: bool
    train_transitions: int
    holdout_transitions: int
    method: str


def poly_step(x: int, coeffs: Sequence[int], p: int) -> int:
    out = 0
    power = 1
    for c in coeffs:
        out = (out + c*power) % p
        power = (power*x) % p
    return out


def poly_generate(seed: int, coeffs: Sequence[int], p: int, count: int) -> list[int]:
    xs=[seed % p]
    for _ in range(count-1):
        xs.append(poly_step(xs[-1], coeffs, p))
    return xs


def recover_polynomial_recurrence_relational(
    seq: Sequence[int],
    degree: int,
    p: int,
    train_outputs: int | None = None,
) -> PolynomialRecovery:
    """
    Nonlinear annihilator theorem:

      x[n+1] = c0 + sum_{j=1}^d cj*x[n]^j

    Subtract consecutive equations:

      x[n+1]-x[n] =
        sum_{j=1}^d cj*(x[n]^j - x[n-1]^j)

    c0 disappears. Solve d unknown structural coefficients c1..cd
    from d independent relational equations, then lift back to c0.
    """
    if len(seq) < degree + 2:
        raise ValueError("need at least degree+2 outputs")
    if train_outputs is None:
        train_outputs = degree + 2
    train_outputs=max(train_outputs, degree+2)
    train_outputs=min(train_outputs,len(seq))

    # Try windows inside training set until a nonsingular dxd relational feature matrix appears.
    coeff_tail=None
    chosen=None
    max_start=train_outputs-(degree+2)
    for start in range(max_start+1):
        A=[]
        y=[]
        # equations indexed n=start+1 ... start+d
        for n in range(start+1,start+degree+1):
            row=[]
            for j in range(1,degree+1):
                row.append((pow(seq[n],j,p)-pow(seq[n-1],j,p))%p)
            A.append(row)
            y.append((seq[n+1]-seq[n])%p)
        sol=solve_square_mod(A,y,p)
        if sol is not None:
            coeff_tail=sol
            chosen=start
            break

    if coeff_tail is None:
        return PolynomialRecovery(
            degree,None,False,False,train_outputs-1,max(0,len(seq)-train_outputs),
            "relational annihilator: singular feature geometry"
        )

    # lift c0 back from one original transition
    x0=seq[chosen]
    x1=seq[chosen+1]
    c0=x1
    for j,c in enumerate(coeff_tail, start=1):
        c0=(c0-c*pow(x0,j,p))%p
    coeffs=(c0,*coeff_tail)

    preds=[poly_step(seq[i],coeffs,p) for i in range(len(seq)-1)]
    actual=list(seq[1:])
    train_t=max(0,train_outputs-1)
    exact_train=preds[:train_t]==actual[:train_t]
    exact_hold=preds[train_t:]==actual[train_t:]

    return PolynomialRecovery(
        degree,tuple(coeffs),exact_train,exact_hold,train_t,len(preds)-train_t,
        "relational annihilator + nonlinear feature lift"
    )


def discover_polynomial_degree(
    seq: Sequence[int],
    p: int,
    max_degree: int=6,
    train_outputs: int | None=None,
):
    """Return the simplest degree that is exact on train AND holdout."""
    candidates=[]
    for d in range(1,max_degree+1):
        if len(seq)<d+2:
            break
        r=recover_polynomial_recurrence_relational(seq,d,p,train_outputs)
        candidates.append(r)
        if r.coefficients is not None and r.exact_train and r.exact_holdout:
            return r,candidates
    return None,candidates


# ---------------------------------------------------------------------------
# LFSR / linear-complexity tomography
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class LFSRRecovery:
    linear_complexity: int
    connection: tuple[int,...]
    train_bits: int
    holdout_bits: int
    exact_holdout: bool
    predicted_holdout: tuple[int,...]


def berlekamp_massey(bits: Sequence[int]) -> tuple[int,list[int]]:
    """
    Binary Berlekamp-Massey.
    Returns L and C with recurrence:
      s[n] = C[1]s[n-1] + ... + C[L]s[n-L] mod 2.
    """
    s=[int(x)&1 for x in bits]
    n=len(s)
    C=[0]*n
    B=[0]*n
    C[0]=B[0]=1
    L=0
    m=-1
    for N in range(n):
        d=s[N]
        for i in range(1,L+1):
            d ^= C[i] & s[N-i]
        if d==0:
            continue
        T=C[:]
        shift=N-m
        for j in range(n-shift):
            if B[j]:
                C[j+shift]^=1
        if 2*L <= N:
            L=N+1-L
            B=T
            m=N
    return L,C[:L+1]


def lfsr_predict(history: Sequence[int], connection: Sequence[int], count: int) -> list[int]:
    L=len(connection)-1
    h=[int(x)&1 for x in history]
    if L==0:
        return [0]*count
    out=[]
    for _ in range(count):
        nxt=0
        for i in range(1,L+1):
            nxt ^= connection[i] & h[-i]
        h.append(nxt)
        out.append(nxt)
    return out


def recover_lfsr_with_holdout(bits: Sequence[int], train_bits: int) -> LFSRRecovery:
    if train_bits<=0 or train_bits>len(bits):
        raise ValueError
    train=list(bits[:train_bits])
    hold=list(bits[train_bits:])
    L,C=berlekamp_massey(train)
    pred=lfsr_predict(train,C,len(hold))
    return LFSRRecovery(
        L,tuple(C),train_bits,len(hold),pred==hold,tuple(pred)
    )


def generate_lfsr(
    state: Sequence[int],
    connection: Sequence[int],
    count: int,
) -> list[int]:
    """
    Produce output sequence beginning with supplied initial state.
    connection length = L+1, C[0]=1.
    """
    L=len(connection)-1
    if len(state)<L:
        raise ValueError("state must contain at least L seed bits")
    s=[int(x)&1 for x in state[:L]]
    if not any(s):
        s[0]=1
    while len(s)<count:
        nxt=0
        for i in range(1,L+1):
            nxt ^= connection[i] & s[-i]
        s.append(nxt)
    return s[:count]


# ---------------------------------------------------------------------------
# Xorshift32 model discovery (non-cryptographic family)
# ---------------------------------------------------------------------------

MASK32=(1<<32)-1

def xorshift32_step(x: int, a: int, b: int, c: int) -> int:
    x &= MASK32
    x ^= (x << a) & MASK32
    x ^= (x >> b)
    x ^= (x << c) & MASK32
    return x & MASK32


@dataclass(frozen=True)
class XorShiftRecovery:
    shifts: tuple[int,int,int] | None
    exact: bool
    tested: int


def recover_xorshift32(seq: Sequence[int], max_shift: int=31) -> XorShiftRecovery:
    if len(seq)<3:
        raise ValueError("need >=3 full-state outputs")
    tested=0
    for a in range(1,max_shift+1):
        for b in range(1,max_shift+1):
            y1_pre=None
            for c in range(1,max_shift+1):
                tested+=1
                ok=True
                for x,y in zip(seq[:-1],seq[1:]):
                    if xorshift32_step(x,a,b,c)!=y:
                        ok=False
                        break
                if ok:
                    return XorShiftRecovery((a,b,c),True,tested)
    return XorShiftRecovery(None,False,tested)


# ---------------------------------------------------------------------------
# Signatures / controls
# ---------------------------------------------------------------------------

def first_differences(seq: Sequence[int], p: int) -> list[int]:
    return [(seq[i+1]-seq[i])%p for i in range(len(seq)-1)]


def relational_order_signature(seq: Sequence[int], p: int, depth: int=5):
    rows=[]
    cur=list(seq)
    for d in range(depth):
        if len(cur)<2:
            break
        distinct=len(set(cur))
        rows.append({
            "order":d,
            "length":len(cur),
            "distinct":distinct,
            "zero_fraction":sum(x==0 for x in cur)/len(cur),
        })
        cur=first_differences(cur,p)
    return rows


def random_prime_field_sequence(p: int, count: int) -> list[int]:
    return [secrets.randbelow(p) for _ in range(count)]


def random_bits(count:int) -> list[int]:
    return [secrets.randbits(1) for _ in range(count)]