//! The fully linear proof of `draft-irtf-cfrg-vdaf-22` §7.3 over the validity circuits of
//! §7.4.1–§7.4.4 (Count, Sum, SumVec, Histogram), with the gadgets of Appendix A and the
//! polynomials in the Lagrange basis of §6.1.3.2.
//!
//! Each of those circuits calls one gadget, so a circuit here names one; the proof,
//! query and decision follow §7.3.3–§7.3.5 for a gadget list of length one. The NTT of
//! §6.1.2 is evaluated directly: the polynomials of these circuits are short.

use super::field::Field;

pub fn next_power_of_2(x: usize) -> usize {
    x.max(1).next_power_of_two()
}

/// wire_poly_len (§7.3.2): the smallest power of two that holds the seed and every call.
pub fn wire_poly_len(calls: usize) -> usize {
    next_power_of_2(1 + calls)
}

/// gadget_poly_len (§7.3.2).
pub fn gadget_poly_len(degree: usize, wire_len: usize) -> usize {
    degree * (wire_len - 1) + 1
}

/// The first `n` powers of the principal `n`-th root of unity.
pub fn nth_root_powers<F: Field>(n: usize) -> Vec<F> {
    let w = F::nth_root(n);
    let mut v = Vec::with_capacity(n);
    let mut x = F::one();
    for _ in 0..n {
        v.push(x);
        x *= w;
    }
    v
}

/// Horner evaluation of a polynomial in the monomial basis.
pub fn poly_eval_mono<F: Field>(p: &[F], x: F) -> F {
    p.iter().rev().fold(F::zero(), |acc, c| acc * x + *c)
}

/// ntt (§6.1.2): the polynomial `p`, by its coefficients, at the `n`-th roots of unity, or at
/// those roots times s = the principal 2n-th root when `set_s`.
pub fn ntt<F: Field>(p: &[F], n: usize, set_s: bool) -> Vec<F> {
    let s = if set_s { F::nth_root(2 * n) } else { F::one() };
    nth_root_powers::<F>(n).into_iter().map(|w| poly_eval_mono(p, s * w)).collect()
}

/// inv_ntt (§6.1.2): the coefficients of the polynomial whose values at the `n`-th roots of
/// unity are `v`.
pub fn inv_ntt<F: Field>(v: &[F], n: usize) -> Vec<F> {
    assert_eq!(v.len(), n);
    let wi = F::nth_root(n).inv();
    let n_inv = F::from_u128(n as u128).inv();
    (0..n)
        .map(|j| {
            let step = wi.pow(j as u128);
            let mut x = F::one();
            let mut acc = F::zero();
            for vi in v {
                acc += *vi * x;
                x *= step;
            }
            acc * n_inv
        })
        .collect()
}

/// double_evaluations (§6.1.3.2): 2n values of the polynomial from its n values.
pub fn double_evaluations<F: Field>(p: &[F]) -> Vec<F> {
    let n = p.len();
    assert!(n.is_power_of_two());
    let odd = ntt(&inv_ntt(p, n), n, true);
    p.iter().zip(odd).flat_map(|(e, o)| [*e, o]).collect()
}

/// poly_mul (§6.1.3.2).
pub fn poly_mul<F: Field>(p: &[F], q: &[F]) -> Vec<F> {
    assert_eq!(p.len(), q.len());
    let (a, b) = (double_evaluations(p), double_evaluations(q));
    a.iter().zip(b).map(|(x, y)| *x * y).collect()
}

/// poly_eval_batched (§6.1.3.2), as the draft writes it.
pub fn poly_eval_batched<F: Field>(polys: &[Vec<F>], x: F) -> Vec<F> {
    let n = polys[0].len();
    assert!(n.is_power_of_two() && polys.iter().all(|p| p.len() == n));
    let nodes = nth_root_powers::<F>(n);
    let mut k = F::one();
    let mut u: Vec<F> = polys.iter().map(|p| p[0]).collect();
    let mut d = nodes[0] - x;
    for i in 1..n {
        k *= d;
        d = nodes[i] - x;
        let t = k * nodes[i];
        for (j, p) in polys.iter().enumerate() {
            u[j] *= d;
            if i < p.len() {
                u[j] += t * p[i];
            }
        }
    }
    let sign = if (n - 1) % 2 == 1 { -F::one() } else { F::one() };
    let factor = sign * F::from_u128(n as u128).inv();
    u.into_iter().map(|v| v * factor).collect()
}

pub fn poly_eval<F: Field>(p: &[F], x: F) -> F {
    poly_eval_batched(&[p.to_vec()], x)[0]
}

/// extend_values_to_power_of_2 (§6.1.3.2), as the draft writes it.
pub fn extend_values_to_power_of_2<F: Field>(p: &mut Vec<F>, n: usize) {
    assert!(n.is_power_of_two() && p.len() <= n);
    let x = nth_root_powers::<F>(n);
    let mut w = vec![F::zero(); n];
    let len = p.len();
    for i in 0..len {
        w[i] = (0..len).filter(|&j| j != i).fold(F::one(), |acc, j| acc * (x[i] - x[j]));
    }
    for k in len..n {
        for i in 0..k {
            w[i] *= x[i] - x[k];
        }
        let (mut y_num, mut y_den) = (F::zero(), F::one());
        for (i, v) in p.iter().enumerate() {
            y_num = y_num * w[i] + y_den * *v;
            y_den *= w[i];
        }
        w[k] = (0..k).fold(F::one(), |acc, j| acc * (x[k] - x[j]));
        p.push(-w[k] * y_num * y_den.inv());
    }
}

/// A gadget of Appendix A.
#[derive(Clone, Debug)]
pub enum Gadget {
    /// Mul (A.1): x · y.
    Mul,
    /// PolyEval (A.2): p(x), `p` by its coefficients, leading zeros stripped; `n` the number
    /// of values its gadget polynomial is computed at.
    PolyEval { p: Vec<i128>, n: usize },
    /// ParallelSum (A.3) of `count` Mul subcircuits.
    ParallelSumMul { count: usize },
}

impl Gadget {
    pub fn poly_eval(p: &[i128], calls: usize) -> Gadget {
        let mut p = p.to_vec();
        while p.len() > 1 && *p.last().unwrap() == 0 {
            p.pop();
        }
        assert!(!p.is_empty());
        let degree = p.len() - 1;
        let n = next_power_of_2(gadget_poly_len(degree, wire_poly_len(calls)));
        Gadget::PolyEval { p, n }
    }

    pub fn arity(&self) -> usize {
        match self {
            Gadget::Mul => 2,
            Gadget::PolyEval { .. } => 1,
            Gadget::ParallelSumMul { count } => 2 * count,
        }
    }

    pub fn degree(&self) -> usize {
        match self {
            Gadget::Mul | Gadget::ParallelSumMul { .. } => 2,
            Gadget::PolyEval { p, .. } => p.len() - 1,
        }
    }

    pub fn eval<F: Field>(&self, inp: &[F]) -> F {
        match self {
            Gadget::Mul => inp[0] * inp[1],
            Gadget::PolyEval { p, .. } => {
                let c: Vec<F> = p.iter().map(|&x| F::from_i128(x)).collect();
                poly_eval_mono(&c, inp[0])
            }
            Gadget::ParallelSumMul { count } => (0..*count).fold(F::zero(), |acc, i| acc + inp[2 * i] * inp[2 * i + 1]),
        }
    }

    /// The gadget over the wire polynomials, all in the Lagrange basis.
    pub fn eval_poly<F: Field>(&self, inp: &[Vec<F>]) -> Vec<F> {
        match self {
            Gadget::Mul => poly_mul(&inp[0], &inp[1]),
            Gadget::PolyEval { p, n } => {
                let len = inp[0].len();
                let mono = inv_ntt(&inp[0], len);
                let c: Vec<F> = p.iter().map(|&x| F::from_i128(x)).collect();
                ntt(&mono, *n, false).into_iter().map(|x| poly_eval_mono(&c, x)).collect()
            }
            Gadget::ParallelSumMul { count } => {
                let out_len = next_power_of_2(gadget_poly_len(2, inp[0].len()));
                let mut sum = vec![F::zero(); out_len];
                for i in 0..*count {
                    let cur = poly_mul(&inp[2 * i], &inp[2 * i + 1]);
                    for (s, c) in sum.iter_mut().zip(cur) {
                        *s += c;
                    }
                }
                sum
            }
        }
    }
}

/// A validity circuit (§7.3.2) with one gadget.
pub trait Valid<F: Field> {
    fn gadget(&self) -> &Gadget;
    fn gadget_calls(&self) -> usize;
    fn meas_len(&self) -> usize;
    fn joint_rand_len(&self) -> usize;
    fn eval_output_len(&self) -> usize;
    fn output_len(&self) -> usize;
    /// The circuit over `meas`, each gadget call made through `g`.
    fn eval(&self, meas: &[F], joint_rand: &[F], num_shares: usize, g: &mut dyn FnMut(&[F]) -> F) -> Vec<F>;
    fn truncate(&self, meas: &[F]) -> Vec<F>;

    fn prove_rand_len(&self) -> usize {
        self.gadget().arity()
    }
    fn query_rand_len(&self) -> usize {
        1 + if self.eval_output_len() > 1 { self.eval_output_len() } else { 0 }
    }
    fn proof_len(&self) -> usize {
        let g = self.gadget();
        g.arity() + gadget_poly_len(g.degree(), wire_poly_len(self.gadget_calls()))
    }
    fn verifier_len(&self) -> usize {
        1 + self.gadget().arity() + 1
    }
}

/// prove (§7.3.3).
pub fn prove<F: Field, V: Valid<F> + ?Sized>(valid: &V, meas: &[F], prove_rand: &[F], joint_rand: &[F]) -> Vec<F> {
    let g = valid.gadget();
    let p = wire_poly_len(valid.gadget_calls());
    let mut wires: Vec<Vec<F>> = (0..g.arity())
        .map(|j| {
            let mut w = vec![F::zero(); p];
            w[0] = prove_rand[j];
            w
        })
        .collect();
    let mut k = 0usize;
    valid.eval(meas, joint_rand, 1, &mut |inp: &[F]| {
        k += 1;
        for (j, x) in inp.iter().enumerate() {
            wires[j][k] = *x;
        }
        g.eval(inp)
    });
    let mut proof: Vec<F> = prove_rand[..g.arity()].to_vec();
    let gp = g.eval_poly(&wires);
    proof.extend_from_slice(&gp[..gadget_poly_len(g.degree(), p)]);
    proof
}

/// Why a query fails: the test point is one of the wire polynomials' points (§7.3.4).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TestPointIsRoot;

/// query (§7.3.4).
pub fn query<F: Field, V: Valid<F> + ?Sized>(
    valid: &V,
    meas: &[F],
    proof: &[F],
    query_rand: &[F],
    joint_rand: &[F],
    num_shares: usize,
) -> Result<Vec<F>, TestPointIsRoot> {
    let g = valid.gadget();
    let p = wire_poly_len(valid.gadget_calls());
    let gp_len = gadget_poly_len(g.degree(), p);
    let (seeds, rest) = proof.split_at(g.arity());
    let mut poly = rest[..gp_len].to_vec();
    let n = next_power_of_2(poly.len());
    extend_values_to_power_of_2(&mut poly, n);
    let size = next_power_of_2(gp_len);
    while poly.len() < size {
        poly = double_evaluations(&poly);
    }
    let step = size / p;
    let mut wires: Vec<Vec<F>> = seeds
        .iter()
        .map(|s| {
            let mut w = vec![F::zero(); p];
            w[0] = *s;
            w
        })
        .collect();
    let mut k = 0usize;
    let out = valid.eval(meas, joint_rand, num_shares, &mut |inp: &[F]| {
        k += 1;
        for (j, x) in inp.iter().enumerate() {
            wires[j][k] = *x;
        }
        poly[k * step]
    });
    let eol = valid.eval_output_len();
    let (v, t) = if eol > 1 {
        let v = query_rand[..eol].iter().zip(&out).fold(F::zero(), |acc, (r, o)| acc + *r * *o);
        (v, query_rand[eol])
    } else {
        (out[0], query_rand[0])
    };
    if t.pow(p as u128) == F::one() {
        return Err(TestPointIsRoot);
    }
    let mut verifier = vec![v];
    verifier.extend(poly_eval_batched(&wires, t));
    verifier.push(poly_eval(&poly, t));
    Ok(verifier)
}

/// decide (§7.3.5).
pub fn decide<F: Field, V: Valid<F> + ?Sized>(valid: &V, verifier: &[F]) -> bool {
    let g = valid.gadget();
    if verifier.len() != valid.verifier_len() || verifier[0] != F::zero() {
        return false;
    }
    let a = g.arity();
    g.eval(&verifier[1..1 + a]) == verifier[1 + a]
}

// ---------------------------------------------------------------------------------------
// The circuits of §7.4.

/// Count (§7.4.1).
pub struct Count {
    gadget: Gadget,
}

impl Default for Count {
    fn default() -> Self {
        Count { gadget: Gadget::Mul }
    }
}

impl<F: Field> Valid<F> for Count {
    fn gadget(&self) -> &Gadget {
        &self.gadget
    }
    fn gadget_calls(&self) -> usize {
        1
    }
    fn meas_len(&self) -> usize {
        1
    }
    fn joint_rand_len(&self) -> usize {
        0
    }
    fn eval_output_len(&self) -> usize {
        1
    }
    fn output_len(&self) -> usize {
        1
    }
    fn eval(&self, meas: &[F], _jr: &[F], _n: usize, g: &mut dyn FnMut(&[F]) -> F) -> Vec<F> {
        let sq = g(&[meas[0], meas[0]]);
        vec![sq - meas[0]]
    }
    fn truncate(&self, meas: &[F]) -> Vec<F> {
        meas.to_vec()
    }
}

/// bit_length of a positive integer.
pub fn bit_length(x: u128) -> usize {
    128 - x.leading_zeros() as usize
}

/// encode_range_checked_int (§7.4.2): `bits` − 1 binary digits of the rest and a last bit of
/// weight max − (2^(bits−1) − 1). `None` when the value is past `max`.
pub fn encode_range_checked_int<F: Field>(value: u128, max: u128) -> Option<Vec<F>> {
    if value > max {
        return None;
    }
    let bits = bit_length(max);
    let rest_ones = (1u128 << (bits - 1)) - 1;
    let last_weight = max - rest_ones;
    let (rest, last) = if value <= rest_ones { (value, 0) } else { (value - last_weight, 1) };
    let mut e: Vec<F> = (0..bits - 1).map(|l| F::from_u128((rest >> l) & 1)).collect();
    e.push(F::from_u128(last));
    Some(e)
}

/// decode_range_checked_int (§7.4.2): linear, so it applies to shares too.
pub fn decode_range_checked_int<F: Field>(e: &[F], max: u128) -> F {
    let bits = bit_length(max);
    let rest_ones = (1u128 << (bits - 1)) - 1;
    let last_weight = max - rest_ones;
    let mut d = F::zero();
    for (l, b) in e[..bits - 1].iter().enumerate() {
        d += F::from_u128(1u128 << l) * *b;
    }
    d + F::from_u128(last_weight) * e[bits - 1]
}

/// Sum (§7.4.2).
pub struct Sum {
    pub max: u128,
    pub bits: usize,
    gadget: Gadget,
}

impl Sum {
    pub fn new(max: u128) -> Sum {
        assert!(max > 0);
        let bits = bit_length(max);
        Sum { max, bits, gadget: Gadget::poly_eval(&[0, -1, 1], bits) }
    }
}

impl<F: Field> Valid<F> for Sum {
    fn gadget(&self) -> &Gadget {
        &self.gadget
    }
    fn gadget_calls(&self) -> usize {
        self.bits
    }
    fn meas_len(&self) -> usize {
        self.bits
    }
    fn joint_rand_len(&self) -> usize {
        0
    }
    fn eval_output_len(&self) -> usize {
        self.bits
    }
    fn output_len(&self) -> usize {
        1
    }
    fn eval(&self, meas: &[F], _jr: &[F], _n: usize, g: &mut dyn FnMut(&[F]) -> F) -> Vec<F> {
        meas.iter().map(|b| g(&[*b])).collect()
    }
    fn truncate(&self, meas: &[F]) -> Vec<F> {
        vec![decode_range_checked_int(meas, self.max)]
    }
}

/// The range check SumVec and Histogram share: ParallelSum over chunks of the measurement,
/// each element times a power of its chunk's joint randomness against the element less
/// 1/num_shares.
fn chunked_range_check<F: Field>(
    meas: &[F],
    joint_rand: &[F],
    num_shares: usize,
    calls: usize,
    chunk: usize,
    g: &mut dyn FnMut(&[F]) -> F,
) -> F {
    let shares_inv = F::from_u128(num_shares as u128).inv();
    let mut out = F::zero();
    for (i, r) in joint_rand.iter().enumerate().take(calls) {
        let mut r_power = *r;
        let mut inputs = Vec::with_capacity(2 * chunk);
        for j in 0..chunk {
            let index = i * chunk + j;
            let m = if index < meas.len() { meas[index] } else { F::zero() };
            inputs.push(r_power * m);
            inputs.push(m - shares_inv);
            r_power *= *r;
        }
        out += g(&inputs);
    }
    out
}

/// SumVec (§7.4.3).
pub struct SumVec {
    pub length: usize,
    pub max: u128,
    pub bits: usize,
    pub chunk: usize,
    calls: usize,
    gadget: Gadget,
}

impl SumVec {
    pub fn new(length: usize, max: u128, chunk: usize) -> SumVec {
        assert!(max > 0 && length > 0 && chunk > 0);
        let bits = bit_length(max);
        let calls = (length * bits).div_ceil(chunk);
        SumVec { length, max, bits, chunk, calls, gadget: Gadget::ParallelSumMul { count: chunk } }
    }
}

impl<F: Field> Valid<F> for SumVec {
    fn gadget(&self) -> &Gadget {
        &self.gadget
    }
    fn gadget_calls(&self) -> usize {
        self.calls
    }
    fn meas_len(&self) -> usize {
        self.length * self.bits
    }
    fn joint_rand_len(&self) -> usize {
        self.calls
    }
    fn eval_output_len(&self) -> usize {
        1
    }
    fn output_len(&self) -> usize {
        self.length
    }
    fn eval(&self, meas: &[F], jr: &[F], n: usize, g: &mut dyn FnMut(&[F]) -> F) -> Vec<F> {
        vec![chunked_range_check(meas, jr, n, self.calls, self.chunk, g)]
    }
    fn truncate(&self, meas: &[F]) -> Vec<F> {
        meas.chunks(self.bits).map(|c| decode_range_checked_int(c, self.max)).collect()
    }
}

/// Histogram (§7.4.4).
pub struct Histogram {
    pub length: usize,
    pub chunk: usize,
    calls: usize,
    gadget: Gadget,
}

impl Histogram {
    pub fn new(length: usize, chunk: usize) -> Histogram {
        assert!(length > 0 && chunk > 0);
        let calls = length.div_ceil(chunk);
        Histogram { length, chunk, calls, gadget: Gadget::ParallelSumMul { count: chunk } }
    }
}

impl<F: Field> Valid<F> for Histogram {
    fn gadget(&self) -> &Gadget {
        &self.gadget
    }
    fn gadget_calls(&self) -> usize {
        self.calls
    }
    fn meas_len(&self) -> usize {
        self.length
    }
    fn joint_rand_len(&self) -> usize {
        self.calls
    }
    fn eval_output_len(&self) -> usize {
        2
    }
    fn output_len(&self) -> usize {
        self.length
    }
    fn eval(&self, meas: &[F], jr: &[F], n: usize, g: &mut dyn FnMut(&[F]) -> F) -> Vec<F> {
        let range = chunked_range_check(meas, jr, n, self.calls, self.chunk, g);
        let shares_inv = F::from_u128(n as u128).inv();
        let sum = meas.iter().fold(-shares_inv, |acc, b| acc + *b);
        vec![range, sum]
    }
    fn truncate(&self, meas: &[F]) -> Vec<F> {
        meas.to_vec()
    }
}

#[cfg(test)]
mod tests {
    use super::super::field::{F128, F64};
    use super::*;

    #[test]
    fn ntt_and_its_inverse_round_trip() {
        let p: Vec<F128> = (1..=8u128).map(F128::from_u128).collect();
        assert_eq!(ntt(&inv_ntt(&p, 8), 8, false), p);
    }

    #[test]
    fn lagrange_evaluation_matches_the_monomial_form() {
        let mono: Vec<F64> = [3u128, 1, 4, 1, 5, 9, 2, 6].iter().map(|&x| F64::from_u128(x)).collect();
        let vals = ntt(&mono, 8, false);
        let x = F64::from_u128(123_456_789);
        assert_eq!(poly_eval(&vals, x), poly_eval_mono(&mono, x));
        let mut short = ntt(&mono[..5], 8, false)[..5].to_vec();
        extend_values_to_power_of_2(&mut short, 8);
        assert_eq!(short, ntt(&mono[..5], 8, false));
    }

    fn run<V: Valid<F128>>(v: &V, meas: &[F128], shares: usize) -> bool {
        let pr: Vec<F128> = (0..v.prove_rand_len()).map(|i| F128::from_u128(1000 + i as u128)).collect();
        let jr: Vec<F128> = (0..v.joint_rand_len()).map(|i| F128::from_u128(77 + i as u128)).collect();
        let qr: Vec<F128> = (0..v.query_rand_len()).map(|i| F128::from_u128(31337 + i as u128)).collect();
        let proof = prove(v, meas, &pr, &jr);
        assert_eq!(proof.len(), v.proof_len());
        // Share the measurement and the proof additively and query each share.
        let mut acc = vec![F128::zero(); v.verifier_len()];
        for s in 0..shares {
            let (m, p): (Vec<F128>, Vec<F128>) = if s + 1 < shares {
                (vec![F128::from_u128(5 + s as u128); meas.len()], vec![F128::from_u128(9 + s as u128); proof.len()])
            } else {
                let k = F128::from_u128((5..5 + shares as u128 - 1).sum());
                let kp = F128::from_u128((9..9 + shares as u128 - 1).sum());
                (meas.iter().map(|x| *x - k).collect(), proof.iter().map(|x| *x - kp).collect())
            };
            let vs = query(v, &m, &p, &qr, &jr, shares).unwrap();
            acc = acc.iter().zip(vs).map(|(a, b)| *a + b).collect();
        }
        decide(v, &acc)
    }

    #[test]
    fn a_valid_measurement_is_accepted_and_an_invalid_one_refused() {
        let h = Histogram::new(8, 3);
        let mut m = vec![F128::zero(); 8];
        m[2] = F128::one();
        assert!(run(&h, &m, 2));
        m[5] = F128::one();
        assert!(!run(&h, &m, 2), "two-hot");
        let sv = SumVec::new(3, 255, 4);
        let mut m: Vec<F128> = [7u128, 0, 255].iter().flat_map(|&x| encode_range_checked_int::<F128>(x, 255).unwrap()).collect();
        assert!(run(&sv, &m, 3));
        m[0] = F128::from_u128(2);
        assert!(!run(&sv, &m, 3), "a digit of 2");
    }
}
