cryptix-bn254 0.1.0

use core::{ops::{Add, Neg}, fmt::Debug};

use cryptix_bigint::property::IsBigInt;
use cryptix_ecc::CurvePoint;
use cryptix_field::field::{montgomery::{MontgomeryOps, Montgomery}, Field, MulIdentity};
use cryptix_field::group::{AbelianGroup, CommunicativeAdd, Group, AddIdentity, AssociativeAdd};

use crate::{galoisfield::{fp2::Fp2Element, BN254, U256, FpElement, fp12::Fp12Element}, fp2};

use super::e1::BN254Fp;

/// Point using Jacobian projective coordinate on E1
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct BN254Fp2 {
    pub x: Fp2Element,
    pub y: Fp2Element,
    pub z: Fp2Element
}

impl Debug for BN254Fp2 {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "E2Point {{\n    ")?;
        write!(f, "x: fp2!(\n        \"")?;
        self.x.0.fmt(f)?;
        write!(f, "\", \n        \"")?;
        self.x.1.fmt(f)?;
        write!(f, "\"\n    ),\n    ")?;
        write!(f, "y: fp2!(\n        \"")?;
        self.y.0.fmt(f)?;
        write!(f, "\", \n        \"")?;
        self.y.1.fmt(f)?;
        write!(f, "\"\n    ),\n    ")?;
        write!(f, "z: fp2!(\n        \"")?;
        self.z.0.fmt(f)?;
        write!(f, "\", \n        \"")?;
        self.z.1.fmt(f)?;
        write!(f, "\"\n    )\n}}")
    }
}

impl BN254Fp2 {
    /// create a instance of e1 point at infinity
    pub const fn inf() -> Self {
        Self { x: Fp2Element::ZERO, y:  Fp2Element::ONE, z:  Fp2Element::ZERO }
    }
}

impl Add for BN254Fp2 {
    type Output = BN254Fp2;

    fn add(self, rhs: Self) -> Self::Output {
        if self.at_inf() { return rhs }
        if rhs.at_inf() { return self }

        let mut t = Self::inf();

        let t0 = self.x.mont_mul(rhs.x);
        let t1 = self.y.mont_mul(rhs.y);
        let t2 = self.z.mont_mul(rhs.z);
        let t3 = self.x + self.y;
        let t4 = rhs.x + rhs.y;            // 5
        
        let t3 = t3.mont_mul(t4);
        let t4 = t0 + t1;
        let t3 = t3 - t4;
        let t4 = self.y + self.z;
        t.x = rhs.y + rhs.z;                           // 10
        
        let t4 = t4.mont_mul(t.x);
        t.x = t1 + t2;
        let t4 = t4 - t.x;
        t.x = self.x + self.z;
        t.y = rhs.x + rhs.z;                           // 15
        
        t.x = t.x.mont_mul(t.y);
        t.y = t0 + t2;
        t.y = t.x - t.y;
        t.x = t0 + t0;
        let t0 = t.x + t0;                 // 20
        
        let t2 = t2.mul_3b();
        t.z = t1 + t2;
        let t1 = t1 - t2;
        t.y = t.y.mul_3b();
        t.x = t4.mont_mul(t.y);                        // 25
        
        let t2 = t3.mont_mul(t1);
        t.x = t2 - t.x;
        t.y = t.y.mont_mul(t0);
        let t1 = t1.mont_mul(t.z);
        t.y = t1 + t.y;                                // 30
        
        let t0 = t0.mont_mul(t3);
        t.z = t.z.mont_mul(t4) + t0;

        t
    }
}

impl AssociativeAdd for BN254Fp2 { }

impl CommunicativeAdd for BN254Fp2 { }

impl AddIdentity for BN254Fp2 {
    const ADD_IDENTITY: Self = Self::inf();
}

impl Neg for BN254Fp2 {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self { x: self.x, y: -self.y, z: self.z }
    }
}

impl Group for BN254Fp2 { }

impl AbelianGroup for BN254Fp2 { }

impl CurvePoint for BN254Fp2 {
    type MulScalar = U256;
    const GENERATOR: Self = BN254Fp2 {
        x: fp2!(
            "173F71DE61DF5B55EF7AB941DCD46337BB92F0233D01BC7CAF45AAAF63659E66",
            "10C76BD06630C4E91FCE93827A42E357413D1E26BC087B6A389BBB43D0D2F4E7"
        ),
        y: fp2!(
            "182C7A581C2D120D9215E25E884B2CC75A0FCF130C7D101225D0B13674482015",
            "1598CFFE781D1A201E0D27E0927F85986CE6262AA3A25124FBB0D5E1FCB5D7AA"
        ),
        z: fp2!(
            "212BA4F27FFFFFF5A2C62EFFFFFFFFCDB939FFFFFFFFFF8A15FFFFFFFFFFFF8E",
            "0000000000000000000000000000000000000000000000000000000000000000"
        )
    };

    fn double(self) -> Self {
        if self.at_inf() { return self }

        let mut t = Self::inf();

        let t0 = self.y.mont_sqr();
        t.z = t0 + t0;
        t.z = t.z + t.z;
        t.z = t.z + t.z;
        let t1 = self.y.mont_mul(self.z);        // 5

        let t2 = self.z.mont_sqr().mul_3b();
        t.x = t2.mont_mul(t.z);
        t.y = t0 + t2;
        t.z = t1.mont_mul(t.z);                      // 10
        
        let t1 = t2 + t2;
        let t2 = t1 + t2;
        let t0 = t0 - t2;
        t.y = t0.mont_mul(t.y) + t.x;                // 15
    
        let t1 = self.x.mont_mul(self.y);
        t.x = t0.mont_mul(t1);
        t.x = t.x + t.x;

        t
    }

    fn scalar_mul(self, k: U256) -> Self {
        const W: usize = 4;
        const LEN: usize = 16;
        // mask is 0xf
        const MASK: <U256 as IsBigInt>::Dig = 0xf;
        let mut kp = [BN254Fp2::inf(); LEN];

        kp[1] = self;
        kp[2] = self.double();

        for i in 3..LEN {
            // i is even
            if i & 1 == 0 {
                kp[i] = kp[i >> 1].double();
            } else {
                kp[i] = kp[i - 1].mix_add(self);
            }
            
            debug_assert!(kp[i].on_curve())
        }

        let ki = ((k[U256::DIG_LEN - 1] >> (W * (U256::DIG_BIT_LEN / W - 1))) & MASK) as usize;
        let mut q = kp[ki];

        for j in (0..(U256::DIG_BIT_LEN / W - 1)).rev() {
            let ki = ((k[U256::DIG_LEN - 1] >> (W * j)) & MASK) as usize;
            for _ in 0..W {
                q = q.double();
            }
            q = q + kp[ki]
        }

        for i in (0..(U256::DIG_LEN - 1)).rev() {
            for j in (0..(U256::DIG_BIT_LEN / W)).rev() {
                let ki = ((k[i] >> (W * j)) & MASK) as usize;
                for _ in 0..W {
                    q = q.double();
                }
                q = q + kp[ki]
            }
        }

        q
    }

    fn at_inf(&self) -> bool {
        self.z.is_zero()
    }

    /// check if y^2z == x^3 + z^3 * (1 - i)
    fn on_curve(&self) -> bool {
        let x3 = self.x.mont_sqr().mont_mul(self.x);
        
        let z3 = self.z.mont_sqr().mont_mul(self.z);
        let z3mi = z3.mont_mul(BN254::B_DIV_XI_MONT);

        let y2z = self.y.mont_sqr().mont_mul(self.z);

        x3 + z3mi == y2z
    }

    fn normalize(self) -> Self {
        let t = self.z.mont_inv();
        let x = self.x.mont_mul(t);
        let y = self.y.mont_mul(t);

        Self { x, y, z: Fp2Element(BN254::R_P, FpElement::ZERO) }
    }

    fn mont_form(self) -> Self {
        Self { 
            x: self.x.mont_form(), 
            y: self.y.mont_form(), 
            z: self.z.mont_form() 
        }
    }

    fn mont_rdc(self) -> Self {
        Self { 
            x: self.x.mont_rdc(), 
            y: self.y.mont_rdc(), 
            z: self.z.mont_rdc() 
        }
    }
}

impl BN254Fp2 {
    pub fn mix_add(self, mut rhs: Self) -> Self {
        // TODO: we show apply this restriction with type
        if rhs.z != BN254::R_P.into() {
            rhs = rhs.normalize();
        }

        if self.at_inf() { return rhs }
        if rhs.at_inf() { return self }

        let mut t = Self::inf();

        let t0 = self.x.mont_mul(rhs.x);
        let t1 = self.y.mont_mul(rhs.y);
        let t3 = rhs.x + rhs.y;
        let t4 = self.x + self.y;
        let t3 = t3.mont_mul(t4);             // 5
        
        let t4 = t0 + t1;
        let t3 = t3 - t4;
        let t4 = self.z.mont_mul(rhs.y);
        let t4 = self.y + t4;
        t.y = self.z.mont_mul(rhs.x);                     // 10

        t.y = t.y + self.x;
        t.x = t0 + t0;
        let t0 = t.x + t0;
        let t2 = self.z.mul_3b();
        t.z = t1 + t2;                                    // 15

        let t1 = t1 - t2;
        t.y = t.y.mul_3b();
        t.x = t.y.mont_mul(t4);
        let t2 = t3.mont_mul(t1);
        t.x = t2 - t.x;                                   // 20

        t.y = t.y.mont_mul(t0);
        let t1 = t1.mont_mul(t.z);
        t.y = t1 + t.y;
        let t0 = t0.mont_mul(t3);
        t.z = t.z.mont_mul(t4) + t0;                      // 25

        t
    }

    /// returns (2self, line_{self, self}(p))
    pub fn dbl_line(mut self, p: BN254Fp) -> (Self, Fp12Element) {
        let mut val = Fp12Element::ZERO;

        val.1.0 = self.x.mont_sqr();
        
        let a = self.x.mont_mul(self.y).hlv();
        let b = self.y.mont_sqr();
        let c = self.z.mont_sqr();

        val.0.1 = c.mont_mul(BN254::B_DIV_XI_MONT);
        let f = val.0.1 + val.0.1;
        val.0.1 = val.0.1 + f;                                         // E

        let f = val.0.1 + val.0.1 + val.0.1;               // F = 3E

        self.x = (b - f).mont_mul(a);                                  // X3
        val.0.0 = (self.y + self.z).mont_sqr() - b - c;
        let g = (b + f).hlv();                             // H

        self.z = b.mont_mul(val.0.0);                                  // Z3
        self.y = g.mont_sqr();

        let a = val.0.1.mont_sqr();
        self.y = self.y - a - a - a;

        val.0.0 = -val.0.0.mont_mul_fp(p.y);
        val.0.1 = val.0.1 - b;
        val.1.0 = val.1.0.mont_mul_fp(p.x);
        val.1.0 = val.1.0 + val.1.0 + val.1.0;
        
        (self, val)
    }

    pub fn mix_add_line(mut self, mut rhs: Self, p: BN254Fp) -> (Self, Fp12Element) {
        // TODO: we show apply this restriction with type
        if rhs.z != BN254::R_P.into() {
            rhs = rhs.normalize();
        }
        
        let mut val = Fp12Element::ZERO;
        
        let a = self.z.mont_mul(rhs.y);
        let b = self.z.mont_mul(rhs.x);

        let theta = self.y - a;
        let lambda = self.x - b;
        let c = theta.mont_sqr();
        let d = lambda.mont_sqr();
        let e = d.mont_mul(lambda);
        let f = self.z.mont_mul(c);
        let g = self.x.mont_mul(d);

        let h = e + f - g - g;
        self.x = lambda.mont_mul(h);                 // X3
        let i = self.y.mont_mul(e);
        self.y = (g - h).mont_mul(theta) - i;        // Y3
        self.z = self.z.mont_mul(e);                 // Z3
        
        val.0.0 = lambda.mont_mul_fp(p.y);
        val.1.0 = -theta.mont_mul_fp(p.x);
        val.0.1 = theta.mont_mul(rhs.x) - lambda.mont_mul(rhs.y);
        
        (self, val)
    }

    pub fn map_frob(self) -> Self {
        Self { 
            x: self.x.conjugate().mont_mul(BN254::XI_MONT[1]), 
            y: self.y.conjugate().mont_mul(BN254::XI_MONT[2]), 
            z: self.z.conjugate() 
        }
    }

    pub fn map2_frob(self) -> Self {
        Self { 
            x: self.x.mont_mul_fp(BN254::P2_FROB_MAP_VAL[1]), 
            y: self.y.mont_mul_fp(BN254::P2_FROB_MAP_VAL[2]), 
            z: self.z 
        }
    }
}