#!/usr/bin/env python3
# Independent reference for the FRI query-index derivation used by the PQ STARK-verify
# precompile 0x0AE8 (direct SP1/Plonky3 inner FRI/STARK verify). This module implements ONE
# self-contained, constant-free slice of component (d) in
# docs/AERE-STARK-VERIFIER-PORT-SPEC.md: turning Fiat-Shamir challenger output into FRI query
# indices, and the per-layer folding-index walk each query follows.
#
# HONEST SCOPE. This is NOT a STARK verifier and does not verify anything. It reproduces the
# public Plonky3 p3-fri index arithmetic (sample_bits + the arity-2 folding walk). It depends on
# NO secret round constants, so it can be validated offline against hand-computed golden vectors
# and cross-checked against independent implementations (Node and Java). Conformance to a real
# SP1 v6.1.0 trace is [MEASURE] (needs the Poseidon2 challenger + an exported proof); see the
# port spec section 8 for the exact validation command. The top-level 0x0AE8 verifier stays
# fail-closed regardless of this module.
#
# Plonky3 semantics reproduced (mark [VERIFY] against p3-fri / p3-challenger at the pinned SP1
# v6.1.0 Plonky3 revision):
#   - sample_bits(bits): low `bits` bits of a sampled base-field element's canonical u32
#     representative. [VERIFY] LSB (not MSB) and that one BabyBear sample supplies enough bits
#     (log_max_height < 31 for the pinned config).
#   - folding walk: arity 2, index_sibling = index ^ 1, index_pair = index >> 1 per layer.
#     [VERIFY] against p3-fri verifier.rs verify_query.
#   - num_fold_rounds = log_max_height - log_final_poly_len; after that many folds the index
#     collapses into the final polynomial's domain [0, 2^log_final_poly_len). [VERIFY] the final
#     poly length for the pinned config.

from dataclasses import dataclass, asdict
from typing import List

# BabyBear prime, for range-checking a canonical field representative fed to sample_bits.
BABYBEAR_P = 2013265921  # 2^31 - 2^27 + 1 = 0x78000001


def sample_bits(canonical_u32: int, bits: int) -> int:
    """Reduce a sampled BabyBear field element (its canonical u32 representative, in [0, p)) to a
    query index in [0, 2^bits) by taking the low `bits` bits. Mirrors Plonky3
    CanSampleBits::sample_bits for a DuplexChallenger over BabyBear.

    [VERIFY] Plonky3 masks the LOW bits of as_canonical_u32(); confirm LSB vs MSB against
    p3-challenger. bits must be < 31 so a single 31-bit BabyBear sample supplies them."""
    if bits < 0:
        raise ValueError("bits must be non-negative")
    if bits >= 31:
        # A single BabyBear element only carries ~31 bits; larger indices would need multiple
        # samples. The pinned SP1 configs keep log_max_height well under 31, so we guard here
        # rather than silently truncate.
        raise ValueError("bits >= 31 needs multiple field samples; not supported for pinned config")
    if canonical_u32 < 0 or canonical_u32 >= BABYBEAR_P:
        raise ValueError("canonical_u32 must be a canonical BabyBear representative in [0, p)")
    return canonical_u32 & ((1 << bits) - 1)


@dataclass(frozen=True)
class FoldStep:
    layer: int          # 0-based FRI folding round
    index: int          # this query's index within the current (folded) domain
    sibling: int        # index ^ 1, the paired evaluation opened alongside `index`
    index_pair: int     # index >> 1, the Merkle-open position in the half-size next layer
    parity: int         # index & 1, which of the sibling pair is this query's own evaluation


def walk(index: int, log_max_height: int, log_final_poly_len: int = 0) -> List[FoldStep]:
    """Produce the per-layer folding-index walk for one FRI query.

    Given the sampled `index` in [0, 2^log_max_height) and the final polynomial log-length, return
    the sequence of FoldStep records, one per folding round. num_fold_rounds =
    log_max_height - log_final_poly_len. This is the index bookkeeping a FRI verifier follows to
    know, at every layer, which sibling pair to open and which folded value is its own. It does NOT
    perform the fold arithmetic or any Merkle/Poseidon2 opening (those are delegated, port spec
    section 5)."""
    if log_max_height < 0 or log_final_poly_len < 0:
        raise ValueError("log heights must be non-negative")
    if log_final_poly_len > log_max_height:
        raise ValueError("log_final_poly_len must not exceed log_max_height")
    if index < 0 or index >= (1 << log_max_height):
        raise ValueError("index out of range for log_max_height")
    num_fold_rounds = log_max_height - log_final_poly_len
    steps: List[FoldStep] = []
    cur = index
    for layer in range(num_fold_rounds):
        steps.append(
            FoldStep(
                layer=layer,
                index=cur,
                sibling=cur ^ 1,
                index_pair=cur >> 1,
                parity=cur & 1,
            )
        )
        cur = cur >> 1
    return steps


def final_index(index: int, log_max_height: int, log_final_poly_len: int = 0) -> int:
    """The terminal index after all folds; must land in the final polynomial domain
    [0, 2^log_final_poly_len)."""
    num_fold_rounds = log_max_height - log_final_poly_len
    return index >> num_fold_rounds


def derive_query_indices(sampled_canonical_u32: List[int], log_max_height: int) -> List[int]:
    """Turn a list of challenger-sampled field representatives (one per query) into query indices.
    In a real verifier the samples come out of the Poseidon2 duplex challenger (delegated); here
    they are supplied so the index reduction can be tested independently of the sponge."""
    return [sample_bits(s, log_max_height) for s in sampled_canonical_u32]


def walk_as_dicts(index: int, log_max_height: int, log_final_poly_len: int = 0):
    return [asdict(s) for s in walk(index, log_max_height, log_final_poly_len)]