//! Example of simplified NTWE scheme (Gärtner 2024)

use gwrizienn::{
	matrix::Matrix,
	ntt::{Ntt, NttInv},
	ring::{Lift, Ring},
	vector::{Vector, VectorRef},
	*,
};
use num_traits::{Inv, One, Zero};

// Implement Zq
ring!(Zq, u32, u32, i32, 50177);

// Implement Rq = Zq/(x^256+1)
// zeta=66 is the first 512-th root of unity mod q
poly!(Rq, 256, Zq, u32, u32, u32, u64, 50177, 66);

// Implement Zb and Rb with a big modulus for NTT in Z2q
// chosen because > N*q^2
ring!(Zb, u64, u128, i128, 644539222529);
poly!(
	Rb,
	256,
	Zb,
	u64,
	u128,
	u64,
	u128,
	644539222529,
	483489047161
);

//Implement Z2q
ring!(Z2q, u32, u64, i64, 100354);

// Implement R2q
// A different macro must be used because 2q is not prime.
poly2!(
	Rq,
	Rb,
	R2q,
	256,
	Z2q,
	u32,
	u64,
	u32,
	u64,
	i64,
	i64,
	50177,
	644539222529
);

const L: usize = 3;
const M: usize = 2;

fn main() {
	let mut rng = rand::thread_rng();
	let uniform = Zq::uniform();
	let ball_c = Z2q::uniform_positive_semiball(1);
	let ball_s = Zq::uniform_ball(1);
	let ball_y = Zb::uniform_ball(55);

	// generate secret key s = [f s0 e]
	let mut s = Vector::<Rq, { L + M + 1 }>::random(ball_s, &mut rng);
	let f0 = &mut s.0[0];
	*f0 *= Zq(2);
	*f0 += Zq::one();
	let s_ntt = s.clone().ntt();
	let f = &s_ntt.0[0];
	let s0: VectorRef<_, L> = s_ntt.get_sub(1);
	let e: VectorRef<_, M> = s_ntt.get_sub(1 + L);
	// generate public key
	let a0 = Matrix::<Rq, M, L>::random(uniform, &mut rng);
	let b = (&a0.clone().ntt() * s0 + e) * &f.clone().inv();
	let mut a = Matrix::<R2q, M, { L + M + 1 }>::zero();
	a.set_column(0, (b.ntt_inv().lift() * Z2q(100352)).get_ref());
	a[0][0] += Z2q(50177);
	a.set_columns(1, &(a0.lift() * Z2q(2)));
	for i in 0..M {
		a[i][1 + L + i] += Z2q(2);
	}
	let a: Matrix<Rb, M, { L + M + 1 }> = a.lift();
	let a = a.ntt();
	// random value for signing
	let y = Vector::<Rb, { L + M + 1 }>::random(ball_y, &mut rng);
	let y = y.ntt();
	// commitment
	let w: Vector<R2q, M> = (&a * &y).ntt_inv().lift();
	// challenge
	let c = R2q::random(ball_c, &mut rng);
	let cb: Rb = c.lift();
	// proof
	let s: Vector<Rb, { L + M + 1 }> = s.lift().lift();
	let z = y + s.ntt() * &cb.ntt();
	// verify
	let mut w2: Vector<R2q, M> = (&a * &z).ntt_inv().lift();
	w2[0] -= c * Z2q(50177);
	assert_eq!(w, w2);
}