#include "PrivateKey.hpp"

BGNPrivateKey::BGNPrivateKey(const BGNPrivateKey& other)
    : a1(other.a1), b1(other.b1), c1(other.c1), d1(other.d1),
    a2(other.a2), b2(other.b2), c2(other.c2), d2(other.d2),
    pi_1_curvegen(other.pi_1_curvegen), pi_2_twistgen(other.pi_2_twistgen),
    pi_T_pairgen(other.pi_T_pairgen), curve_memoizer(other.curve_memoizer),
    twist_memoizer(other.twist_memoizer), pair_memoizer(other.pair_memoizer),
    curve_max_checked(other.curve_max_checked),
    twist_max_checked(other.twist_max_checked),
    pair_max_checked(other.pair_max_checked)
{ }


Scalar BGNPrivateKey::decrypt(const CurveBipoint& ciphertext)
{
    CurveBipoint pi_1_ciphertext = pi_1(ciphertext); 

    auto lookup = curve_memoizer.find(pi_1_ciphertext);
    if (lookup != curve_memoizer.end())
    {
        return lookup->second;
    }

    curve_max_checked++;
    CurveBipoint i = pi_1_curvegen * curve_max_checked;

    while (i != pi_1_ciphertext)
    {
        curve_memoizer[i] = curve_max_checked;

        i = i + pi_1_curvegen;
        curve_max_checked++;
    }
    curve_memoizer[i] = curve_max_checked;

    return curve_max_checked;
}

Scalar BGNPrivateKey::decrypt(const TwistBipoint& ciphertext)
{
    TwistBipoint pi_2_ciphertext = pi_2(ciphertext);

    auto lookup = twist_memoizer.find(pi_2_ciphertext);
    if (lookup != twist_memoizer.end())
    {
        return lookup->second;
    }

    twist_max_checked++;
    TwistBipoint i = pi_2_twistgen * twist_max_checked;
    while (i != pi_2_ciphertext)
    {
        twist_memoizer[i] = twist_max_checked;

        i = i + pi_2_twistgen;
        twist_max_checked++;
    }
    twist_memoizer[i] = twist_max_checked;

    return twist_max_checked;
}


Scalar BGNPrivateKey::decrypt(const Quadripoint& ciphertext)
{
    Quadripoint pi_T_ciphertext = pi_T(ciphertext); 

    auto lookup = pair_memoizer.find(pi_T_ciphertext);
    if (lookup != pair_memoizer.end())
    {
        return lookup->second;
    }

    pair_max_checked++;
    Quadripoint i = pi_T_pairgen * pair_max_checked;
    while (i != pi_T_ciphertext)
    {
        pair_memoizer[i] = pair_max_checked;

        i = i + pi_T_pairgen;
        pair_max_checked++;
    }
    pair_memoizer[i] = pair_max_checked;

    return pair_max_checked;
}

std::ostream& operator<<(std::ostream& os, const BGNPrivateKey& output)
{
    os << output.a1 << output.b1 << output.c1 << output.d1;
    os << output.a2 << output.b2 << output.c2 << output.d2;
    os << output.pi_1_curvegen;
    os << output.pi_2_twistgen;
    os << output.pi_T_pairgen;

    return os;
}

std::istream& operator>>(std::istream& is, BGNPrivateKey& input)
{
    is >> input.a1 >> input.b1 >> input.c1 >> input.d1;
    is >> input.a2 >> input.b2 >> input.c2 >> input.d2;
    is >> input.pi_1_curvegen;
    is >> input.pi_2_twistgen;
    is >> input.pi_T_pairgen;

    input.curve_max_checked = Scalar(0);
    input.twist_max_checked = Scalar(0);
    input.pair_max_checked = Scalar(0);

    input.curve_memoizer.clear();
    input.twist_memoizer.clear();
    input.pair_memoizer.clear();

    input.curve_memoizer[input.pi_1_curvegen * input.curve_max_checked] = input.curve_max_checked;
    input.twist_memoizer[input.pi_2_twistgen * input.twist_max_checked] = input.twist_max_checked;
    input.pair_memoizer[input.pi_T_pairgen * input.pair_max_checked] = input.pair_max_checked;

    return is;
}

BGNPrivateKey::BGNPrivateKey()
{ }

void BGNPrivateKey::set(const BGNPublicKey& pub_key, const Scalar& a1, const Scalar& b1, const Scalar& c1, const Scalar& d1, const Scalar& a2, const Scalar& b2, const Scalar& c2, const Scalar& d2)
{
    this->a1 = a1;
    this->b1 = b1;
    this->c1 = c1;
    this->d1 = d1;

    this->a2 = a2;
    this->b2 = b2;
    this->c2 = c2;
    this->d2 = d2;

    this->pi_1_curvegen = pi_1(pub_key.get_bipoint_curvegen());
    this->pi_2_twistgen = pi_2(pub_key.get_bipoint_twistgen());
    this->pi_T_pairgen  = pi_T(pairing(pub_key.get_bipoint_curvegen(), pub_key.get_bipoint_twistgen()));

    this->curve_max_checked = Scalar(0);
    this->twist_max_checked = Scalar(0);
    this->pair_max_checked = Scalar(0);

    this->curve_memoizer[this->pi_1_curvegen * this->curve_max_checked] = this->curve_max_checked;
    this->twist_memoizer[this->pi_2_twistgen * this->twist_max_checked] = this->twist_max_checked;
    this->pair_memoizer[this->pi_T_pairgen * this->pair_max_checked] = this->pair_max_checked;
}

CurveBipoint BGNPrivateKey::pi_1(const CurveBipoint& input) const
{
    CurveBipoint retval;
    curvepoint_fp_t temp0, temp1;
    

    b1.mult(temp0, input[0]);
    c1.mult(temp0, temp0);
    curvepoint_fp_neg(temp0, temp0);

    a1.mult(temp1, input[1]);
    c1.mult(temp1, temp1);

    curvepoint_fp_add_vartime(temp0, temp0, temp1);
    curvepoint_fp_set(retval[0], temp0);


    b1.mult(temp0, input[0]);
    d1.mult(temp0, temp0);
    curvepoint_fp_neg(temp0, temp0);

    a1.mult(temp1, input[1]);
    d1.mult(temp1, temp1);

    curvepoint_fp_add_vartime(temp0, temp0, temp1);
    curvepoint_fp_set(retval[1], temp0);


    return retval;
}

TwistBipoint BGNPrivateKey::pi_2(const TwistBipoint& input) const
{
    TwistBipoint retval;
    twistpoint_fp2_t temp0, temp1;
    

    b2.mult(temp0, input[0]);
    c2.mult(temp0, temp0);
    twistpoint_fp2_neg(temp0, temp0);

    a2.mult(temp1, input[1]);
    c2.mult(temp1, temp1);

    twistpoint_fp2_add_vartime(temp0, temp0, temp1);
    twistpoint_fp2_set(retval[0], temp0);


    b2.mult(temp0, input[0]);
    d2.mult(temp0, temp0);
    twistpoint_fp2_neg(temp0, temp0);

    a2.mult(temp1, input[1]);
    d2.mult(temp1, temp1);

    twistpoint_fp2_add_vartime(temp0, temp0, temp1);
    twistpoint_fp2_set(retval[1], temp0);


    return retval;     
}

Quadripoint BGNPrivateKey::pi_T(const Quadripoint& input) const
{
    Quadripoint retval;
    fp12e_t temp0, temp1, temp2, temp3;


    b1.mult(temp0, input[0]);
    c1.mult(temp0, temp0);
    b2.mult(temp0, temp0);
    c2.mult(temp0, temp0);
    
    b1.mult(temp1, input[1]);
    c1.mult(temp1, temp1);
    a2.mult(temp1, temp1);
    c2.mult(temp1, temp1);
    fp12e_invert(temp1, temp1);

    a1.mult(temp2, input[2]);
    c1.mult(temp2, temp2);
    b2.mult(temp2, temp2);
    c2.mult(temp2, temp2);
    fp12e_invert(temp2, temp2);
    
    a1.mult(temp3, input[3]);
    c1.mult(temp3, temp3);
    a2.mult(temp3, temp3);
    c2.mult(temp3, temp3);

    fp12e_mul(temp0, temp0, temp1);
    fp12e_mul(temp1, temp2, temp3);
    fp12e_mul(temp0, temp0, temp1);
    fp12e_set(retval[0], temp0);


    b1.mult(temp0, input[0]);
    c1.mult(temp0, temp0);
    b2.mult(temp0, temp0);
    d2.mult(temp0, temp0);
    
    b1.mult(temp1, input[1]);
    c1.mult(temp1, temp1);
    a2.mult(temp1, temp1);
    d2.mult(temp1, temp1);
    fp12e_invert(temp1, temp1);

    a1.mult(temp2, input[2]);
    c1.mult(temp2, temp2);
    b2.mult(temp2, temp2);
    d2.mult(temp2, temp2);
    fp12e_invert(temp2, temp2);
    
    a1.mult(temp3, input[3]);
    c1.mult(temp3, temp3);
    a2.mult(temp3, temp3);
    d2.mult(temp3, temp3);

    fp12e_mul(temp0, temp0, temp1);
    fp12e_mul(temp1, temp2, temp3);
    fp12e_mul(temp0, temp0, temp1);
    fp12e_set(retval[1], temp0);
    

    b1.mult(temp0, input[0]);
    d1.mult(temp0, temp0);
    b2.mult(temp0, temp0);
    c2.mult(temp0, temp0);
    
    b1.mult(temp1, input[1]);
    d1.mult(temp1, temp1);
    a2.mult(temp1, temp1);
    c2.mult(temp1, temp1);
    fp12e_invert(temp1, temp1);

    a1.mult(temp2, input[2]);
    d1.mult(temp2, temp2);
    b2.mult(temp2, temp2);
    c2.mult(temp2, temp2);
    fp12e_invert(temp2, temp2);
    
    a1.mult(temp3, input[3]);
    d1.mult(temp3, temp3);
    a2.mult(temp3, temp3);
    c2.mult(temp3, temp3);
    
    fp12e_mul(temp0, temp0, temp1);
    fp12e_mul(temp1, temp2, temp3);
    fp12e_mul(temp0, temp0, temp1);
    fp12e_set(retval[2], temp0);
    

    b1.mult(temp0, input[0]);
    d1.mult(temp0, temp0);
    b2.mult(temp0, temp0);
    d2.mult(temp0, temp0);
    
    b1.mult(temp1, input[1]);
    d1.mult(temp1, temp1);
    a2.mult(temp1, temp1);
    d2.mult(temp1, temp1);
    fp12e_invert(temp1, temp1);

    a1.mult(temp2, input[2]);
    d1.mult(temp2, temp2);
    b2.mult(temp2, temp2);
    d2.mult(temp2, temp2);
    fp12e_invert(temp2, temp2);
    
    a1.mult(temp3, input[3]);
    d1.mult(temp3, temp3);
    a2.mult(temp3, temp3);
    d2.mult(temp3, temp3);

    fp12e_mul(temp0, temp0, temp1);
    fp12e_mul(temp1, temp2, temp3);
    fp12e_mul(temp0, temp0, temp1);
    fp12e_set(retval[3], temp0);
    

    return retval;
}