extern crate ff;
use crate::matrix::Matrix;
use ff::PrimeField;
use rand_core::{CryptoRng, RngCore};
use std::fmt::Display;

mod matrix;

struct HashKey<PF>
where
    PF: PrimeField,
{
    synthetic_max: u64,
    threshold: u64,
    key: Vec<Vec<PF>>,
}

#[derive(Debug)]
struct Hash<PF>(Vec<PF>)
where
    PF: PrimeField;

impl<PF> HashKey<PF>
where
    PF: PrimeField,
{
    pub fn new<R: RngCore + CryptoRng + Copy>(
        rng: R,
        synthetic_max: u64,
        threshold: u64,
    ) -> HashKey<PF> {
        let mut k = vec![];
        for _i in 0..synthetic_max {
            let mut p_i = vec![];
            for _j in 0..threshold {
                p_i.push(PF::random(rng));
            }
            k.push(p_i);
        }
        HashKey {
            synthetic_max,
            threshold,
            key: k,
        }
    }

    pub fn generate_hash(&self, value: PF) -> Hash<PF> {
        let mut hash = vec![value];

        for i in 0..self.synthetic_max as usize {
            // Here evaluate our polynomial key[i]*X key[i]*X^2 key[i]*X^3 etc....
            let mut result = PF::zero();
            for j in 0..self.threshold as usize {
                result = result + (self.key[i][j].clone() * value.pow_vartime(&[j as u64]));
            }
            hash.push(result);
        }
        Hash(hash)
    }

    pub fn random<R: RngCore + CryptoRng + Copy>(&self, rng: R) -> Hash<PF> {
        let mut hash = vec![PF::random(rng)];
        for _i in 0..self.synthetic_max as usize {
            hash.push(PF::random(rng));
        }
        Hash(hash)
    }
}

struct Solver<PF>
where
    PF: PrimeField,
{
    synthetic_max: u64,
    threshold: u64,
    expanded_hashes: Vec<Hash<PF>>,
}

impl<PF> Solver<PF>
where
    PF: PrimeField + PartialOrd + Display,
{
    pub fn new(synthetic_max: u64, threshold: u64) -> Solver<PF> {
        Solver {
            synthetic_max,
            threshold,
            expanded_hashes: vec![],
        }
    }

    /// Add a new hash to be evaluated by the solver...
    pub fn add_hash(&mut self, hash: Hash<PF>) {
        println!("Hash {}: {:?}", self.expanded_hashes.len(), hash);

        // To expand a Hash we take DHF(x0,p1..ps) and produce DHF(1,x0,x0^1..x0^t-1,p1..ps)
        let mut expanded_hash = vec![];
        for i in 0..self.threshold {
            expanded_hash.push(hash.0[0].pow_vartime(&[i]))
        }
        for element in hash.0.iter().skip(1) {
            expanded_hash.push(element.clone())
        }
        self.expanded_hashes.push(Hash(expanded_hash))
    }

    /// Solve the system and return the indexes of the true hashes (or error if the system is unsolvable)
    pub fn attempt_solve(&self) -> Result<Vec<usize>, ()> {
        // Arrange the hashes into an augmented for matrix...
        let m = (self.synthetic_max + self.threshold) as usize;
        let mut augmented_matrix =
            Matrix::new(m + self.expanded_hashes.len(), self.expanded_hashes.len());
        for j in 0..self.expanded_hashes.len() {
            for i in 0..m {
                augmented_matrix.update(i, j, self.expanded_hashes[j as usize].0[i as usize]);
            }
        }
        // Append the identity matrix...
        for (i, r) in (m..m + self.expanded_hashes.len()).enumerate() {
            augmented_matrix.update(r as usize, i, PF::one());
        }

        // Transpose for row reduction...
        augmented_matrix = augmented_matrix.transpose();

        // Pretty Print...
        for i in 0..augmented_matrix.rows() {
            for j in 0..augmented_matrix.cols() {
                print!("{} ", augmented_matrix.at(i, j as usize));
            }
            println!(";")
        }

        // Do the actual row reduction...
        let cols = augmented_matrix.cols();
        let rows = augmented_matrix.rows();

        let mut h = 0;
        let mut k = 0;

        while h < rows && k < cols {
            let mut i_max = h;
            let mut i_max_v = augmented_matrix.at(h, k);
            for x in h + 1..rows {
                if augmented_matrix.at(x, k) > i_max_v {
                    i_max = x;
                    i_max_v = augmented_matrix.at(x, k)
                }
            }
            if augmented_matrix.at(i_max, k).is_zero() {
                k += 1;
            } else {
                augmented_matrix.swap_rows(h, i_max);
                for i in h + 1..rows {
                    let f = augmented_matrix.at(i, k) * augmented_matrix.at(h, k).invert().unwrap();
                    augmented_matrix.update(i, k, PF::zero());
                    for j in k + 1..cols {
                        let val = augmented_matrix.at(i, j) - augmented_matrix.at(h, j) * f;
                        augmented_matrix.update(i, j, val);
                    }
                }
                h += 1;
                k += 1;
            }
        }

        println!();
        println!();
        println!();

        // Transpose back to column form...
        augmented_matrix = augmented_matrix.transpose();

        // Pretty Print for checking...
        for i in 0..augmented_matrix.rows() {
            for j in 0..augmented_matrix.cols() {
                print!("{} ", augmented_matrix.at(i, j as usize));
            }
            println!(";");
            if i + 1 == m {
                println!("----------------------------------------");
            }
        }

        // Calculate the Kernel Basis Vectors...these are the columns where the
        // Original matrix is now all 0s
        let mut count = 0;
        for col in (0..self.expanded_hashes.len()).rev() {
            let mut is_null = true;
            for row in 0..m {
                if !augmented_matrix.at(row, col).is_zero() {
                    is_null = false;
                    break;
                }
            }
            if is_null == false {
                break;
            } else {
                count += 1;
            }
        }

        if count == 0 {
            // Unsolvable...
            return Err(());
        }

        // We found some vectors...now we can derive a solution...
        println!("Found {} Kernel Basis Vectors...", count);
        let mut solution = vec![];

        // Check all the basis vectors...
        let basis_state = self.expanded_hashes.len() - count;
        for i in 0..self.expanded_hashes.len() {
            let mut is_zero = true;
            for b in 0..count {
                is_zero = is_zero & augmented_matrix.at(i + m, basis_state + b).is_zero()
            }
            if !is_zero {
                solution.push(i)
            }
        }

        println!("Solution: {:?}", solution);
        Ok(solution)
    }
}

#[cfg(test)]
mod tests {
    use crate::ff::Field;
    use crate::{HashKey, Solver};
    use ff::PrimeField;
    use rand_core::OsRng;
    use std::fmt::{Display, Formatter};

    // Generate a Prime Field to Test with...
    // This is way larger than the 64 bit prime the paper specifies, but wth
    #[derive(PrimeField)]
    #[PrimeFieldModulus = "52435875175126190479447740508185965837690552500527637822603658699938581184513"]
    #[PrimeFieldGenerator = "7"]
    #[PrimeFieldReprEndianness = "little"]
    struct Fp([u64; 4]);

    impl Display for Fp {
        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
            write!(f, "{:10}", self.0[0])
        }
    }

    #[test]
    fn it_works() {
        let rng = OsRng;
        let s = 10u64;
        let t = 30u64;
        let dhf = HashKey::<Fp>::new(rng, s, t);

        let mut solver = Solver::new(s, t);
        for i in 0..40 {
            // These are the indexes which are to be random...you can try swapping them around..
            // REMEMBER: the number of true hash needs to be AT LEAST `t+1` AND the number of fake hashes
            // must be AT-MOST `s`.
            // If |random hashes| > s this this procedure will *not* work
            // There is an extended algorithm for extract assuming |true hashes| > t though - we do not implement it.
            if i != 2 && i != 5 && i != 8 && i < 34 {
                let x0 = Fp::random(rng);
                let hash = dhf.generate_hash(x0);
                solver.add_hash(hash);
            } else {
                solver.add_hash(dhf.random(rng));
            }
        }

        assert_eq!(
            solver.attempt_solve().unwrap(),
            vec![
                0, 1, 3, 4, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
                25, 26, 27, 28, 29, 30, 31, 32, 33
            ]
        );
    }
}