#include <stdlib.h>
#include "mpz_multi_powm_wu.h"

static size_t choose_window(size_t bit_count);
static void init_powers(mpz_t ***powers, mpz_t *bases, size_t base_count, mpz_t modulus, size_t power_count);
static void clear_powers(mpz_t **powers, size_t base_count, size_t power_count);
static void init_factors(mpz_t **factors, size_t count);
static void clear_factors(mpz_t *factors, size_t count);
static void calculate_factors(mpz_t *exponents, size_t exponent_count, mpz_t modulus, size_t bit_count, size_t window, mpz_t **powers, mpz_t *factors);
static void multiply_factors(mpz_t value, mpz_t *factors, size_t count, mpz_t modulus);

void mpz_multi_powm_wu(mpz_t value, mpz_t *bases, mpz_t *exponents, size_t count, mpz_t modulus, size_t bit_count)
{
    size_t window = choose_window(bit_count);
    size_t power_count = (size_t)(1 << (window - 1));
    mpz_t **powers;
    init_powers(&powers, bases, count, modulus, power_count);
    mpz_t *factors;
    init_factors(&factors, bit_count);

    calculate_factors(exponents, count, modulus, bit_count, window, powers, factors);
    multiply_factors(value, factors, bit_count, modulus);

    clear_powers(powers, count, power_count);
    clear_factors(factors, bit_count);
}

/*
Chooses appropriate window for a given number of bits as described in
http://dasan.sejong.ac.kr/~chlim/pub/multi_exp.ps.

bit_count - the number of bits.
*/
static size_t choose_window(size_t bit_count)
{
    size_t window;
    if (bit_count <= 24)
    {
        window = 2;
    }
    else if (bit_count <= 80)
    {
        window = 3;
    }
    else if (bit_count <= 240)
    {
        window = 4;
    }
    else if (bit_count <= 672)
    {
        window = 5;
    }
    else
    {
        window = 6;
    }
    return window;
}

/*
Initializes odd powers of given bases. All exponentiations are perfomed
modulo some numeber.

powers - the initialized powers
bases - the bases
base_count - the number of bases
modulus - the modulus
power_count - the number of odd powers for each base.
*/
static void init_powers(mpz_t ***powers, mpz_t *bases, size_t base_count, mpz_t modulus, size_t power_count)
{
    size_t size = mpz_sizeinbase(modulus, 2);
    mpz_t square;
    mpz_init2(square, size);
    mpz_t product;
    mpz_init2(product, size * 2);

    *powers = calloc(base_count, sizeof(mpz_t*));
    int i;
    for (i = 0; i < base_count; i++)
    {
        (*powers)[i] = calloc(power_count, sizeof(mpz_t));
        if (power_count > 0)
        {
            mpz_init2((*powers)[i][0], size);
            mpz_mod((*powers)[i][0], bases[i], modulus);
            mpz_powm_ui(square, (*powers)[i][0], 2, modulus);
            int j;
            for (j = 1; j < power_count; j++)
            {
                mpz_init2((*powers)[i][j], size);
                mpz_mul(product, (*powers)[i][j - 1], square);
                mpz_mod((*powers)[i][j], product, modulus);
            }
        }
    }

    mpz_clear(square);
    mpz_clear(product);
}

/*
Clears a given array of powers.

powers - the powers
base_count - the number of bases
powers_count - the number od powers for each base.
*/
static void clear_powers(mpz_t **powers, size_t base_count, size_t power_count)
{
    int i;
    for (i = 0; i < base_count; i++)
    {
        int j;
        for (j = 0; j < power_count; j++)
        {
            mpz_clear(powers[i][j]);
        }
        free(powers[i]);
    }
    free(powers);
}

/*
Initializes a given array of factors with 1.

factors - the factors
count - the number of factors.
*/
static void init_factors(mpz_t **factors, size_t count)
{
    *factors = calloc(count, sizeof(mpz_t));
    int i;
    for (i = 0; i < count; i++)
    {
        mpz_init_set_si((*factors)[i], 1);
    }
}

/*
Clears a given array of factors.

factors - the factors
count - the number of factors.
*/
static void clear_factors(mpz_t *factors, size_t count)
{
    int i;
    for (i = 0; i < count; i++)
    {
        mpz_clear(factors[i]);
    }
    free(factors);
}

/*
Processes exponents and calculates the factors of the future multi exponent.

exponents - the exponents
exponent_count - the number of exponents
modulus - the modulus
bit_count - the number of bits in the longest exponent
window - the window
powers - the precalculated odd powers of the bases
factors - the factors array to be updated.
*/
static void calculate_factors(mpz_t *exponents, size_t exponent_count, mpz_t modulus, size_t bit_count, size_t window, mpz_t **powers, mpz_t *factors)
{
    mpz_t exponent;
    mpz_init2(exponent, bit_count);
    mpz_t product;
    mpz_init2(product, bit_count * 2);
    int mask = (1 << window) - 1;

    int i;
    for (i = 0; i < exponent_count; i++)
    {
        mpz_set(exponent, exponents[i]);
        int index = 0;
        while (mpz_cmp_si(exponent, 0) > 0)
        {
            mp_bitcnt_t shift = mpz_scan1(exponent, 0);
            if (shift > 0)
            {
                mpz_tdiv_q_2exp(exponent, exponent, shift);
                index += (int)shift;
            }

            int j = ((int)mpz_get_si(exponent) & mask) >> 1;
            mpz_mul(product, factors[index], powers[i][j]);
            mpz_mod(factors[index], product, modulus);

            mpz_tdiv_q_2exp(exponent, exponent, window);
            index += (int)window;
        }
    }

    mpz_clear(exponent);
    mpz_clear(product);
}

/*
Multiplies the factors modulo some number.

value - the returned result
factors - the factors
count - the number of factors
modulus - the modulus.
*/
static void multiply_factors(mpz_t value, mpz_t *factors, size_t count, mpz_t modulus)
{
    mpz_t product;
    mpz_init2(product, mpz_sizeinbase(modulus, 2) * 2);

    mpz_set_si(value, 1);

    int i;
    for (i = (int)count - 1; i >= 0; i--)
    {
        mpz_powm_ui(value, value, 2, modulus);
        if (mpz_cmp_si(factors[i], 1) > 0)
        {
            mpz_mul(product, value, factors[i]);
            mpz_mod(value, product, modulus);
        }
    }

    mpz_clear(product);
}
