"""Sequential six-state Sieve, ported from the HPC graph measurement.

The C engine builds an :class:`HPCGraph` and runs a column-first sequential
measurement with real back-action.  Only the following semantics are used by
the Q2_K pipeline, so V2 implements them directly on a chain:

* local Born probabilities ``|amplitude|^2``;
* a multiplicative neighbour mix with ``same_penalty`` for equal bins;
* a log-domain slack filter;
* deterministic MAP collapse (the production setting);
* back-action that biases still-unmeasured neighbours away from the outcome.

This is mathematically identical to the production path while removing the
graph allocation, edge bookkeeping and triality view machinery.
"""

from __future__ import annotations

import numpy as np


def sieve_measure_chain(
    amplitudes: np.ndarray,
    edges: list[tuple[int, int]],
    same_penalty: float = 0.85,
    log_slack: float = 1.5,
    deterministic: bool = True,
    rng: np.random.Generator | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    amps = np.array(amplitudes, dtype=np.float64, copy=True)
    n_sites, n_states = amps.shape
    if n_states != 6:
        raise ValueError("sieve expects six symbolic states per site")
    adj: list[list[int]] = [[] for _ in range(n_sites)]
    for eid, (a, b) in enumerate(edges):
        adj[a].append(eid)
        adj[b].append(eid)
    alive = np.ones(len(edges), dtype=bool)
    marginals = np.zeros((n_sites, 6), dtype=np.float64)
    measured = np.zeros(n_sites, dtype=np.int64)
    generator = rng or np.random.default_rng(271828)

    for site in range(n_sites - 1, -1, -1):
        local = amps[site] * amps[site]
        total = float(local.sum())
        probs0 = local / total if total > 1.0e-30 else np.full(6, 1.0 / 6.0)
        log_score = np.log(probs0 + 1.0e-30)

        for eid in adj[site]:
            if not alive[eid]:
                continue
            a, b = edges[eid]
            partner = b if a == site else a
            q = amps[partner] * amps[partner]
            qt = float(q.sum())
            q = q / qt if qt > 1.0e-30 else np.full(6, 1.0 / 6.0)
            mix = 1.0 - (1.0 - same_penalty) * q
            log_score += np.log(mix + 1.0e-30)

        best = float(log_score.max())
        probs = np.where(best - log_score > log_slack, 0.0, np.exp(log_score - best))
        total = float(probs.sum())
        probs = probs / total if total > 1.0e-30 else np.full(6, 1.0 / 6.0)
        marginals[site] = probs

        if deterministic:
            outcome = int(np.argmax(probs))
        else:
            outcome = int(np.searchsorted(np.cumsum(probs), generator.random(), side="right"))
            outcome = min(outcome, 5)
        measured[site] = outcome

        amps[site] = 0.0
        amps[site, outcome] = 1.0
        for eid in adj[site]:
            if not alive[eid]:
                continue
            a, b = edges[eid]
            partner = b if a == site else a
            amps[partner, outcome] *= same_penalty
            norm = float(np.linalg.norm(amps[partner]))
            if norm > 1.0e-30:
                amps[partner] /= norm
        for eid in adj[site]:
            alive[eid] = False
        adj[site] = []

    return marginals, measured


def block_chain_edges(n_blocks: int) -> list[tuple[int, int]]:
    """The production two-site-per-block chain, including cross-block links."""
    edges: list[tuple[int, int]] = []
    for i in range(n_blocks):
        s0, s1 = 2 * i, 2 * i + 1
        edges.append((s0, s1))
        if i + 1 < n_blocks:
            edges.append((s0, 2 * (i + 1)))
            edges.append((s1, 2 * (i + 1) + 1))
    return edges


def subblock_chain_edges(n_subblocks: int) -> list[tuple[int, int]]:
    return [(i, i + 1) for i in range(n_subblocks - 1)]