use bls12_381::*;
use rayon::prelude::*;
use crate::utils;

fn lambda(i: i64, shares: i64) -> (Scalar, bool) {
    assert!(i >= 0);
    assert!(shares > 0);
    assert!(shares >= i);
    let mut temp = 1.0;
    for j in 1..shares + 1 {
        if j != i {
            temp = temp * (-(j as f64) / ((i as f64) - (j as f64)));
        }
    }
    let scalar = Scalar::from(temp.round().abs() as u64);
    if temp < 0.0 {
        return (scalar, true);
    } else {
        return (scalar, false);
    }
}

pub fn recover_shares(shares: &Vec<Gt>, n: i64) -> Gt {
    shares.par_iter().enumerate().map(|(i, share)| {
        let (l, b) = lambda((i + 1) as i64, n);
        let mut temp = share * l;
        if b {
            temp = -temp
        }
        temp
    }).reduce(|| Gt::identity(), |a, b| a + b)
}

fn recover_shares_scalar(shares: &Vec<Scalar>, n: i64) -> Scalar {
    let mut target = Scalar::zero();
    for i in 0..shares.len() {
        let (l, b) = lambda((i + 1) as i64, n);
        let mut temp = shares[i] * l;
        if b {
            temp = -temp
        }
        target = target + temp;
    }
    return target;
}

fn eval_at(poly: &Vec<Gt>, x: i64) -> Gt {
    let mut y = Gt::identity();
    for j in 0..poly.len() {
        y = y + (poly[j] * Scalar::from((x.pow(j as u32)) as u64));
    }
    return y;
}

fn eval_at_scalar(poly: &Vec<Scalar>, x: i64) -> Scalar {
    let mut y = Scalar::zero();
    for j in 0..poly.len() {
        y = y + (poly[j] * Scalar::from((x.pow(j as u32)) as u64));
    }
    return y;
}

pub fn gen_shares(secret: Gt, n: i64, t: i64) -> Vec<Gt> {
    assert!(t > 0);
    assert!(n >= t);
    assert!(t <= n);
    let mut poly = Vec::new();
    poly.push(secret);
    (0..t - 1).for_each(|_| {
        poly.push(utils::random_gt());
    });
    let mut shares = Vec::new();
    for i in 0..n {
        shares.push(eval_at(&poly, i + 1));
    }
    return shares;
}

pub fn gen_shares_scalar(secret: Scalar, n: i64, t: i64) -> Vec<Scalar> {
    assert!(t > 0);
    assert!(n >= t);
    assert!(t <= n);
    let mut poly = Vec::new();
    poly.push(secret);
    (0..t - 1).for_each(|_| {
        poly.push(utils::random_scalar());
    });
    let mut shares = Vec::new();
    for i in 0..n {
        shares.push(eval_at_scalar(&poly, i + 1));
    }
    return shares;
}


#[cfg(test)]
fn test_shamir_scalar(t:i64,n:i64){
    let scalar = Scalar::from(329183724587293224);

    let scalarshares = gen_shares_scalar(scalar,n,t);

    let gt = utils::random_scalar();

    let shares = gen_shares_scalar(gt,n,t);
    let gt2 = recover_shares_scalar(&shares[0..((t+1) as usize)].to_vec(),n);
    assert_eq!(gt,gt2);

    let mut newershares = Vec::new();
    for i in 0..shares.len() {
        newershares.push(shares[i]+scalarshares[i]);
    }
    let gt3 = recover_shares_scalar(&newershares[0..((t+1) as usize)].to_vec(),n);
    assert_eq!(gt3,gt+scalar);
}


#[cfg(test)]
fn test_shamir_scalar_into_gt(t:i64,n:i64){
    let scalar = utils::random_scalar();

    let scalarshares = gen_shares_scalar(scalar,n,t);

    let gt2 = recover_shares_scalar(&scalarshares[0..((t+1) as usize)].to_vec(),n);
    

    assert_eq!(scalar,gt2);


    let gt_elm1 = utils::random_gt();

    let gt_shares1 :Vec<Gt> = scalarshares.iter().map(|x| gt_elm1 * x).collect();
    
    let gt41 = recover_shares(&gt_shares1[0..((t+1) as usize)].to_vec(),n);
    assert_eq!(gt_elm1*scalar,gt41);
}


#[test]
fn test(){
    let t = 3;
    let n = 4;
    test_shamir_scalar(t,n);
    test_shamir_scalar_into_gt(t,n);
}