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

static size_t choose_window(size_t bit_count);
static size_t get_window_count(size_t base_count, size_t window);
static void init_factors(mpz_t ***factors, mpz_t *bases, size_t base_count, mpz_t modulus, size_t window, size_t window_count);
static void clear_factors(mpz_t **factors, size_t base_count, size_t window, size_t window_count);
static void calculate_multi_powm(mpz_t value, mpz_t *exponents, size_t exponent_count, mpz_t modulus, size_t bit_count, size_t window, size_t window_count, mpz_t **factors);

void mpz_multi_powm_ll(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 window_count = get_window_count(count, window);
    mpz_t **factors;
    init_factors(&factors, bases, count, modulus, window, window_count);

    calculate_multi_powm(value, exponents, count, modulus, bit_count, window, window_count, factors);

    clear_factors(factors, count, window, window_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 <= 10)
    {
	window = 2;
    }
    else if (bit_count <= 24)
    {
	window = 3;
    }
    else if (bit_count <= 60)
    {
        window = 4;
    }
    else if (bit_count <= 144)
    {
	window = 5;
    }
    else if (bit_count <= 342)
    {
	window = 6;
    }
    else if (bit_count <= 797)
    {
        window = 7;
    }
    else if (bit_count <= 1828)
    {
	window = 8;
    }
    else
    {
	window = 9;
    }
    return window;
}

/*
Returns the number of windows for given number of bases and  window.

base_count - the number of bases
window - the window.
*/
static size_t get_window_count(size_t base_count, size_t window)
{
    return (base_count + window - 1) / window;
}

/*
Initializes factors from given bases, modulus, window and a number of windows.

factors - the factors to be initialized
bases - the bases
base_count - the number of bases
modulus - the modulus
window - the window
window_count - the number of windows.
*/
static void init_factors(mpz_t ***factors, mpz_t *bases, size_t base_count, mpz_t modulus, size_t window, size_t window_count)
{
    mpz_t product;
    mpz_init2(product, mpz_sizeinbase(modulus, 2) * 2);

    *factors = calloc(window_count, sizeof(mpz_t*));
    int i;
    for (i = 0; i < window_count; i++)
    {
        size_t bases_in_window = window;
        if (i == window_count - 1)
        {
            bases_in_window = base_count - (window_count - 1) * window;
        }
        size_t factor_count = 1 << bases_in_window;
        (*factors)[i] = calloc(factor_count, sizeof(mpz_t));
        int offset = i * (int)window;
        int j;
        for (j = 0; j < factor_count; j++)
        {
            mpz_set_si(product, 1);
            int k;
            for (k = 0; k < bases_in_window; k++)
            {
                 int mask = 1 << (bases_in_window - 1 - k);
                 if ((j & mask) > 0)
                 {
                     mpz_mul(product, product, bases[offset + k]);
                     mpz_mod(product, product, modulus);
                 }
            }
            mpz_init_set((*factors)[i][j], product);
        }
    }

    mpz_clear(product);
}

/*
Clears given factors.

factors - the factors
base_count - the naumber of bases
window - the window
window_count - the number of windows.
*/
static void clear_factors(mpz_t **factors, size_t base_count, size_t window, size_t window_count)
{
    int i;
    for (i = 0; i < window_count; i++)
    {
        size_t bases_in_window = window;
        if (i == window_count - 1)
        {
            bases_in_window = base_count - (window_count - 1) * window;
        }
        size_t factor_count = 1 << bases_in_window;
        int j;
        for (j = 0; j < factor_count; j++)
        {
            mpz_clear(factors[i][j]);
        }
        free(factors[i]);
    }
    free(factors);
}

/*
Calculates multi-exponentiation for given exponents, modulus, number of bits, window, nember of windows and factors.

value - the value
exponents - the exponents
exponent_count - the number of exponents
modulus the modulus
bit_count - the number of bits in the greatest exponent
window - the window
window_count - the number of windows
factors - the windowed exponents.
*/
static void calculate_multi_powm(mpz_t value, mpz_t *exponents, size_t exponent_count, mpz_t modulus, size_t bit_count, size_t window, size_t window_count, mpz_t **factors)
{
    mpz_t product;
    mpz_init2(product, mpz_sizeinbase(modulus, 2));

    mpz_set_si(value, 1);
    int i;
    for(i = (int)bit_count - 1; i >= 0; i--)
    {
        mpz_powm_ui(value, value, 2, modulus);

        int j;
        for (j = 0; j < window_count; j++)
        {
            int offset = j * (int)window;
            int exponents_in_window = (int)window;
            if (j == window_count - 1)
            {
                exponents_in_window = (int)exponent_count - offset;
            }
            int index = 0;
            int k;
            for (k = 0; k < exponents_in_window; k++)
            {
                index <<= 1;
                index |= mpz_tstbit(exponents[offset + k], (mp_bitcnt_t)i);
            }
            mpz_mul(product, value, factors[j][index]);
            mpz_mod(value, product, modulus);
        }
    }

    mpz_clear(product);
}
