use crate::{poly::*, params::*, discrete_gaussian::*, gadget::*, arith::*, util::*, number_theory::*};
use std::iter::once;

fn serialize_polymatrix(vec: &mut Vec<u8>, a: &PolyMatrixRaw) {
    for i in 0..a.rows * a.cols * a.params.poly_len {
        vec.extend_from_slice(&u64::to_ne_bytes(a.data[i]));
    }
}

fn serialize_vec_polymatrix(vec: &mut Vec<u8>, a: &Vec<PolyMatrixRaw>) {
    for i in 0..a.len() {
        serialize_polymatrix(vec, &a[i]);
    }
}

pub struct PublicParameters<'a> {
    v_packing: Vec<PolyMatrixNTT<'a>>,                  // Ws
    v_expansion_left: Option<Vec<PolyMatrixNTT<'a>>>,
    v_expansion_right: Option<Vec<PolyMatrixNTT<'a>>>,
    v_conversion: Option<Vec<PolyMatrixNTT<'a>>>,              // V
}

impl<'a> PublicParameters<'a> {
    pub fn init(params: &'a Params) -> Self {
        if params.expand_queries {
            PublicParameters { 
                v_packing: Vec::new(), 
                v_expansion_left: Some(Vec::new()), 
                v_expansion_right: Some(Vec::new()), 
                v_conversion: Some(Vec::new())
            }
        } else {
            PublicParameters { 
                v_packing: Vec::new(), 
                v_expansion_left: None,
                v_expansion_right: None,
                v_conversion: None,
            }
        }
    }

    fn from_ntt_alloc_vec(v: &Vec<PolyMatrixNTT<'a>>) -> Option<Vec<PolyMatrixRaw<'a>>> {
        Some(v.iter().map(from_ntt_alloc).collect())
    }

    fn from_ntt_alloc_opt_vec(v: &Option<Vec<PolyMatrixNTT<'a>>>) -> Option<Vec<PolyMatrixRaw<'a>>> {
        Some(v.as_ref()?.iter().map(from_ntt_alloc).collect())
    }

    pub fn to_raw(&self) -> Vec<Option<Vec<PolyMatrixRaw>>> {
        vec![
            Self::from_ntt_alloc_vec(&self.v_packing),
            Self::from_ntt_alloc_opt_vec(&self.v_expansion_left),
            Self::from_ntt_alloc_opt_vec(&self.v_expansion_right),
            Self::from_ntt_alloc_opt_vec(&self.v_conversion),
        ]
    }

    pub fn serialize(&self) -> Vec<u8> {
        let mut data = Vec::new();
        for v in self.to_raw().iter() {
            if v.is_some() {
                serialize_vec_polymatrix(&mut data, v.as_ref().unwrap());
            }
        }
        data
    }
}

pub struct Query<'a> {
    ct: Option<PolyMatrixRaw<'a>>,
    v_buf: Option<Vec<u64>>,
    v_ct: Option<Vec<PolyMatrixRaw<'a>>>,
}

impl<'a> Query<'a> {
    pub fn empty() -> Self {
        Query { ct: None, v_ct: None, v_buf: None }
    }

    pub fn serialize(&self) -> Vec<u8> {
        let mut data = Vec::new();
        if self.ct.is_some() {
            let ct = self.ct.as_ref().unwrap();
            serialize_polymatrix(&mut data, &ct);
        }
        if self.v_buf.is_some() {
            let v_buf = self.v_buf.as_ref().unwrap();
            data.extend(v_buf.iter().map(|x| { x.to_ne_bytes() }).flatten());
        }
        if self.v_ct.is_some() {
            let v_ct = self.v_ct.as_ref().unwrap();
            for x in v_ct {
                serialize_polymatrix(&mut data, x);
            }
        }
        data
    }
}

pub struct Client<'a> {
    params: &'a Params,
    sk_gsw: PolyMatrixRaw<'a>,
    sk_reg: PolyMatrixRaw<'a>,
    sk_gsw_full: PolyMatrixRaw<'a>,
    sk_reg_full: PolyMatrixRaw<'a>,
    dg: DiscreteGaussian,
    g: usize,
    stop_round: usize,
}

fn matrix_with_identity<'a> (p: &PolyMatrixRaw<'a>) -> PolyMatrixRaw<'a> {
    assert_eq!(p.cols, 1);
    let mut r = PolyMatrixRaw::zero(p.params, p.rows, p.rows + 1);
    r.copy_into(p, 0, 0);
    r.copy_into(&PolyMatrixRaw::identity(p.params, p.rows, p.rows), 0, 1);
    r
}

fn params_with_moduli(params: &Params, moduli: &Vec<u64>) -> Params {
    Params::init(
        params.poly_len, 
        moduli, 
        params.noise_width,
        params.n,
        params.pt_modulus,
        params.q2_bits,
        params.t_conv,
        params.t_exp_left,
        params.t_exp_right,
        params.t_gsw,
        params.expand_queries,
        params.db_dim_1,
        params.db_dim_2,
        params.instances,
        params.db_item_size,
    )
}

impl<'a> Client<'a> {
    pub fn init(params: &'a Params) -> Self {
        let sk_gsw_dims = params.get_sk_gsw();
        let sk_reg_dims = params.get_sk_reg();
        let sk_gsw = PolyMatrixRaw::zero(params, sk_gsw_dims.0, sk_gsw_dims.1);
        let sk_reg = PolyMatrixRaw::zero(params, sk_reg_dims.0, sk_reg_dims.1);
        let sk_gsw_full = matrix_with_identity(&sk_gsw);
        let sk_reg_full = matrix_with_identity(&sk_reg);
        let dg = DiscreteGaussian::init(params);

        let further_dims = params.db_dim_2;
        let num_expanded = 1usize << params.db_dim_1;
        let num_bits_to_gen = params.t_gsw * further_dims + num_expanded;
        let g = log2_ceil_usize(num_bits_to_gen);
        let stop_round = log2_ceil_usize(params.t_gsw * further_dims);
        Self {
            params,
            sk_gsw,
            sk_reg,
            sk_gsw_full,
            sk_reg_full,
            dg,
            g,
            stop_round,
        }
    }

    fn get_fresh_gsw_public_key(&mut self, m: usize) -> PolyMatrixRaw<'a> {
        let params = self.params;
        let n = params.n;

        let a = PolyMatrixRaw::random(params, 1, m);
        let e = PolyMatrixRaw::noise(params, n, m, &mut self.dg);
        let a_inv = -&a;
        let b_p = &self.sk_gsw.ntt() * &a.ntt();
        let b = &e.ntt() + &b_p;
        let p = stack(&a_inv, &b.raw());
        p
    }

    fn get_regev_sample(&mut self) -> PolyMatrixNTT<'a> {
        let params = self.params;
        let a = PolyMatrixRaw::random(params, 1, 1);
        let e = PolyMatrixRaw::noise(params, 1, 1, &mut self.dg);
        let b_p = &self.sk_reg.ntt() * &a.ntt();
        let b = &e.ntt() + &b_p;
        let mut p = PolyMatrixNTT::zero(params, 2, 1);
        p.copy_into(&(-&a).ntt(), 0, 0);
        p.copy_into(&b, 1, 0);
        p
    }

    fn get_fresh_reg_public_key(&mut self, m: usize) -> PolyMatrixNTT<'a> {
        let params = self.params;

        let mut p = PolyMatrixNTT::zero(params, 2, m);

        for i in 0..m {
            p.copy_into(&self.get_regev_sample(), 0, i);
        }
        p
    }

    fn encrypt_matrix_gsw(&mut self, ag: &PolyMatrixNTT<'a>) -> PolyMatrixNTT<'a> {
        let mx = ag.cols;
        let p = self.get_fresh_gsw_public_key(mx);
        let res = &(p.ntt()) + &(ag.pad_top(1));
        res
    }
    
    fn encrypt_matrix_reg(&mut self, a: &PolyMatrixNTT<'a>) -> PolyMatrixNTT<'a> {
        let m = a.cols;
        let p = self.get_fresh_reg_public_key(m);
        &p + &a.pad_top(1)
    }

    fn generate_expansion_params(&mut self, num_exp: usize, m_exp: usize) -> Vec<PolyMatrixNTT<'a>> {
        let params = self.params;
        let g_exp = build_gadget(params, 1, m_exp);
        let g_exp_ntt = g_exp.ntt();
        let mut res = Vec::new();

        for i in 0..num_exp {
            let t = (params.poly_len / (1 << i)) + 1;
            let tau_sk_reg = automorph_alloc(&self.sk_reg, t);
            let prod = &tau_sk_reg.ntt() * &g_exp_ntt;
            let w_exp_i = self.encrypt_matrix_reg(&prod);
            res.push(w_exp_i);
        }
        res
    }

    pub fn generate_keys(&mut self) -> PublicParameters {
        let params = self.params;
        self.dg.sample_matrix(&mut self.sk_gsw);
        self.dg.sample_matrix(&mut self.sk_reg);
        self.sk_gsw_full = matrix_with_identity(&self.sk_gsw);
        self.sk_reg_full = matrix_with_identity(&self.sk_reg);
        let sk_reg_ntt = to_ntt_alloc(&self.sk_reg);
        let m_conv = params.m_conv();

        let mut pp = PublicParameters::init(params);
        
        // Params for packing
        let gadget_conv = build_gadget(params, 1, m_conv);
        let gadget_conv_ntt = to_ntt_alloc(&gadget_conv);
        for i in 0..params.n {
            let scaled = scalar_multiply_alloc(&sk_reg_ntt, &gadget_conv_ntt);
            let mut ag = PolyMatrixNTT::zero(params, params.n, m_conv);
            ag.copy_into(&scaled, i, 0);
            let w = self.encrypt_matrix_gsw(&ag);
            pp.v_packing.push(w);
        }

        if params.expand_queries {
            // Params for expansion
            
            pp.v_expansion_left = Some(self.generate_expansion_params(self.g, params.t_exp_left));
            pp.v_expansion_right = Some(self.generate_expansion_params(self.stop_round + 1, params.t_exp_right));

            // Params for converison
            let g_conv = build_gadget(params, 2, 2 * m_conv);
            let sk_reg_ntt = self.sk_reg.ntt();
            let sk_reg_squared_ntt = &sk_reg_ntt * &sk_reg_ntt;
            pp.v_conversion = Some(Vec::from_iter(once(PolyMatrixNTT::zero(params, 2, 2 * m_conv))));
            for i in 0..2*m_conv {
                let sigma;
                if i % 2 == 0 {
                    let val = g_conv.get_poly(0, i)[0];
                    sigma = &sk_reg_squared_ntt * &single_poly(params, val).ntt();
                } else {
                    let val = g_conv.get_poly(1, i)[0];
                    sigma = &sk_reg_ntt * &single_poly(params, val).ntt();
                }
                let ct = self.encrypt_matrix_reg(&sigma);
                pp.v_conversion.as_mut().unwrap()[0].copy_into(&ct, 0, i);
            }
        }

        pp
    }

    // reindexes a vector of regev ciphertexts, to help server
    fn reorient_reg_ciphertexts(&self, out: &mut [u64], v_reg: &Vec<PolyMatrixNTT>) {
        let params = self.params;
        let poly_len = params.poly_len;
        let crt_count = params.crt_count;

        assert_eq!(crt_count, 2);
        assert!(log2(params.moduli[0]) <= 32);

        let num_reg_expanded = 1<<params.db_dim_1;
        let ct_rows = v_reg[0].rows;
        let ct_cols = v_reg[0].cols;

        assert_eq!(ct_rows, 2);
        assert_eq!(ct_cols, 1);

        for j in 0..num_reg_expanded {
            for r in 0..ct_rows {
                for m in 0.. ct_cols {
                    for z in 0..params.poly_len {
                        let idx_a_in = r * (ct_cols*crt_count*poly_len) + m * (crt_count*poly_len);
                        let idx_a_out =     z * (num_reg_expanded*ct_cols*ct_rows)
                                                + j * (ct_cols*ct_rows)
                                                + m * (ct_rows)
                                                + r;
                        let val1 = v_reg[j].data[idx_a_in + z] % params.moduli[0];
                        let val2 = v_reg[j].data[idx_a_in + params.poly_len + z] % params.moduli[1];

                        out[idx_a_out] = val1 | (val2 << 32);
                    }
                }
            }
        }
    }

    pub fn generate_query(&mut self, idx_target: usize) -> Query<'a> {
        let params = self.params;
        let further_dims = params.db_dim_2;
        let idx_dim0= idx_target / (1 << further_dims);
        let idx_further  = idx_target % (1 << further_dims);
        let scale_k = params.modulus / params.pt_modulus;
        let bits_per = get_bits_per(params, params.t_gsw);

        let mut query = Query::empty();
        if params.expand_queries {
            // pack query into single ciphertext
            let mut sigma = PolyMatrixRaw::zero(params, 1, 1);
            sigma.data[2*idx_dim0] = scale_k;
            for i in 0..further_dims as u64 {
                let bit: u64 = ((idx_further as u64) & (1 << i)) >> i;
                for j in 0..params.t_gsw {
                    let val = (1u64 << (bits_per * j)) * bit;
                    let idx = (i as usize) * params.t_gsw + (j as usize);
                    sigma.data[2*idx + 1] = val;
                }
            }
            let inv_2_g_first = invert_uint_mod(1 << self.g, params.modulus).unwrap();
            let inv_2_g_rest = invert_uint_mod(1 << (self.stop_round+1), params.modulus).unwrap();

            for i in 0..params.poly_len/2 {
                sigma.data[2*i]   = multiply_uint_mod(sigma.data[2*i], inv_2_g_first, params.modulus);
                sigma.data[2*i+1] = multiply_uint_mod(sigma.data[2*i+1], inv_2_g_rest, params.modulus);
            }

            query.ct = Some(from_ntt_alloc(&self.encrypt_matrix_reg(&to_ntt_alloc(&sigma))));
        } else {
            let num_expanded = 1 << params.db_dim_1;
            let mut sigma_v = Vec::<PolyMatrixNTT>::new();

            // generate regev ciphertexts
            let reg_cts_buf_words = num_expanded * 2 * params.poly_len;
            let mut reg_cts_buf = vec![0u64; reg_cts_buf_words];
            let mut reg_cts = Vec::<PolyMatrixNTT>::new();
            for i in 0..num_expanded {
                let value = ((i == idx_dim0) as u64) * scale_k;
                let sigma = PolyMatrixRaw::single_value(&params, value);
                reg_cts.push(self.encrypt_matrix_reg(&to_ntt_alloc(&sigma)));
            }
            // reorient into server's preferred indexing
            self.reorient_reg_ciphertexts(reg_cts_buf.as_mut_slice(), &reg_cts);

            // generate GSW ciphertexts
            for i in 0..further_dims {
                let bit = ((idx_further as u64) & (1 << (i as u64))) >> (i as u64);
                let mut ct_gsw = PolyMatrixNTT::zero(&params, 2, 2 * params.t_gsw);

                for j in 0..params.t_gsw {
                    let value = (1u64 << (bits_per * j)) * bit;
                    let sigma = PolyMatrixRaw::single_value(&params, value);
                    let sigma_ntt = to_ntt_alloc(&sigma);
                    let ct = &self.encrypt_matrix_reg(&sigma_ntt);
                    ct_gsw.copy_into(ct, 0, 2*j + 1);
                    let prod = &to_ntt_alloc(&self.sk_reg) * &sigma_ntt;
                    let ct = &self.encrypt_matrix_reg(&prod);
                    ct_gsw.copy_into(ct, 0, 2*j);
                }
                sigma_v.push(ct_gsw);
            }

            query.v_buf = Some(reg_cts_buf);
            query.v_ct = Some(sigma_v.iter().map(|x| { from_ntt_alloc(x) }).collect());
        }
        query
    }
    
    pub fn decode_response(&self, data: &[u8]) -> Vec<u8> {
        /*
            0. NTT over q2 the secret key

            1. read first row in q2_bit chunks
            2. read rest in q1_bit chunks
            3. NTT over q2 the first row
            4. Multiply the results of (0) and (3)
            5. Divide and round correctly
        */
        let params = self.params;
        let p = params.pt_modulus;
        let p_bits = log2_ceil(params.pt_modulus);
        let q1 = 4 * params.pt_modulus;
        let q1_bits = log2_ceil(q1) as usize;
        let q2 = Q2_VALUES[params.q2_bits as usize];
        let q2_bits = params.q2_bits as usize;

        let q2_params = params_with_moduli(params, &vec![q2]);

        // this only needs to be done during keygen
        let mut sk_gsw_q2 = PolyMatrixRaw::zero(&q2_params, params.n, 1);
        for i in 0..params.poly_len * params.n {
            sk_gsw_q2.data[i] = recenter(self.sk_gsw.data[i], params.modulus, q2);
        }
        let mut sk_gsw_q2_ntt = PolyMatrixNTT::zero(&q2_params, params.n, 1);
        to_ntt(&mut sk_gsw_q2_ntt, &sk_gsw_q2);

        let mut result = PolyMatrixRaw::zero(&params, params.instances * params.n, params.n);

        let mut bit_offs = 0;
        for instance in 0..params.instances {
            // this must be done during decoding
            let mut first_row = PolyMatrixRaw::zero(&q2_params, 1, params.n);
            let mut rest_rows = PolyMatrixRaw::zero(&params, params.n, params.n);
            for i in 0..params.n * params.poly_len {
                first_row.data[i] = read_arbitrary_bits(data, bit_offs, q2_bits);
                bit_offs += q2_bits;
            }
            for i in 0..params.n * params.n * params.poly_len {
                rest_rows.data[i] = read_arbitrary_bits(data, bit_offs, q1_bits);
                bit_offs += q1_bits;
            }

            let mut first_row_q2 = PolyMatrixNTT::zero(&q2_params, 1, params.n);
            to_ntt(&mut first_row_q2, &first_row);

            let sk_prod = (&sk_gsw_q2_ntt * &first_row_q2).raw();

            let q1_i64 = q1 as i64;
            let q2_i64 = q2 as i64;
            let p_i128 = p as i128;
            for i in 0..params.n * params.n * params.poly_len {
                let mut val_first = sk_prod.data[i] as i64;
                if val_first >= q2_i64/2 {
                    val_first -= q2_i64;
                }
                let mut val_rest = rest_rows.data[i] as i64;
                if val_rest >= q1_i64/2 {
                    val_rest -= q1_i64;
                }

                let denom = (q2 * (q1 / p)) as i64;

                let mut r = val_first * q1_i64;
                r += val_rest * q2_i64;

                // divide r by q2, rounding
                let sign: i64 = if r >= 0 { 1 } else { -1 };
                let mut res = ((r + sign*(denom/2)) as i128) / (denom as i128);
                res = (res + (denom as i128/p_i128)*(p_i128) + 2*(p_i128)) % (p_i128);
                let idx = instance * params.n * params.n * params.poly_len + i;
                result.data[idx] = res as u64;
            }
        }

        // println!("{:?}", result.data);
        let trials = params.n * params.n;
        let chunks = params.instances * trials;
        let bytes_per_chunk = f64::ceil(params.db_item_size as f64 / chunks as f64) as usize;
        let logp = log2(params.pt_modulus);
        let modp_words_per_chunk = f64::ceil((bytes_per_chunk * 8) as f64 / logp as f64) as usize;
        result.to_vec(p_bits as usize, modp_words_per_chunk)
    }
}