// Copyright (C) 2019-2021 Aleo Systems Inc.
// This file is part of the snarkVM library.

// The snarkVM library is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// The snarkVM library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with the snarkVM library. If not, see <https://www.gnu.org/licenses/>.

use crate::errors::EncodingError;
use snarkvm_curves::traits::{
    pairing_engine::{AffineCurve, ProjectiveCurve},
    Group,
    MontgomeryParameters,
    TwistedEdwardsParameters,
};
use snarkvm_fields::{Field, LegendreSymbol, One, SquareRootField, Zero};
use snarkvm_utilities::{to_bytes, FromBytes, ToBytes};

use std::{cmp, marker::PhantomData, ops::Neg};

pub struct Elligator2<P: MontgomeryParameters + TwistedEdwardsParameters, G: Group + ProjectiveCurve> {
    _parameters: PhantomData<P>,
    _group: PhantomData<G>,
}

impl<P: MontgomeryParameters + TwistedEdwardsParameters, G: Group + ProjectiveCurve> Elligator2<P, G> {
    const A: P::BaseField = <P as MontgomeryParameters>::COEFF_A;
    const B: P::BaseField = <P as MontgomeryParameters>::COEFF_B;
    const D: P::BaseField = <P as TwistedEdwardsParameters>::COEFF_D;

    /// Returns the encoded group element for a given base field element.
    #[allow(clippy::many_single_char_names)]
    pub fn encode(input: &P::BaseField) -> Result<(<G as ProjectiveCurve>::Affine, bool), EncodingError> {
        // The input base field must be nonzero, otherwise inverses will fail.
        if input.is_zero() {
            return Err(EncodingError::InputMustBeNonzero);
        }

        // We define as convention for the input to be of high sign.
        let sign_high = input > &input.neg();
        let input = if sign_high { *input } else { input.neg() };

        // Compute the parameters for the alternate Montgomery form: v^2 == u^3 + A * u^2 + B * u.
        let (a, b) = {
            let a = Self::A * Self::B.inverse().unwrap();
            let b = P::BaseField::one() * Self::B.square().inverse().unwrap();
            (a, b)
        };

        // Compute the mapping from Fq to E(Fq) as an alternate Montgomery element (u, v).
        let (u, v) = {
            // Let r = element.
            let r = input;

            // Let u = D.
            // TODO (howardwu): change to 11.
            let u = Self::D;

            // Let ur2 = u * r^2;
            let ur2 = r.square() * u;

            {
                // Verify u is a quadratic nonresidue.
                #[cfg(debug_assertions)]
                assert!(u.legendre().is_qnr());

                // Verify 1 + ur^2 != 0.
                assert_ne!(P::BaseField::one() + ur2, P::BaseField::zero());

                // Verify A^2 * ur^2 != B(1 + ur^2)^2.
                let a2 = a.square();
                assert_ne!(a2 * ur2, (P::BaseField::one() + ur2).square() * b);
            }

            // Let v = -A / (1 + ur^2).
            let v = (P::BaseField::one() + ur2).inverse().unwrap() * (-a);

            // Let e = legendre(v^3 + Av^2 + Bv).
            let v2 = v.square();
            let v3 = v2 * v;
            let av2 = a * v2;
            let bv = b * v;
            let e = (v3 + (av2 + bv)).legendre();

            // Let x = ev - ((1 - e) * A/2).
            let two = P::BaseField::one().double();
            let x = match e {
                LegendreSymbol::Zero => -(a * two.inverse().unwrap()),
                LegendreSymbol::QuadraticResidue => v,
                LegendreSymbol::QuadraticNonResidue => (-v) - a,
            };

            // Let y = -e * sqrt(x^3 + Ax^2 + Bx).
            let x2 = x.square();
            let x3 = x2 * x;
            let ax2 = a * x2;
            let bx = b * x;
            let value = (x3 + (ax2 + bx)).sqrt().unwrap();
            let y = match e {
                LegendreSymbol::Zero => P::BaseField::zero(),
                LegendreSymbol::QuadraticResidue => -value,
                LegendreSymbol::QuadraticNonResidue => value,
            };

            (x, y)
        };

        // Ensure (u, v) is a valid alternate Montgomery element.
        {
            // Enforce v^2 == u^3 + A * u^2 + B * u
            let v2 = v.square();
            let u2 = u.square();
            let u3 = u2 * u;
            assert_eq!(v2, u3 + (a * u2) + (b * u));
        }

        // Convert the alternate Montgomery element (u, v) to Montgomery element (s, t).
        let (s, t) = {
            let s = u * Self::B;
            let t = v * Self::B;

            // Ensure (s, t) is a valid Montgomery element
            #[cfg(debug_assertions)]
            {
                // Enforce B * t^2 == s^3 + A * s^2 + s
                let t2 = t.square();
                let s2 = s.square();
                let s3 = s2 * s;
                assert_eq!(Self::B * t2, s3 + (Self::A * s2) + s);
            }

            (s, t)
        };

        // Convert the Montgomery element (s, t) to the twisted Edwards element (x, y).
        let (x, y) = {
            let x = s * t.inverse().unwrap();

            let numerator = s - P::BaseField::one();
            let denominator = s + P::BaseField::one();
            let y = numerator * denominator.inverse().unwrap();

            (x, y)
        };

        Ok((<G as ProjectiveCurve>::Affine::read(&to_bytes![x, y]?[..])?, sign_high))
    }

    #[allow(clippy::many_single_char_names)]
    pub fn decode(
        group_element: &<G as ProjectiveCurve>::Affine,
        sign_high: bool,
    ) -> Result<P::BaseField, EncodingError> {
        // The input group element must be nonzero, otherwise inverses will fail.
        if group_element.is_zero() {
            return Err(EncodingError::InputMustBeNonzero);
        }

        let x = P::BaseField::read(&to_bytes![group_element.to_x_coordinate()]?[..])?;
        let y = P::BaseField::read(&to_bytes![group_element.to_y_coordinate()]?[..])?;

        // Compute the parameters for the alternate Montgomery form: v^2 == u^3 + A * u^2 + B * u.
        let (a, b) = {
            let a = Self::A * Self::B.inverse().unwrap();
            let b = P::BaseField::one() * Self::B.square().inverse().unwrap();
            (a, b)
        };

        // Convert the twisted Edwards element (x, y) to the alternate Montgomery element (u, v)
        let (u_reconstructed, v_reconstructed) = {
            let numerator = P::BaseField::one() + y;
            let denominator = P::BaseField::one() - y;

            let u = numerator * (denominator.inverse().unwrap());
            let v = numerator * ((denominator * x).inverse().unwrap());

            // Ensure (u, v) is a valid Montgomery element
            #[cfg(debug_assertions)]
            {
                // Enforce B * v^2 == u^3 + A * u^2 + u
                let v2 = v.square();
                let u2 = u.square();
                let u3 = u2 * u;
                assert_eq!(Self::B * v2, u3 + (Self::A * u2) + u);
            }

            let u = u * Self::B.inverse().unwrap();
            let v = v * Self::B.inverse().unwrap();

            // Ensure (u, v) is a valid alternate Montgomery element.
            {
                // Enforce v^2 == u^3 + A * u^2 + B * u
                let v2 = v.square();
                let u2 = u.square();
                let u3 = u2 * u;
                assert_eq!(v2, u3 + (a * u2) + (b * u));
            }

            (u, v)
        };

        let x = u_reconstructed;

        // TODO (howardwu): change to 11.
        // Let u = D.
        let u = Self::D;

        {
            // Verify u is a quadratic nonresidue.
            #[cfg(debug_assertions)]
            assert!(u.legendre().is_qnr());

            // Verify that x != -A.
            assert_ne!(x, -a);

            // Verify that if y is 0, then x is 0.
            if y.is_zero() {
                assert!(x.is_zero());
            }

            // Verify -ux(x + A) is a residue.
            assert_eq!((-(u * x) * (x + a)).legendre(), LegendreSymbol::QuadraticResidue);
        }

        let exists_in_sqrt_fq2 = v_reconstructed.square().sqrt().unwrap() == v_reconstructed;

        let element = if exists_in_sqrt_fq2 {
            // Let value = sqrt(-x / ((x + A) * u)).
            let numerator = -x;
            let denominator = (x + a) * u;
            (numerator * denominator.inverse().unwrap()).sqrt().unwrap()
        } else {
            // Let value2 = sqrt(-(x + A) / ux)).
            let numerator = -x - a;
            let denominator = x * u;
            (numerator * denominator.inverse().unwrap()).sqrt().unwrap()
        };

        let element = if sign_high {
            cmp::max(element, -element)
        } else {
            cmp::min(element, -element)
        };

        #[cfg(debug_assertions)]
        assert!(&Self::encode(&element)?.0 == group_element);

        Ok(element)
    }
}