//! The attack battery of extensions/synthesis.md §7.5.7: each attack an integer function of
//! the frames, so its result is the same bytes on every machine.

use super::{words, Frames};
use crate::int_video_enc::{encode, EncParams};
use crate::synthesis::{denoise3, upscale2};

/// An attack of the battery.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Attack {
    /// The frames as marked.
    None,
    /// `wai.video.int_motion`'s encoder at step `q`, each frame on its own (one keyframe).
    IntMotion(u32),
    /// The base of a staged picture: each 2 × 2 block's rounded mean, encoded as
    /// `wai.video.int_motion` at step `q`, then `wai-synthesis:upscale2`.
    StagedBase(u32),
    /// `wai-synthesis:denoise3`, applied `n` times.
    Denoise(u8),
    /// Area resampling to `⌊w·num/den⌋` × `⌊h·num/den⌋`.
    Scale(u32, u32),
    /// Area resampling of the width alone, to `⌊w·num/den⌋`.
    Aspect(u32, u32),
    /// `wai-synthesis:upscale2`.
    Upscale2,
    /// A centred crop keeping `100 − p` percent of each side.
    Crop(u32),
    /// `p` percent of the frame removed from one side: 0 the left, 1 the right, 2 the top,
    /// 3 the bottom.
    CropSide(u8, u32),
    /// The columns reversed.
    Mirror,
    /// A quarter turn clockwise.
    Rotate90,
    /// A rotation by a small angle about the centre, bilinear, by the `(cos, sin)` of
    /// [`ROTATIONS`] at that index.
    Rotate(usize),
    /// A simulated capture of a screen: area resampling by `num/den`, a rotation by
    /// [`ROTATIONS`] at an index (clockwise where the flag is set), the display curve
    /// `rnd(v·(v + 255) / 510)`, then noise uniform in `−4..=4`.
    Screen(u32, u32, usize, bool),
    /// Noise uniform in `−s..=s` on every sample.
    Noise(u32),
    /// A constant added to every sample.
    Brightness(i32),
    /// `128 + (v − 128)·num/den`, rounded half away from zero.
    Contrast(u32, u32),
    /// The central quarter of the picture's area replaced by the means of its 16 × 16
    /// blocks.
    LocalBlur,
    /// Area resampling to half, then `wai.video.int_motion` at step `q`.
    HalfThenIntMotion(u32),
    /// A clip encoded as `wai.video.int_motion` with motion: blocks of 16, a search radius
    /// of 4, half-pel vectors, step `q`.
    ClipIntMotion(u32),
    /// Every other frame of a clip, from the first.
    FrameDrop,
    /// The first frame of a clip alone.
    OneFrame,
}

/// The rotations' angles in tenths of a degree, and their cosines and sines at `2^16`.
pub const ROTATIONS: [(u32, i64, i64); 4] = [(5, 65534, 572), (10, 65526, 1144), (20, 65496, 2287), (50, 65287, 5712)];

/// The attacks applied to pictures, by name, in the table's order.
pub const IMAGE_ATTACKS: [(&str, Attack); 33] = [
    ("none", Attack::None),
    ("int-motion-q4", Attack::IntMotion(4)),
    ("int-motion-q8", Attack::IntMotion(8)),
    ("int-motion-q16", Attack::IntMotion(16)),
    ("int-motion-q32", Attack::IntMotion(32)),
    ("int-motion-q48", Attack::IntMotion(48)),
    ("staged-base-q8", Attack::StagedBase(8)),
    ("denoise3", Attack::Denoise(1)),
    ("denoise3-x4", Attack::Denoise(4)),
    ("scale-3/4", Attack::Scale(3, 4)),
    ("scale-1/2", Attack::Scale(1, 2)),
    ("scale-1/3", Attack::Scale(1, 3)),
    ("scale-1/4", Attack::Scale(1, 4)),
    ("scale-3/2", Attack::Scale(3, 2)),
    ("upscale2", Attack::Upscale2),
    ("aspect-3/4", Attack::Aspect(3, 4)),
    ("half-then-int-motion-q16", Attack::HalfThenIntMotion(16)),
    ("crop-1%", Attack::Crop(1)),
    ("crop-2%", Attack::Crop(2)),
    ("crop-5%", Attack::Crop(5)),
    ("crop-10%", Attack::Crop(10)),
    ("mirror", Attack::Mirror),
    ("rotate-90", Attack::Rotate90),
    ("rotate-0.5", Attack::Rotate(0)),
    ("rotate-1", Attack::Rotate(1)),
    ("rotate-2", Attack::Rotate(2)),
    ("noise-4", Attack::Noise(4)),
    ("noise-16", Attack::Noise(16)),
    ("brightness+32", Attack::Brightness(32)),
    ("brightness-32", Attack::Brightness(-32)),
    ("contrast-1/2", Attack::Contrast(1, 2)),
    ("contrast-3/2", Attack::Contrast(3, 2)),
    ("local-blur", Attack::LocalBlur),
];

/// The attacks the resynchronising scheme adds to the battery (§7.9.6): a crop of 25%, crops
/// from each side, a rotation of 5° and two simulated captures of a screen.
pub const SYNC_ATTACKS: [(&str, Attack); 20] = [
    ("crop-25%", Attack::Crop(25)),
    ("crop-left-2%", Attack::CropSide(0, 2)),
    ("crop-right-2%", Attack::CropSide(1, 2)),
    ("crop-top-2%", Attack::CropSide(2, 2)),
    ("crop-bottom-2%", Attack::CropSide(3, 2)),
    ("crop-left-5%", Attack::CropSide(0, 5)),
    ("crop-right-5%", Attack::CropSide(1, 5)),
    ("crop-top-5%", Attack::CropSide(2, 5)),
    ("crop-bottom-5%", Attack::CropSide(3, 5)),
    ("crop-left-10%", Attack::CropSide(0, 10)),
    ("crop-right-10%", Attack::CropSide(1, 10)),
    ("crop-top-10%", Attack::CropSide(2, 10)),
    ("crop-bottom-10%", Attack::CropSide(3, 10)),
    ("crop-left-25%", Attack::CropSide(0, 25)),
    ("crop-right-25%", Attack::CropSide(1, 25)),
    ("crop-top-25%", Attack::CropSide(2, 25)),
    ("crop-bottom-25%", Attack::CropSide(3, 25)),
    ("rotate-5", Attack::Rotate(3)),
    ("screen-a", Attack::Screen(5, 4, 1, false)),
    ("screen-b", Attack::Screen(4, 5, 2, true)),
];

/// The attacks applied to clips, by name, in the table's order.
pub const CLIP_ATTACKS: [(&str, Attack); 7] = [
    ("none", Attack::None),
    ("clip-int-motion-q8", Attack::ClipIntMotion(8)),
    ("clip-int-motion-q24", Attack::ClipIntMotion(24)),
    ("frame-drop", Attack::FrameDrop),
    ("one-frame", Attack::OneFrame),
    ("scale-1/2", Attack::Scale(1, 2)),
    ("noise-16", Attack::Noise(16)),
];

/// The attack a name of the battery's lists names, and a crop from one side by any percent
/// below 100 (`crop-left-40%`), which a corpus derives frames by.
pub fn attack_named(name: &str) -> Option<Attack> {
    if let Some((_, a)) = IMAGE_ATTACKS.iter().chain(SYNC_ATTACKS.iter()).chain(CLIP_ATTACKS.iter()).find(|(n, _)| *n == name) {
        return Some(*a);
    }
    for (side, prefix) in ["crop-left-", "crop-right-", "crop-top-", "crop-bottom-"].iter().enumerate() {
        if let Some(p) = name.strip_prefix(prefix).and_then(|r| r.strip_suffix('%')).and_then(|p| p.parse::<u32>().ok()) {
            return (p < 100).then_some(Attack::CropSide(side as u8, p));
        }
    }
    None
}

/// The frames a corpus forges a lattice into: each frame tiled by its own top-left quarter,
/// `⌊w/2⌋ × ⌊h/2⌋`, so its content repeats with the period a tiled mark has and carries no
/// mark (§7.9, the corpus's negative control).
pub fn periodic(f: &Frames) -> Frames {
    let (w, h) = (f.w, f.h);
    let (qw, qh) = (w / 2, h / 2);
    each_frame(f, w, h, |fr| {
        let mut out = Vec::with_capacity(w * h * 3);
        for y in 0..h {
            for x in 0..w {
                let p = ((y % qh) * w + x % qw) * 3;
                out.extend_from_slice(&fr[p..p + 3]);
            }
        }
        out
    })
}

fn each_frame(f: &Frames, w: usize, h: usize, g: impl Fn(&[u8]) -> Vec<u8>) -> Frames {
    let mut data = Vec::with_capacity(w * h * 3 * f.n);
    for i in 0..f.n {
        data.extend(g(f.frame(i)));
    }
    Frames { data, w, h, n: f.n }
}

/// `wai.video.int_motion` of one frame at step `q`: a keyframe, whose reconstruction is
/// `clamp(128 + r·q)` with `r` the sample less 128 divided by `q`, rounded half away from
/// zero (integer-payloads §3).
pub fn int_motion(frame: &[u8], w: usize, h: usize, q: u32) -> Vec<u8> {
    let p = EncParams { block: 1, search: 0, mv_precision: 1, q };
    encode(&[frame.to_vec()], h as u32, w as u32, &p).expect("a frame the encoder takes").recon.remove(0)
}

/// For each destination sample along an axis of `n` source samples resampled to `m`: the
/// source samples it covers and their overlaps, in units of `1/m` of a source sample; the
/// overlaps of one destination sample sum to `n`.
fn overlaps(n: usize, m: usize) -> Vec<Vec<(usize, u64)>> {
    (0..m)
        .map(|d| {
            let (lo, hi) = (d * n, (d + 1) * n);
            (lo / m..hi.div_ceil(m))
                .filter_map(|s| {
                    let o = hi.min((s + 1) * m) as i64 - lo.max(s * m) as i64;
                    (o > 0).then_some((s, o as u64))
                })
                .collect()
        })
        .collect()
}

/// Area resampling of one frame to `nw` × `nh`: each destination sample the mean of the
/// source it covers, weighted by overlap, rounded half up.
pub fn area(frame: &[u8], w: usize, h: usize, nw: usize, nh: usize) -> Vec<u8> {
    let (ox, oy) = (overlaps(w, nw), overlaps(h, nh));
    let den = (w * h) as u64;
    let mut out = vec![0u8; nw * nh * 3];
    for (dy, ry) in oy.iter().enumerate() {
        for (dx, rx) in ox.iter().enumerate() {
            for c in 0..3 {
                let mut num = 0u64;
                for &(sy, wy) in ry {
                    for &(sx, wx) in rx {
                        num += wy * wx * frame[(sy * w + sx) * 3 + c] as u64;
                    }
                }
                out[(dy * nw + dx) * 3 + c] = ((2 * num + den) / (2 * den)) as u8;
            }
        }
    }
    out
}

/// `n / d` rounded to the nearest integer, halves away from zero.
fn rdiv(n: i64, d: i64) -> i64 {
    n.signum() * ((2 * n.abs() + d) / (2 * d))
}

/// A rotation of one frame about its centre by the angle whose cosine and sine at `2^16`
/// are `(c, s)`: destination `(x, y)` reads the source at `(X, Y)`, where, at `2^16`,
/// `X = (c·u + s·v + (w − 1)·2^16) >> 1` and `Y = (−s·u + c·v + (h − 1)·2^16) >> 1` with
/// `u = 2x + 1 − w` and `v = 2y + 1 − h`, bilinearly from the four samples around it, a
/// sample outside the frame being 0, rounded half up.
pub fn rotate(frame: &[u8], w: usize, h: usize, c: i64, s: i64) -> Vec<u8> {
    let at = |x: i64, y: i64, ch: usize| -> i64 {
        if x < 0 || y < 0 || x >= w as i64 || y >= h as i64 { 0 } else { frame[(y as usize * w + x as usize) * 3 + ch] as i64 }
    };
    let one = 1i64 << 16;
    let mut out = vec![0u8; w * h * 3];
    for y in 0..h {
        for x in 0..w {
            let (u, v) = (2 * x as i64 + 1 - w as i64, 2 * y as i64 + 1 - h as i64);
            let px = (c * u + s * v + (w as i64 - 1) * one) >> 1;
            let py = (-s * u + c * v + (h as i64 - 1) * one) >> 1;
            let (xi, fx, yi, fy) = (px >> 16, px & 0xffff, py >> 16, py & 0xffff);
            for ch in 0..3 {
                let v = at(xi, yi, ch) * (one - fx) * (one - fy)
                    + at(xi + 1, yi, ch) * fx * (one - fy)
                    + at(xi, yi + 1, ch) * (one - fx) * fy
                    + at(xi + 1, yi + 1, ch) * fx * fy;
                out[(y * w + x) * 3 + ch] = ((v + (1 << 31)) >> 32) as u8;
            }
        }
    }
    out
}

/// Apply an attack to frames, with `seed` for the noise's stream.
pub fn apply(a: Attack, f: &Frames, seed: u32) -> Frames {
    let (w, h) = (f.w, f.h);
    match a {
        Attack::None => f.clone(),
        Attack::IntMotion(q) => each_frame(f, w, h, |fr| int_motion(fr, w, h, q)),
        Attack::StagedBase(q) => each_frame(f, w, h, |fr| {
            let (hw, hh) = (w / 2, h / 2);
            let mut half = vec![0u8; hw * hh * 3];
            for y in 0..hh {
                for x in 0..hw {
                    for c in 0..3 {
                        let s = |dx: usize, dy: usize| fr[((2 * y + dy) * w + 2 * x + dx) * 3 + c] as u32;
                        half[(y * hw + x) * 3 + c] = ((s(0, 0) + s(1, 0) + s(0, 1) + s(1, 1) + 2) >> 2) as u8;
                    }
                }
            }
            upscale2(&int_motion(&half, hw, hh, q), hw, hh)
        }),
        Attack::Denoise(n) => each_frame(f, w, h, |fr| (0..n).fold(fr.to_vec(), |b, _| denoise3(&b, w, h))),
        Attack::Scale(num, den) => {
            let (nw, nh) = (w * num as usize / den as usize, h * num as usize / den as usize);
            each_frame(f, nw, nh, |fr| area(fr, w, h, nw, nh))
        }
        Attack::Aspect(num, den) => {
            let nw = w * num as usize / den as usize;
            each_frame(f, nw, h, |fr| area(fr, w, h, nw, h))
        }
        Attack::Upscale2 => each_frame(f, 2 * w, 2 * h, |fr| upscale2(fr, w, h)),
        Attack::Crop(p) => {
            let (dx, dy) = (w * p as usize / 200, h * p as usize / 200);
            let (nw, nh) = (w - 2 * dx, h - 2 * dy);
            each_frame(f, nw, nh, |fr| {
                let mut out = Vec::with_capacity(nw * nh * 3);
                for y in dy..dy + nh {
                    out.extend_from_slice(&fr[(y * w + dx) * 3..(y * w + dx + nw) * 3]);
                }
                out
            })
        }
        Attack::CropSide(side, p) => {
            let (cut_x, cut_y) = (w * p as usize / 100, h * p as usize / 100);
            let (x0, y0, nw, nh) = match side {
                0 => (cut_x, 0, w - cut_x, h),
                1 => (0, 0, w - cut_x, h),
                2 => (0, cut_y, w, h - cut_y),
                _ => (0, 0, w, h - cut_y),
            };
            each_frame(f, nw, nh, |fr| {
                let mut out = Vec::with_capacity(nw * nh * 3);
                for y in y0..y0 + nh {
                    out.extend_from_slice(&fr[(y * w + x0) * 3..(y * w + x0 + nw) * 3]);
                }
                out
            })
        }
        Attack::Screen(num, den, ri, clockwise) => {
            let (nw, nh) = (w * num as usize / den as usize, h * num as usize / den as usize);
            let (_, c, s) = ROTATIONS[ri];
            let s = if clockwise { -s } else { s };
            let shown = each_frame(f, nw, nh, |fr| {
                rotate(&area(fr, w, h, nw, nh), nw, nh, c, s).iter().map(|&v| ((2 * v as u32 * (v as u32 + 255) + 510) / 1020) as u8).collect()
            });
            apply(Attack::Noise(4), &shown, seed)
        }
        Attack::Mirror => each_frame(f, w, h, |fr| {
            let mut out = Vec::with_capacity(w * h * 3);
            for y in 0..h {
                for x in (0..w).rev() {
                    out.extend_from_slice(&fr[(y * w + x) * 3..(y * w + x) * 3 + 3]);
                }
            }
            out
        }),
        Attack::Rotate90 => each_frame(f, h, w, |fr| {
            // Clockwise: destination (x', y') of h × w is source (y', h − 1 − x').
            let mut out = Vec::with_capacity(w * h * 3);
            for yp in 0..w {
                for xp in 0..h {
                    let (sx, sy) = (yp, h - 1 - xp);
                    out.extend_from_slice(&fr[(sy * w + sx) * 3..(sy * w + sx) * 3 + 3]);
                }
            }
            out
        }),
        Attack::Rotate(i) => {
            let (_, c, s) = ROTATIONS[i];
            each_frame(f, w, h, |fr| rotate(fr, w, h, c, s))
        }
        Attack::Noise(s) => {
            let mut r = words(b"wai-mark:battery noise", seed);
            let span = 2 * s + 1;
            Frames { data: f.data.iter().map(|&v| (v as i64 + (r.next().expect("an endless stream") % span) as i64 - s as i64).clamp(0, 255) as u8).collect(), ..f.clone() }
        }
        Attack::Brightness(b) => Frames { data: f.data.iter().map(|&v| (v as i64 + b as i64).clamp(0, 255) as u8).collect(), ..f.clone() },
        Attack::Contrast(num, den) => {
            Frames { data: f.data.iter().map(|&v| (128 + rdiv((v as i64 - 128) * num as i64, den as i64)).clamp(0, 255) as u8).collect(), ..f.clone() }
        }
        Attack::LocalBlur => each_frame(f, w, h, |fr| {
            let mut out = fr.to_vec();
            let (x0, x1, y0, y1) = (w / 4, 3 * w / 4, h / 4, 3 * h / 4);
            for by in (y0..y1).step_by(16) {
                for bx in (x0..x1).step_by(16) {
                    let (ex, ey) = ((bx + 16).min(x1), (by + 16).min(y1));
                    let n = ((ex - bx) * (ey - by)) as u64;
                    for c in 0..3 {
                        let mut sum = 0u64;
                        for y in by..ey {
                            for x in bx..ex {
                                sum += fr[(y * w + x) * 3 + c] as u64;
                            }
                        }
                        let m = ((2 * sum + n) / (2 * n)) as u8;
                        for y in by..ey {
                            for x in bx..ex {
                                out[(y * w + x) * 3 + c] = m;
                            }
                        }
                    }
                }
            }
            out
        }),
        Attack::HalfThenIntMotion(q) => {
            let (nw, nh) = (w / 2, h / 2);
            each_frame(f, nw, nh, |fr| int_motion(&area(fr, w, h, nw, nh), nw, nh, q))
        }
        Attack::ClipIntMotion(q) => {
            let p = EncParams { block: 16, search: 4, mv_precision: 2, q };
            let frames: Vec<Vec<u8>> = (0..f.n).map(|i| f.frame(i).to_vec()).collect();
            let enc = encode(&frames, h as u32, w as u32, &p).expect("a clip the encoder takes");
            Frames { data: enc.recon.concat(), ..f.clone() }
        }
        Attack::FrameDrop => {
            let kept: Vec<usize> = (0..f.n).step_by(2).collect();
            Frames { data: kept.iter().flat_map(|&i| f.frame(i).to_vec()).collect(), n: kept.len(), ..f.clone() }
        }
        Attack::OneFrame => Frames { data: f.frame(0).to_vec(), n: 1, ..f.clone() },
    }
}

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

    #[test]
    fn area_resampling_keeps_a_flat_frame_and_halves_exactly() {
        let flat = vec![77u8; 10 * 6 * 3];
        assert!(area(&flat, 10, 6, 7, 4).iter().all(|&v| v == 77));
        let mut f = vec![0u8; 4 * 2 * 3];
        for (i, v) in f.iter_mut().enumerate() {
            *v = (i * 10) as u8;
        }
        let half = area(&f, 4, 2, 2, 1);
        // Each half is the mean of its 2 × 2 block, per channel.
        let want = |x0: usize, c: usize| ((f[x0 * 3 + c] as u32 + f[(x0 + 1) * 3 + c] as u32 + f[(4 + x0) * 3 + c] as u32 + f[(5 + x0) * 3 + c] as u32 + 2) / 4) as u8;
        assert_eq!(half, vec![want(0, 0), want(0, 1), want(0, 2), want(2, 0), want(2, 1), want(2, 2)]);
    }

    #[test]
    fn the_overlaps_of_a_destination_sum_to_the_source() {
        for (n, m) in [(512, 384), (512, 170), (100, 300), (7, 3)] {
            for o in overlaps(n, m) {
                assert_eq!(o.iter().map(|x| x.1).sum::<u64>(), n as u64);
            }
        }
    }

    #[test]
    fn a_zero_rotation_is_the_identity() {
        let f: Vec<u8> = (0..20 * 12 * 3).map(|i| (i * 7 % 251) as u8).collect();
        assert_eq!(rotate(&f, 20, 12, 1 << 16, 0), f);
    }

    #[test]
    fn a_keyframe_is_scalar_quantisation() {
        let f: Vec<u8> = (0..16 * 16 * 3).map(|i| (i * 13 % 256) as u8).collect();
        let r = int_motion(&f, 16, 16, 8);
        for (a, b) in f.iter().zip(&r) {
            assert_eq!(*b as i64, (128 + rdiv(*a as i64 - 128, 8) * 8).clamp(0, 255));
        }
    }
}
