//! The registered watermark scheme `wai-mark:tent64s` (extensions/synthesis.md §7.9): the
//! tent pattern, the code and the embedding of `wai-mark:tent64`, on a **tiled** layout, and a
//! detector that finds the grid again after a crop, a small rotation or a capture of a screen.
//!
//! - **Tiles.** The 64 × 64 grid is cut into 2 × 2 tiles of 32 × 32 cells. Every tile holds
//!   the same layout: bit `j` is carried by 4 positions of the tile, so by 16 cells of the
//!   frame, the same count as `tent64`. The mark is then periodic over the frame, with a period
//!   of half its width and half its height, and the four cells of a bit in one tile let the
//!   detector check a candidate alignment against itself, with no key and no known signal.
//! - **Embedding.** §7.5.4 exactly, under this layout: no energy is spent on synchronisation.
//! - **Detection, stage 1.** §7.5.5 under this layout: the grid in place, in the 8 orientations.
//! - **Detection, stage 2**, where stage 1 finds nothing: the tile lattice from the
//!   autocorrelation of a sign plane matched to the cell pattern (exact, by a number-theoretic
//!   transform); the frames folded onto one tile along it; the cells' phase by their energy;
//!   the orientation and the tile's offset by how well the four cells of every bit agree; and
//!   the code, at most [`DECODES`] times.
//!
//! Everything is integer arithmetic, so the detector's result is a function of the frames,
//! the same on every machine and in every implementation of §7.9.

use super::{decode, embed_with, read_word_with, tent, Layout, CODE_BITS, ORIENTATIONS};
use super::{admits, luma, stream_of, CELLS, GRID};
use crate::jcs::jcs;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::sync::OnceLock;

/// The scheme's term (§7.4).
pub const SCHEME: &str = "wai-mark:tent64s";
/// The label of the layout's stream.
pub const STREAM_LABEL: &[u8] = b"wai-mark:tent64s stream";
/// Cells along each side of a tile.
pub const TILE: usize = 32;
/// Positions in a tile.
pub const POSITIONS: usize = TILE * TILE;
/// Positions of a tile that carry each bit.
pub const PER_TILE: usize = 4;
/// The frames stage 2 reads at most, from the first.
pub const STAGE2_FRAMES: usize = 8;
/// The longest side of stage 2's working plane, above which the frames are reduced.
pub const WORKING_SIDE: usize = 1024;
/// The clamp of the robust values, in units of 1/25 of a luma level.
pub const ROBUST_CLAMP: i64 = 100;
/// The lattice candidates stage 2 folds along.
pub const CANDIDATES: usize = 8;
/// The hypotheses stage 2 decodes at most.
pub const DECODES: usize = 4;
/// The search bounds of a lattice vector, as fractions of the plane's side: from 7/20 to 3/4.
pub const SPAN: ((usize, usize), (usize, usize)) = ((7, 20), (3, 4));
/// The largest slope of a lattice vector from its axis, as `|dy| ≤ 3·dx / 20`.
pub const SLOPE: (i64, i64) = (3, 20);
/// The sub-sample unit of a refined lattice vector: 1/64 of a sample.
pub const SUBSAMPLE: i64 = 64;
/// The least and the most bins a cell of the folded tile has along each side.
pub const BINS: (i64, i64) = (4, 8);

/// The scheme's definition: what `tent64`'s pins, with this scheme's layout and detector.
pub fn definition() -> Value {
    let mut d = super::definition();
    let m = d.as_object_mut().expect("an object");
    m.insert("scheme".into(), json!(SCHEME));
    m.insert("specification".into(), json!("extensions/synthesis.md section 7.9"));
    m.insert("stream".into(), json!(std::str::from_utf8(STREAM_LABEL).expect("ASCII")));
    m.insert("tiles".into(), json!([GRID / TILE, GRID / TILE]));
    m.insert("cells_per_tile_bit".into(), json!(PER_TILE));
    m.insert(
        "resync".into(),
        json!({
            "frames": STAGE2_FRAMES,
            "working_side": WORKING_SIDE,
            "robust": [25, 5, ROBUST_CLAMP],
            "span": [[SPAN.0.0, SPAN.0.1], [SPAN.1.0, SPAN.1.1]],
            "slope": [SLOPE.0, SLOPE.1],
            "subsample": SUBSAMPLE,
            "bins": [BINS.0, BINS.1],
            "candidates": CANDIDATES,
            "decodes": DECODES,
        }),
    );
    d
}

/// The definition's bytes: its JCS text.
pub fn definition_text() -> String {
    jcs(&definition()).expect("a definition in the subset")
}

/// The scheme's digest.
pub fn digest() -> [u8; 32] {
    Sha256::digest(definition_text().as_bytes()).into()
}

/// The tile's own layout: the shuffle of its 1024 positions, their signs and the mask.
pub struct TileLayout {
    /// `perm[4·j + k]` is the `k`-th position of bit `j`; the last 32 carry nothing.
    pub perm: Vec<u16>,
    /// Each position's sign, +1 or −1.
    pub sign: Vec<i64>,
    /// The mask over the code's bits.
    pub mask: [u8; CODE_BITS],
}

/// The tile layout, from the scheme's stream.
pub fn tile_layout() -> &'static TileLayout {
    static L: OnceLock<TileLayout> = OnceLock::new();
    L.get_or_init(|| {
        let mut r = stream_of(STREAM_LABEL);
        let mut perm: Vec<u16> = (0..POSITIONS as u16).collect();
        for i in (1..POSITIONS).rev() {
            let j = r.next().expect("an endless stream") as usize % (i + 1);
            perm.swap(i, j);
        }
        let sign: Vec<i64> = (0..POSITIONS).map(|_| if r.next().expect("an endless stream") & 1 == 0 { 1 } else { -1 }).collect();
        let mut mask = [0u8; CODE_BITS];
        for m in mask.iter_mut() {
            *m = (r.next().expect("an endless stream") & 1) as u8;
        }
        TileLayout { perm, sign, mask }
    })
}

/// The frame cell of tile position `p` in tile `t` (0 to 3: the tiles left to right, then top
/// to bottom).
fn cell_of(p: usize, t: usize) -> usize {
    let (row, col) = (p / TILE, p % TILE);
    let (tx, ty) = (t % 2, t / 2);
    GRID * (TILE * ty + row) + TILE * tx + col
}

/// The tiled layout as a frame layout of §7.5.3's shape, so that `tent64`'s embedding and its
/// in-place reading serve this scheme: bit `j`'s 16 cells are its 4 positions in each tile.
pub fn layout() -> &'static Layout {
    static L: OnceLock<Layout> = OnceLock::new();
    L.get_or_init(|| {
        let tl = tile_layout();
        let mut perm = Vec::with_capacity(CELLS);
        for j in 0..CODE_BITS {
            for k in 0..PER_TILE {
                for t in 0..4 {
                    perm.push(cell_of(tl.perm[PER_TILE * j + k] as usize, t) as u16);
                }
            }
        }
        for k in CODE_BITS * PER_TILE..POSITIONS {
            for t in 0..4 {
                perm.push(cell_of(tl.perm[k] as usize, t) as u16);
            }
        }
        let mut sign = vec![0i64; CELLS];
        let mut bit_of_cell = vec![None; CELLS];
        for p in 0..POSITIONS {
            for t in 0..4 {
                sign[cell_of(p, t)] = tl.sign[p];
            }
        }
        for j in 0..CODE_BITS {
            for n in 0..16 {
                bit_of_cell[perm[16 * j + n] as usize] = Some(j as u16);
            }
        }
        Layout { perm, sign, mask: tl.mask, bit_of_cell }
    })
}

/// The marked frames under `payload` (§7.9.3): §7.5.4 under the tiled layout.
pub fn embed(unmarked: &[u8], w: usize, h: usize, frames: usize, payload: &[u8; 8]) -> (Vec<u8>, Vec<super::FrameReport>) {
    embed_with(layout(), unmarked, w, h, frames, payload)
}

/// The floor's SSE under the tiled layout (§7.5.6's floor, this scheme's carrying cells).
pub fn floor_sse(w: usize, h: usize) -> u64 {
    super::floor_sse_with(layout(), w, h)
}

/// What stage 2 found, for the corpus's vectors: the lattice it folded along, in 1/64 of a
/// sample of the working plane, the cells' phase, and the hypothesis that decoded.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Resync {
    /// `[X1, Y1, X2, Y2]`.
    pub lattice: [i64; 4],
    /// The bins of a cell along each side of the folded tile, `(x, y)`.
    pub bins: (usize, usize),
    /// The cells' phase in bins, `(x, y)`.
    pub phase: (usize, usize),
    /// The tile's offset in cells, `(x, y)`.
    pub offset: (usize, usize),
    /// The candidate's index, from 0.
    pub candidate: usize,
    /// The hypothesis's rank among those decoded, from 0.
    pub rank: usize,
    /// Its agreement count.
    pub agreement: u32,
}

/// What the detector read.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Detected {
    pub payload: [u8; 8],
    /// 1, the grid in place; 2, the grid found again.
    pub stage: u8,
    /// The orientation it decoded in: of the frames in stage 1, of the folded tile in stage 2.
    pub orientation: u8,
    /// The bits the code corrected.
    pub errors: usize,
    /// Stage 2's finding.
    pub resync: Option<Resync>,
}

/// The key-free detector of §7.9.4: stage 1, then stage 2. Frames the scheme does not admit
/// carry no mark.
pub fn detect(samples: &[u8], w: usize, h: usize, frames: usize) -> Option<Detected> {
    if !admits(w as u64, h as u64, frames as u64) || samples.len() != w * h * 3 * frames {
        return None;
    }
    let lay = layout();
    for o in 0..ORIENTATIONS {
        if let Some((payload, errors)) = decode(&read_word_with(lay, samples, w, h, frames, o)) {
            return Some(Detected { payload, stage: 1, orientation: o, errors, resync: None });
        }
    }
    stage2(samples, w, h, frames)
}

// ---------------------------------------------------------------------------------------
// Stage 2
// ---------------------------------------------------------------------------------------

/// `n / d` rounded to the nearest integer, halves away from zero, for any non-zero `d`.
fn rnd(n: i64, d: i64) -> i64 {
    let (n, d) = if d < 0 { (-n, -d) } else { (n, d) };
    n.signum() * ((2 * n.abs() + d) / (2 * d))
}

/// The working plane of one frame: its luma, reduced by `k` (§7.9.4 step 1).
fn working_plane(rgb: &[u8], w: usize, h: usize, k: usize) -> (Vec<i64>, usize, usize) {
    let y: Vec<i64> = rgb.as_chunks::<3>().0.iter().map(|p| luma(p[0], p[1], p[2])).collect();
    if k == 1 {
        return (y, w, h);
    }
    let (ww, hh) = (w / k, h / k);
    let kk = (k * k) as i64;
    let mut out = vec![0i64; ww * hh];
    for oy in 0..hh {
        for ox in 0..ww {
            let mut s = 0i64;
            for dy in 0..k {
                for dx in 0..k {
                    s += y[(oy * k + dy) * w + ox * k + dx];
                }
            }
            out[oy * ww + ox] = (s + kk / 2) / kk;
        }
    }
    (out, ww, hh)
}

/// The robust values (step 2): `25·Y` less the sum of the 5 × 5 samples around, the plane's
/// edge repeated, clamped to ±[`ROBUST_CLAMP`].
fn robust(y: &[i64], w: usize, h: usize) -> Vec<i64> {
    // The sum over each row's 5-wide window first, then over 5 rows.
    let cx = |x: i64| x.clamp(0, w as i64 - 1) as usize;
    let cy = |v: i64| v.clamp(0, h as i64 - 1) as usize;
    let mut row = vec![0i64; w * h];
    for r in 0..h {
        for x in 0..w {
            row[r * w + x] = (-2..=2).map(|d| y[r * w + cx(x as i64 + d)]).sum();
        }
    }
    let mut out = vec![0i64; w * h];
    for r in 0..h {
        for x in 0..w {
            let s: i64 = (-2..=2).map(|d| row[cy(r as i64 + d) * w + x]).sum();
            out[r * w + x] = (25 * y[r * w + x] - s).clamp(-ROBUST_CLAMP, ROBUST_CLAMP);
        }
    }
    out
}

/// The sign plane (step 3): the sign of the cell pattern's response, over the samples where
/// its window fits. Returns the plane and its width and height.
fn sign_plane(r: &[i64], w: usize, h: usize) -> (Vec<i64>, usize, usize) {
    let (lx, ly) = ((w / GRID).max(2), (h / GRID).max(2));
    let (wq, hq) = (w - lx + 1, h - ly + 1);
    let fx: Vec<i64> = (0..lx as i64).map(|u| tent(u, lx as i64)).collect();
    let fy: Vec<i64> = (0..ly as i64).map(|u| tent(u, ly as i64)).collect();
    let mut horiz = vec![0i64; wq * h];
    for row in 0..h {
        for x in 0..wq {
            horiz[row * wq + x] = (0..lx).map(|u| fx[u] * r[row * w + x + u]).sum();
        }
    }
    let mut q = vec![0i64; wq * hq];
    for y in 0..hq {
        for x in 0..wq {
            let s: i64 = (0..ly).map(|v| fy[v] * horiz[(y + v) * wq + x]).sum();
            q[y * wq + x] = s.signum();
        }
    }
    (q, wq, hq)
}

// ---- the autocorrelation, exact, by a number-theoretic transform modulo 998244353 ----

const P: u64 = 998_244_353;
const ROOT: u64 = 3;

fn pow_mod(mut b: u64, mut e: u64) -> u64 {
    let mut r = 1u64;
    b %= P;
    while e > 0 {
        if e & 1 == 1 {
            r = r * b % P;
        }
        b = b * b % P;
        e >>= 1;
    }
    r
}

fn ntt(a: &mut [u64], invert: bool) {
    let n = a.len();
    let mut j = 0;
    for i in 1..n {
        let mut bit = n >> 1;
        while j & bit != 0 {
            j ^= bit;
            bit >>= 1;
        }
        j ^= bit;
        if i < j {
            a.swap(i, j);
        }
    }
    let mut len = 2;
    while len <= n {
        let mut wl = pow_mod(ROOT, (P - 1) / len as u64);
        if invert {
            wl = pow_mod(wl, P - 2);
        }
        for i in (0..n).step_by(len) {
            let mut wv = 1u64;
            for k in 0..len / 2 {
                let u = a[i + k];
                let v = a[i + k + len / 2] * wv % P;
                a[i + k] = if u + v >= P { u + v - P } else { u + v };
                a[i + k + len / 2] = if u >= v { u - v } else { u + P - v };
                wv = wv * wl % P;
            }
        }
        len <<= 1;
    }
    if invert {
        let ni = pow_mod(n as u64, P - 2);
        for x in a.iter_mut() {
            *x = *x * ni % P;
        }
    }
}

/// A 2-D transform of a `px` × `py` array, rows then columns.
fn ntt2(a: &mut [u64], px: usize, py: usize, invert: bool) {
    for row in a.chunks_mut(px) {
        ntt(row, invert);
    }
    let mut col = vec![0u64; py];
    for x in 0..px {
        for y in 0..py {
            col[y] = a[y * px + x];
        }
        ntt(&mut col, invert);
        for y in 0..py {
            a[y * px + x] = col[y];
        }
    }
}

/// The autocorrelation of sign planes of `wq` × `hq`, summed over them (step 4): `A(dx, dy)`
/// at `[(dy mod py)·px + (dx mod px)]`.
struct Acf {
    a: Vec<i64>,
    px: usize,
    py: usize,
    wq: usize,
    hq: usize,
}

impl Acf {
    fn new(planes: &[Vec<i64>], wq: usize, hq: usize) -> Acf {
        let (px, py) = ((2 * wq).next_power_of_two(), (2 * hq).next_power_of_two());
        let mut acc = vec![0u64; px * py];
        for q in planes {
            let (mut f, mut g) = (vec![0u64; px * py], vec![0u64; px * py]);
            for y in 0..hq {
                for x in 0..wq {
                    let v = (q[y * wq + x] + P as i64) as u64 % P;
                    f[y * px + x] = v;
                    // g is q reversed through the origin, so f ∗ g is the correlation.
                    g[((py - y) % py) * px + (px - x) % px] = v;
                }
            }
            ntt2(&mut f, px, py, false);
            ntt2(&mut g, px, py, false);
            for i in 0..px * py {
                acc[i] = (acc[i] + f[i] * g[i]) % P;
            }
        }
        ntt2(&mut acc, px, py, true);
        let a = acc.into_iter().map(|v| if v > P / 2 { v as i64 - P as i64 } else { v as i64 }).collect();
        Acf { a, px, py, wq, hq }
    }

    /// `A(dx, dy)`, and 0 outside the plane's offsets.
    fn at(&self, dx: i64, dy: i64) -> i64 {
        if dx.unsigned_abs() as usize >= self.wq || dy.unsigned_abs() as usize >= self.hq {
            return 0;
        }
        self.a[(dy.rem_euclid(self.py as i64) as usize) * self.px + dx.rem_euclid(self.px as i64) as usize]
    }

    /// The pairs of samples at an offset: `(Wq − |dx|)·(Hq − |dy|)`, and 0 outside.
    fn overlap(&self, dx: i64, dy: i64) -> i64 {
        if dx.unsigned_abs() as usize >= self.wq || dy.unsigned_abs() as usize >= self.hq {
            return 0;
        }
        (self.wq as i64 - dx.abs()) * (self.hq as i64 - dy.abs())
    }
}

/// A lattice candidate: `b1 = (dx, dy)` and `b2 = (g, e)`, `g = rnd(−dy·e, dx)`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct Candidate {
    dx: i64,
    dy: i64,
    e: i64,
    g: i64,
}

/// Step 5: the candidates in the order of their score, each more than 2 from every earlier
/// one in `dx`, `dy` or `e`; ties in the order listed, `dx`, then `dy`, then `e`. The score of
/// a candidate is `⌊65536·ΣA / isqrt(Σov)⌋` over its four lattice points `b1`, `b2`,
/// `b1 + b2` and `b1 − b2`, `ov` each point's overlap.
///
/// A suppressed candidate is within 2 of a taken one in all three, so each taken one
/// suppresses at most 125; the first [`KEPT`] in the order therefore hold every candidate the
/// suppression can reach before [`CANDIDATES`] are taken, and only those are kept.
fn candidates(acf: &Acf) -> Vec<Candidate> {
    use std::cmp::Reverse;
    use std::collections::BinaryHeap;
    let (wq, hq) = (acf.wq as i64, acf.hq as i64);
    let ((lo_n, lo_d), (hi_n, hi_d)) = SPAN;
    // ⌈lo_n·n / lo_d⌉ for n ≥ 0.
    let lo = |n: i64| (lo_n as i64 * n + lo_d as i64 - 1) / lo_d as i64;
    let hi = |n: i64| hi_n as i64 * n / hi_d as i64;
    // The heap's top is the worst kept: the least score, then the latest listed.
    let mut heap: BinaryHeap<(Reverse<i64>, u64, Candidate)> = BinaryHeap::new();
    let mut index = 0u64;
    for dx in lo(wq)..=hi(wq) {
        let m = SLOPE.0 * dx / SLOPE.1;
        for dy in -m..=m {
            for e in lo(hq)..=hi(hq) {
                let c = Candidate { dx, dy, e, g: rnd(-dy * e, dx) };
                let pts = [(c.dx, c.dy), (c.g, c.e), (c.dx + c.g, c.dy + c.e), (c.dx - c.g, c.dy - c.e)];
                let a: i64 = pts.iter().map(|&(x, y)| acf.at(x, y)).sum();
                let ov: i64 = pts.iter().map(|&(x, y)| acf.overlap(x, y)).sum();
                let s = (65536 * a).div_euclid((ov as u64).isqrt().max(1) as i64);
                heap.push((Reverse(s), index, c));
                if heap.len() > KEPT {
                    heap.pop();
                }
                index += 1;
            }
        }
    }
    let mut kept: Vec<(i64, u64, Candidate)> = heap.into_iter().map(|(Reverse(s), i, c)| (s, i, c)).collect();
    kept.sort_by(|a, b| b.0.cmp(&a.0).then(a.1.cmp(&b.1)));
    let mut out: Vec<Candidate> = Vec::new();
    for (_, _, c) in kept {
        if out.iter().all(|t| (t.dx - c.dx).abs().max((t.dy - c.dy).abs()).max((t.e - c.e).abs()) > 2) {
            out.push(c);
            if out.len() == CANDIDATES {
                break;
            }
        }
    }
    out
}

/// The candidates kept before suppression: more than 125 times [`CANDIDATES`] less one.
const KEPT: usize = 1024;

/// The sub-sample offset of a peak along one axis, in 1/64: the vertex of the parabola
/// through `l`, `c`, `r`, clamped to ±32; 0 where they make no peak.
fn vertex(l: i64, c: i64, r: i64) -> i64 {
    let d = l - 2 * c + r;
    if d >= 0 {
        return 0;
    }
    rnd(32 * (r - l), -d).clamp(-32, 32)
}

/// Step 6: a vector refined to 1/64 of a sample.
fn refine(acf: &Acf, dx: i64, dy: i64) -> (i64, i64) {
    let c = acf.at(dx, dy);
    (
        SUBSAMPLE * dx + vertex(acf.at(dx - 1, dy), c, acf.at(dx + 1, dy)),
        SUBSAMPLE * dy + vertex(acf.at(dx, dy - 1), c, acf.at(dx, dy + 1)),
    )
}

/// The cell statistics of the folded tile at a phase: `S[b][a]`, row `b`, column `a`; the tile
/// is `32·mx` bins wide and `32·my` high, a cell `mx` by `my`.
fn tile_statistics(f: &[i64], (mx, my): (usize, usize), phase: (usize, usize)) -> Vec<i64> {
    let (nx, ny) = (TILE * mx, TILE * my);
    let tx: Vec<i64> = (0..mx as i64).map(|u| tent(u, mx as i64)).collect();
    let ty: Vec<i64> = (0..my as i64).map(|u| tent(u, my as i64)).collect();
    let mut s = vec![0i64; POSITIONS];
    for b in 0..TILE {
        for a in 0..TILE {
            let mut acc = 0i64;
            for (v, &tyv) in ty.iter().enumerate() {
                let row = (b * my + phase.1 + v) % ny;
                let r: i64 = tx.iter().enumerate().map(|(u, &txu)| txu * f[row * nx + (a * mx + phase.0 + u) % nx]).sum();
                acc += tyv * r;
            }
            s[b * TILE + a] = acc;
        }
    }
    s
}

/// The tile's statistics seen in orientation `o`: `S_o[y][x] = ±S[y2][x2]`, as §7.5.5 maps a
/// frame, negated once for each reversal.
fn oriented(s: &[i64], o: u8) -> Vec<i64> {
    let sg = if ((o & 1) ^ ((o >> 1) & 1)) == 1 { -1 } else { 1 };
    let mut out = vec![0i64; POSITIONS];
    for y in 0..TILE {
        for x in 0..TILE {
            let x1 = if o & 1 != 0 { TILE - 1 - x } else { x };
            let y1 = if o & 2 != 0 { TILE - 1 - y } else { y };
            let (x2, y2) = if o & 4 != 0 { (y1, x1) } else { (x1, y1) };
            out[y * TILE + x] = sg * s[y2 * TILE + x2];
        }
    }
    out
}

/// The value at position `p` of a tile offset by `(x0, y0)` cells.
fn at_offset(s: &[i64], p: usize, x0: usize, y0: usize) -> i64 {
    s[((p / TILE + y0) % TILE) * TILE + (p % TILE + x0) % TILE]
}

/// The agreement of a hypothesis: over every bit, the pairs of its four cells whose signed
/// statistics `ε·S` are both above 0 or both not.
fn agreement(s: &[i64], x0: usize, y0: usize) -> u32 {
    let tl = tile_layout();
    let mut n = 0u32;
    for j in 0..CODE_BITS {
        let c: [bool; PER_TILE] = std::array::from_fn(|k| {
            let p = tl.perm[PER_TILE * j + k] as usize;
            tl.sign[p] * at_offset(s, p, x0, y0) > 0
        });
        for a in 0..PER_TILE {
            for b in a + 1..PER_TILE {
                n += (c[a] == c[b]) as u32;
            }
        }
    }
    n
}

/// The word a hypothesis reads: `T_j = Σ ε·S` over bit `j`'s four positions, its sign, unmasked.
fn word(s: &[i64], x0: usize, y0: usize) -> [u8; CODE_BITS] {
    let tl = tile_layout();
    std::array::from_fn(|j| {
        let t: i64 = (0..PER_TILE)
            .map(|k| {
                let p = tl.perm[PER_TILE * j + k] as usize;
                tl.sign[p] * at_offset(s, p, x0, y0)
            })
            .sum();
        (t > 0) as u8 ^ tl.mask[j]
    })
}

/// Stage 2 (§7.9.4): the grid found again, or `None`.
pub fn stage2(samples: &[u8], w: usize, h: usize, frames: usize) -> Option<Detected> {
    stage2_run(samples, w, h, frames, true).0
}

/// The words stage 2 decodes, in their order, all [`DECODES`] of them whether or not one gives
/// a payload: what §7.9.5's model is checked against.
pub fn stage2_words(samples: &[u8], w: usize, h: usize, frames: usize) -> Vec<[u8; CODE_BITS]> {
    stage2_run(samples, w, h, frames, false).1
}

/// Stage 2's intermediate results, for the corpus to state: every candidate kept, its
/// vectors, its refined lattice, its bins and its phase, and the hypotheses it decodes, each
/// `(agreement, tile, orientation, x0, y0)`, the tile an index into the candidates kept.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Trace {
    pub candidates: Vec<([i64; 4], Resync)>,
    pub decoded: Vec<(u32, usize, u8, usize, usize)>,
}

/// [`Trace`] of the frames' stage 2, whatever stage 1 reads; `None` for frames the scheme does
/// not admit.
pub fn stage2_trace(samples: &[u8], w: usize, h: usize, frames: usize) -> Option<Trace> {
    if !admits(w as u64, h as u64, frames as u64) || samples.len() != w * h * 3 * frames {
        return None;
    }
    Some(stage2_run(samples, w, h, frames, false).2)
}

fn stage2_run(samples: &[u8], w: usize, h: usize, frames: usize, stop: bool) -> (Option<Detected>, Vec<[u8; CODE_BITS]>, Trace) {
    let k = w.max(h).div_ceil(WORKING_SIDE);
    let fb = w * h * 3;
    let used = frames.min(STAGE2_FRAMES);
    let mut rob = Vec::with_capacity(used);
    let mut signs = Vec::with_capacity(used);
    let (mut ww, mut hh, mut wq, mut hq) = (0, 0, 0, 0);
    for f in 0..used {
        let (y, w2, h2) = working_plane(&samples[f * fb..(f + 1) * fb], w, h, k);
        let r = robust(&y, w2, h2);
        let (q, w3, h3) = sign_plane(&r, w2, h2);
        (ww, hh, wq, hq) = (w2, h2, w3, h3);
        rob.push(r);
        signs.push(q);
    }
    let acf = Acf::new(&signs, wq, hq);
    let cands = candidates(&acf);
    // Each candidate's folded tile at its phase.
    let mut tiles: Vec<(Resync, Vec<i64>)> = Vec::new();
    let mut trace = Trace::default();
    for (ci, c) in cands.iter().enumerate() {
        let (x1, y1) = refine(&acf, c.dx, c.dy);
        let (x2, y2) = refine(&acf, c.g, c.e);
        let det = x1 * y2 - x2 * y1;
        if det <= 0 {
            continue;
        }
        let bins = |v: i64| (v / (SUBSAMPLE * TILE as i64)).clamp(BINS.0, BINS.1) as usize;
        let (mx, my) = (bins(x1), bins(y2));
        let (nx, ny) = (TILE * mx, TILE * my);
        // Folded by the parity of the tile each sample falls in, so the tiles can be compared.
        let mut f = vec![vec![0i64; nx * ny]; 4];
        let (sx, sy) = (SUBSAMPLE * nx as i64, SUBSAMPLE * ny as i64);
        for r in &rob {
            for y in 0..hh as i64 {
                for x in 0..ww as i64 {
                    let u = (sx * (y2 * x - x2 * y)).div_euclid(det);
                    let v = (sy * (x1 * y - y1 * x)).div_euclid(det);
                    let class = (u.div_euclid(nx as i64).rem_euclid(2) + 2 * v.div_euclid(ny as i64).rem_euclid(2)) as usize;
                    f[class][v.rem_euclid(ny as i64) as usize * nx + u.rem_euclid(nx as i64) as usize] += r[y as usize * ww + x as usize];
                }
            }
        }
        // The phase whose cells agree most across the tiles: Σ (S² − Σ S_c²), S = Σ S_c.
        let mut best: Option<(i128, (usize, usize), Vec<i64>)> = None;
        for py in 0..my {
            for px in 0..mx {
                let parts: Vec<Vec<i64>> = f.iter().map(|fc| tile_statistics(fc, (mx, my), (px, py))).collect();
                let s: Vec<i64> = (0..POSITIONS).map(|i| parts.iter().map(|p| p[i]).sum()).collect();
                let e: i128 = (0..POSITIONS).map(|i| s[i] as i128 * s[i] as i128 - parts.iter().map(|p| p[i] as i128 * p[i] as i128).sum::<i128>()).sum();
                if best.as_ref().is_none_or(|b| e > b.0) {
                    best = Some((e, (px, py), s));
                }
            }
        }
        let (_, phase, s) = best.expect("at least one phase");
        let res = Resync { lattice: [x1, y1, x2, y2], bins: (mx, my), phase, offset: (0, 0), candidate: ci, rank: 0, agreement: 0 };
        trace.candidates.push(([c.dx, c.dy, c.g, c.e], res));
        tiles.push((res, s));
    }
    // Every hypothesis: a candidate, an orientation and an offset, by agreement.
    let mut hyps: Vec<(u32, usize, u8, usize, usize)> = Vec::new();
    for (ti, (_, s)) in tiles.iter().enumerate() {
        for o in 0..ORIENTATIONS {
            let so = oriented(s, o);
            for y0 in 0..TILE {
                for x0 in 0..TILE {
                    hyps.push((agreement(&so, x0, y0), ti, o, x0, y0));
                }
            }
        }
    }
    // Highest agreement first; ties in the order listed (candidate, orientation, y0, x0).
    hyps.sort_by_key(|h| std::cmp::Reverse(h.0));
    let mut words = Vec::new();
    trace.decoded = hyps.iter().take(DECODES).copied().collect();
    for (rank, &(count, ti, o, x0, y0)) in hyps.iter().take(DECODES).enumerate() {
        let (res, s) = &tiles[ti];
        let wd = word(&oriented(s, o), x0, y0);
        words.push(wd);
        if let Some((payload, errors)) = decode(&wd)
            && stop
        {
            let resync = Resync { offset: (x0, y0), rank, agreement: count, ..*res };
            return (Some(Detected { payload, stage: 2, orientation: o, errors, resync: Some(resync) }), words, trace);
        }
    }
    (None, words, trace)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::synthesis::mark;

    #[test]
    fn the_tiled_layout_puts_each_bit_in_four_positions_of_every_tile() {
        let lay = layout();
        let tl = tile_layout();
        let mut seen = vec![0u8; CELLS];
        for j in 0..CODE_BITS {
            let cells = &lay.perm[16 * j..16 * j + 16];
            for &c in cells {
                seen[c as usize] += 1;
                assert_eq!(lay.bit_of_cell[c as usize], Some(j as u16));
            }
            for t in 0..4 {
                let in_tile = cells.iter().filter(|&&c| (c as usize / GRID) / TILE * 2 + (c as usize % GRID) / TILE == t).count();
                assert_eq!(in_tile, PER_TILE);
            }
        }
        assert!(lay.perm.iter().all(|&c| seen[c as usize] <= 1));
        assert_eq!(lay.bit_of_cell.iter().filter(|b| b.is_none()).count(), 128);
        // Every tile holds the same signs.
        for p in 0..POSITIONS {
            for t in 0..4 {
                assert_eq!(lay.sign[cell_of(p, t)], tl.sign[p]);
            }
        }
        assert_eq!(mark::decode(&tl.mask), None);
    }

    #[test]
    fn the_correlation_is_exact() {
        let (wq, hq) = (9, 7);
        let q: Vec<i64> = (0..wq * hq).map(|i| [(-1i64), 0, 1][(i * 7 + i / 3) % 3]).collect();
        let acf = Acf::new(std::slice::from_ref(&q), wq, hq);
        for dy in -(hq as i64) + 1..hq as i64 {
            for dx in -(wq as i64) + 1..wq as i64 {
                let mut s = 0;
                for y in 0..hq as i64 {
                    for x in 0..wq as i64 {
                        let (x2, y2) = (x + dx, y + dy);
                        if x2 >= 0 && y2 >= 0 && x2 < wq as i64 && y2 < hq as i64 {
                            s += q[(y * wq as i64 + x) as usize] * q[(y2 * wq as i64 + x2) as usize];
                        }
                    }
                }
                assert_eq!(acf.at(dx, dy), s, "({dx}, {dy})");
            }
        }
    }

    #[test]
    fn rounding_and_the_vertex() {
        assert_eq!(rnd(5, 2), 3);
        assert_eq!(rnd(-5, 2), -3);
        assert_eq!(rnd(5, -2), -3);
        assert_eq!(vertex(1, 4, 1), 0);
        assert_eq!(vertex(3, 4, 1), -16);
        assert_eq!(vertex(1, 4, 3), 16);
        assert_eq!(vertex(1, 1, 1), 0);
    }

    #[cfg(feature = "synthesis_mark")]
    fn picture(w: usize, h: usize, seed: u32) -> Vec<u8> {
        crate::synthesis_mark::leaves(seed, w, h)
    }

    /// The battery and its pictures are `synthesis_mark`'s.
    #[cfg(feature = "synthesis_mark")]
    #[test]
    fn a_crop_and_a_rotation_are_found_again() {
        use crate::synthesis_mark::battery::{apply, Attack};
        use crate::synthesis_mark::Frames;
        let (w, h) = (256, 192);
        let payload = [3u8, 1, 2, 3, 4, 5, 6, 7];
        let (marked, _) = embed(&picture(w, h, 5), w, h, 1, &payload);
        let d = detect(&marked, w, h, 1).expect("a mark");
        assert_eq!((d.payload, d.stage), (payload, 1));
        let f = Frames::one(marked, w, h);
        for a in [Attack::CropSide(0, 10), Attack::CropSide(3, 25), Attack::Crop(10), Attack::Rotate(2)] {
            let y = apply(a, &f, 0);
            let d = detect(&y.data, y.w, y.h, 1).unwrap_or_else(|| panic!("{a:?}"));
            assert_eq!(d.payload, payload, "{a:?}");
        }
    }
}
