// main.rs
use rand_distr::{Normal, Distribution};
use rand::{thread_rng, Rng};
use std::process;
use rand::prelude::SliceRandom;

fn main() {
    let mut n = 5;
    let mut B = alpha(n as f64);
    let mut p = 7;
    let mut m = 5;
    let mut private_key = private_key(m, n, p);
    let mut public_key = public_key(m, n, p, B, &private_key);
    let mut message = vec![1 as u64,0,0,1,1,0,1];
    let mut encrypted = encrypt_vector(&message, m, n, &public_key, p);
    let mut decrypted = decrypt_vector(&encrypted, &private_key, p);
    println!("decrypted = {:?}", decrypted);
}

fn encrypt_vector(message:&Vec<u64>, m:u64, n:u64, public_key:&Vec<(Vec<u64>, u64)>, p:u64) -> Vec<(Vec<u64>, u64)> {
    let mut encrypted_bits = Vec::new();
    for bit in message {
        let mut encrypted_bit = encrypt_bit(*bit, m, n, public_key, p);
        encrypted_bits.push(encrypted_bit);
    }
    encrypted_bits
}
fn decrypt_vector(encrypted_bits:&Vec<(Vec<u64>, u64)>, private_key:&Vec<u64>, p:u64) -> Vec<u64> {
    let mut decrypted_bits = Vec::new();
    for bit in encrypted_bits {
        let mut decrypted_bit = decrypt_bit(bit, private_key, p);
        decrypted_bits.push(decrypted_bit);
    }
    decrypted_bits
}

fn decrypt_bit(encrypted_bit:&(Vec<u64>, u64), private_key:&Vec<u64>, p:u64) -> u64 {
    let mut a = encrypted_bit.0.clone();
    let mut b = encrypted_bit.1.clone();
    let mut inner_product = inner_product(&a, &private_key, p);
    let mut point = b as i64 - inner_product as i64;
    //println!("point = {}", point);
    let mut q = (p as f64 / 2 as f64).floor();
    //println!("q = {}", q);
    let mut distance_to_0 = (point as i64 - 0 as i64).abs();
    let mut distance_to_q = (point.rem_euclid(p as i64) as i64 - q as i64).abs();
    //println!("distance to 0 = {}", distance_to_0);
    //println!("distance to q = {}", distance_to_q);
    if distance_to_0 < distance_to_q {
        return 0;
    } else {
        return 1;
    }
}

fn encrypt_bit(bit:u64, m:u64, n:u64, public_key:&Vec<(Vec<u64>, u64)>, p:u64) -> (Vec<u64>, u64) {
    let mut subset = choose_subset(m); 
    let mut x = vec![0; n as usize];
    let mut y = 0 as u64;
    for index in subset {
        let mut a = public_key[index as usize].0.clone();
        let mut b = public_key[index as usize].1.clone();
        x = add_vector(&x, &a, p);
        y += b;
        
    }
    if bit == 0 {
        y = y.rem_euclid(p);
        return (x,y);
    } else {
        y += (p as f64 / 2 as f64).floor() as u64;
        y = y.rem_euclid(p);
        return (x, y);
    }
}

fn add_vector(a:&Vec<u64>, b:&Vec<u64>, p:u64) -> Vec<u64> {
    let mut sum = Vec::new();
    for i in 0..a.len() {
        sum.push((a[i] + b[i]).rem_euclid(p));
    }
    sum
}

fn choose_subset(m:u64) -> Vec<u64> {
    let mut set = Vec::new();
    let mut num = rand::thread_rng().gen_range(0..2_u64.pow(m as u32)) as u64;
    for i in 0..m {
        if num%2==1 {
            set.push(i);
        }
        num = num >> 1;
    }
    set
}

fn private_key(m:u64, n:u64, p:u64) -> Vec<u64> {
    let mut private_key = draw_vector_from_Znp(n,p);
    private_key
}

fn public_key(m:u64, n:u64, p:u64, B:f64, private_key:&Vec<u64>) -> Vec<(Vec<u64>, u64)> {
    let mut s = private_key.clone();
    let mut public_key = Vec::new();
    for i in 0..m {
        let mut a = draw_vector_from_Znp(n, p);
        let mut b = inner_product(&a, &s, p) + draw_element_from_X(p, B);
        public_key.push((a,b));
    }
    public_key
} 

fn inner_product(a:&Vec<u64>, b:&Vec<u64>, p:u64) -> u64 {
    let mut sum = 0;
    if a.len() != b.len() {
        println!("Error in inner product - lengths were different");
        process::exit(1);
    }
    for i in 0..a.len() {
        sum += a[i] * b[i];
    }
    return (sum).rem_euclid(p);
}

fn draw_vector_from_Znp(n:u64, p:u64) -> Vec<u64> {
    let mut vec = Vec::new();
    for i in 0..n {
        vec.push(rand::thread_rng().gen_range(0..p));
    }
    vec
}

fn draw_element_from_X(p:u64, B:f64) -> u64 {
    let mut mean = 0.0;
    let mut std_dev = B/(2 as f64*std::f64::consts::PI).sqrt(); // mean 0, standard deviation B/sqrt(2pi)
    let normal = Normal::new(mean, std_dev).unwrap(); // create normal distribution 
    let mut v = normal.sample(&mut rand::thread_rng()); // sample from normal distribution
    let mut w = (v).rem_euclid(1.0); // reduce result modulo 1
    let mut x = w * p as f64; // multiply by p
    let mut y = (x.round()).rem_euclid(p as f64); // round to nearest integer modulo p
    return y as u64;
}

// return B E R+
fn alpha(n:f64) -> f64 {
    return 1 as f64/(n.sqrt() * n.log2() * n.log2());
}