//! Example of simplified Dilithium with q=8380417

use gwrizienn::{
	matrix::Matrix,
	ntt::{Ntt, NttInv},
	ring::Ring,
	vector::Vector,
	*,
};

// Implement Zq
ring!(Zq, u32, u64, i64, 8380417);

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

fn high_bits<const N: usize>(mut v: Vector<Rq, N>) -> Vector<Rq, N> {
	for vi in v.0.iter_mut() {
		for vij in vi.0.iter_mut() {
			vij.0 -= vij.0 % 190464;
		}
	}
	v
}

fn main() {
	let mut rng = rand::thread_rng();
	let uniform = Zq::uniform();
	let ball_c = Zq::uniform_ball(1);
	let ball_s = Zq::uniform_ball(2);
	let ball_y = Zq::uniform_ball(131071);

	// generate secret key
	let a = Matrix::<Rq, 4, 4>::random(uniform, &mut rng);
	let s1 = Vector::<Rq, 4>::random(ball_s, &mut rng);
	let s2 = Vector::<Rq, 4>::random(ball_s, &mut rng);
	// random value for signing
	let y = Vector::<Rq, 4>::random(ball_y, &mut rng);
	// challenge
	let c = Rq::random(ball_c, &mut rng);

	// use NTT for fast multiplication
	let a = a.ntt();
	let s1 = s1.ntt();
	let s2 = s2.ntt();
	let y = y.ntt();
	let c = c.ntt();

	// generate public key
	let t = &a * &s1 + s2;
	// commitment
	let w = &a * &y;
	// proof
	let z = y + s1 * &c;

	// verify
	assert_eq!(
		high_bits((&a * &z - t * &c).ntt_inv()),
		high_bits(w.ntt_inv())
	);

	// let uniform = Zq::uniform();
	// let ball_c = Zq::uniform_ball(1);
	// let ball_s = Zq::uniform_ball(2);
	// let ball_y = Zq::uniform_ball(131071);
	//
	// let a = Matrix::<Rq, 4, 4>::random(uniform, &mut rng).ntt();
	// let s1 = Vector::<Rq, 4>::random(ball_s, &mut rng).ntt();
	// let s2 = Vector::<Rq, 4>::random(ball_s, &mut rng).ntt();
	// let y = Vector::<Rq, 4>::random(ball_y, &mut rng).ntt();
	// let c = Rq::random(ball_c, &mut rng).ntt();
	//
	// let t = &a * &s1 + s2;
	// let w = &a * &y;
	// let z = y + s1 * &c;
}