//! This module implements the Nova traits for pallas::Point, pallas::Scalar, vesta::Point, vesta::Scalar.
use crate::{
  poseidon::{PoseidonRO, PoseidonROCircuit},
  traits::{ChallengeTrait, CompressedGroup, Group},
};
use digest::{ExtendableOutput, Input};
use ff::Field;
use merlin::Transcript;
use num_bigint::BigInt;
use num_traits::Num;
use pasta_curves::{
  self,
  arithmetic::{CurveAffine, CurveExt, Group as OtherGroup},
  group::{Curve, Group as AnotherGroup, GroupEncoding},
  pallas, vesta, Ep, Eq,
};
use rand::SeedableRng;
use rand_chacha::ChaCha20Rng;
use rayon::prelude::*;
use sha3::Shake256;
use std::{io::Read, ops::Mul};

//////////////////////////////////////Pallas///////////////////////////////////////////////

/// A wrapper for compressed group elements that come from the pallas curve
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PallasCompressedElementWrapper {
  repr: [u8; 32],
}

impl PallasCompressedElementWrapper {
  /// Wraps repr into the wrapper
  pub fn new(repr: [u8; 32]) -> Self {
    Self { repr }
  }
}

impl Group for pallas::Point {
  type Base = pallas::Base;
  type Scalar = pallas::Scalar;
  type CompressedGroupElement = PallasCompressedElementWrapper;
  type PreprocessedGroupElement = pallas::Affine;
  type RO = PoseidonRO<Self::Base, Self::Scalar>;
  type ROCircuit = PoseidonROCircuit<Self::Base>;

  fn vartime_multiscalar_mul(
    scalars: &[Self::Scalar],
    bases: &[Self::PreprocessedGroupElement],
  ) -> Self {
    if scalars.len() >= 128 {
      pasta_msm::pallas(bases, scalars)
    } else {
      scalars
        .par_iter()
        .zip(bases)
        .map(|(scalar, base)| base.mul(scalar))
        .reduce(Ep::group_zero, |x, y| x + y)
    }
  }

  fn preprocessed(&self) -> Self::PreprocessedGroupElement {
    self.to_affine()
  }

  fn compress(&self) -> Self::CompressedGroupElement {
    PallasCompressedElementWrapper::new(self.to_bytes())
  }

  fn from_label(label: &'static [u8], n: usize) -> Vec<Self::PreprocessedGroupElement> {
    let mut shake = Shake256::default();
    shake.input(label);
    let mut reader = shake.xof_result();
    let mut gens: Vec<Self::PreprocessedGroupElement> = Vec::new();
    let mut uniform_bytes = [0u8; 32];
    for _ in 0..n {
      reader.read_exact(&mut uniform_bytes).unwrap();
      let hash = Ep::hash_to_curve("from_uniform_bytes");
      gens.push(hash(&uniform_bytes).to_affine());
    }
    gens
  }

  fn to_coordinates(&self) -> (Self::Base, Self::Base, bool) {
    let coordinates = self.to_affine().coordinates();
    if coordinates.is_some().unwrap_u8() == 1 {
      (*coordinates.unwrap().x(), *coordinates.unwrap().y(), false)
    } else {
      (Self::Base::zero(), Self::Base::zero(), true)
    }
  }

  fn get_curve_params() -> (Self::Base, Self::Base, BigInt) {
    let A = Self::Base::zero();
    let B = Self::Base::from(5);
    let order = BigInt::from_str_radix(
      "40000000000000000000000000000000224698fc0994a8dd8c46eb2100000001",
      16,
    )
    .unwrap();

    (A, B, order)
  }

  fn zero() -> Self {
    pallas::Point::group_zero()
  }

  fn get_generator() -> Self {
    pallas::Point::generator()
  }
}

impl ChallengeTrait for pallas::Scalar {
  fn challenge(label: &'static [u8], transcript: &mut Transcript) -> Self {
    let mut key: <ChaCha20Rng as SeedableRng>::Seed = Default::default();
    transcript.challenge_bytes(label, &mut key);
    let mut rng = ChaCha20Rng::from_seed(key);
    pallas::Scalar::random(&mut rng)
  }
}

impl CompressedGroup for PallasCompressedElementWrapper {
  type GroupElement = pallas::Point;

  fn decompress(&self) -> Option<pallas::Point> {
    Some(Ep::from_bytes(&self.repr).unwrap())
  }
  fn as_bytes(&self) -> &[u8] {
    &self.repr
  }
}

//////////////////////////////////////Vesta////////////////////////////////////////////////

/// A wrapper for compressed group elements that come from the vesta curve
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct VestaCompressedElementWrapper {
  repr: [u8; 32],
}

impl VestaCompressedElementWrapper {
  /// Wraps repr into the wrapper
  pub fn new(repr: [u8; 32]) -> Self {
    Self { repr }
  }
}

impl Group for vesta::Point {
  type Base = vesta::Base;
  type Scalar = vesta::Scalar;
  type CompressedGroupElement = VestaCompressedElementWrapper;
  type PreprocessedGroupElement = vesta::Affine;
  type RO = PoseidonRO<Self::Base, Self::Scalar>;
  type ROCircuit = PoseidonROCircuit<Self::Base>;

  fn vartime_multiscalar_mul(
    scalars: &[Self::Scalar],
    bases: &[Self::PreprocessedGroupElement],
  ) -> Self {
    if scalars.len() >= 128 {
      pasta_msm::vesta(bases, scalars)
    } else {
      scalars
        .par_iter()
        .zip(bases)
        .map(|(scalar, base)| base.mul(scalar))
        .reduce(Eq::group_zero, |x, y| x + y)
    }
  }

  fn compress(&self) -> Self::CompressedGroupElement {
    VestaCompressedElementWrapper::new(self.to_bytes())
  }

  fn preprocessed(&self) -> Self::PreprocessedGroupElement {
    self.to_affine()
  }

  fn from_label(label: &'static [u8], n: usize) -> Vec<Self::PreprocessedGroupElement> {
    let mut shake = Shake256::default();
    shake.input(label);
    let mut reader = shake.xof_result();
    let mut gens: Vec<Self::PreprocessedGroupElement> = Vec::new();
    let mut uniform_bytes = [0u8; 32];
    for _ in 0..n {
      reader.read_exact(&mut uniform_bytes).unwrap();
      let hash = Eq::hash_to_curve("from_uniform_bytes");
      gens.push(hash(&uniform_bytes).to_affine());
    }
    gens
  }

  fn to_coordinates(&self) -> (Self::Base, Self::Base, bool) {
    let coordinates = self.to_affine().coordinates();
    if coordinates.is_some().unwrap_u8() == 1 {
      (*coordinates.unwrap().x(), *coordinates.unwrap().y(), false)
    } else {
      (Self::Base::zero(), Self::Base::zero(), true)
    }
  }

  fn get_curve_params() -> (Self::Base, Self::Base, BigInt) {
    let A = Self::Base::zero();
    let B = Self::Base::from(5);
    let order = BigInt::from_str_radix(
      "40000000000000000000000000000000224698fc094cf91b992d30ed00000001",
      16,
    )
    .unwrap();

    (A, B, order)
  }

  fn zero() -> Self {
    vesta::Point::group_zero()
  }

  fn get_generator() -> Self {
    vesta::Point::generator()
  }
}

impl ChallengeTrait for vesta::Scalar {
  fn challenge(label: &'static [u8], transcript: &mut Transcript) -> Self {
    let mut key: <ChaCha20Rng as SeedableRng>::Seed = Default::default();
    transcript.challenge_bytes(label, &mut key);
    let mut rng = ChaCha20Rng::from_seed(key);
    vesta::Scalar::random(&mut rng)
  }
}

impl CompressedGroup for VestaCompressedElementWrapper {
  type GroupElement = vesta::Point;

  fn decompress(&self) -> Option<vesta::Point> {
    Some(Eq::from_bytes(&self.repr).unwrap())
  }
  fn as_bytes(&self) -> &[u8] {
    &self.repr
  }
}