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

static size_t choose_window(size_t count, size_t bit_count);
static size_t get_window_count(size_t bit_count, size_t window);
static void init_windowed_exponents(int ***windowed_exponents, mpz_t *exponents, size_t exponent_count, size_t bit_count, size_t window, size_t window_count);
static void calculate_multi_powm(mpz_t value, mpz_t *bases, size_t bases_count, mpz_t modulus, size_t window, size_t window_count, int **windowed_exponents);
static void clear_windowed_exponents(int **windowed_exponents, size_t count);

void mpz_multi_powm_bg(mpz_t value, mpz_t *bases, mpz_t *exponents, size_t count, mpz_t modulus, size_t bit_count)
{
    size_t window = choose_window(count, bit_count);
    size_t window_count = get_window_count(bit_count, window);
    int **windowed_exponents;
    init_windowed_exponents(&windowed_exponents, exponents, count, bit_count, window, window_count);

    calculate_multi_powm(value, bases, count, modulus, window, window_count, windowed_exponents);

    clear_windowed_exponents(windowed_exponents, count);
}

/*
Chooses thw window for given base count and bit count.

base_count - the number of bases
bit_count - the number of bits in the greatest exponent.
*/
static size_t choose_window(size_t base_count, size_t bit_count)
{
    size_t window;
    if (base_count <= 11)
    {
        window = 2;
    }
    else if (base_count <= 25)
    {
        window = 3;
    }
    else if (base_count <= 61)
    {
        window = 4;
    }
    else if (base_count <= 148)
    {
        window = 5;
    }
    else if (base_count <= 324)
    {
        window = 6;
    }
    else if (base_count <= 776)
    {
        window = 7;
    }
    else if (base_count <= 1892)
    {
        window = 8;
    }
    else if (base_count <= 3826)
    {
        window = 9;
    }
    else if (base_count <= 12269)
    {
        window = 10;
    }
    else if (base_count <= 23513)
    {
        window = 11;
    }
    else
    {
        window = 12;
    }
    return window;
}

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

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

/*
Initializes the windowed exponents from given exponents, window and number of windows.

windowed_exponents - the initialized windowed exponents
exponents - the exponents
exponent_count - the number of exponents
window - the windows
window_count - the number od windows.
*/
static void init_windowed_exponents(int ***windowed_exponents, mpz_t *exponents, size_t exponent_count, size_t bit_count, size_t window, size_t window_count)
{
    mpz_t exponent;
    mpz_init2(exponent, bit_count);

    *windowed_exponents = calloc(exponent_count, sizeof(int*));

    int mask = (1 << window) - 1;
    int i;
    for (i = 0; i < exponent_count; i++)
    {
        (*windowed_exponents)[i] = calloc(window_count, sizeof(int));
        mpz_set(exponent, exponents[i]);
        int j;
        for (j = 0; j < window_count; j++)
        {
            (*windowed_exponents)[i][j] = (int)mpz_get_si(exponent) & mask;
            mpz_tdiv_q_2exp(exponent, exponent, window);
        }
    }

    mpz_clear(exponent);
}

/*
Calculates multi-exponentiation for given bases, modulus, window, nember of windows and windowed exponents.

value - the value
bases - the bases
base_count - the number of bases
modulus the modulus
window - the window
window_count - the number of windows
windowed_exponents - the windowed exponents.
*/
static void calculate_multi_powm(mpz_t value, mpz_t *bases, size_t bases_count, mpz_t modulus, size_t window, size_t window_count, int **windowed_exponents)
{
    mpz_t product;
    mpz_init2(product, mpz_sizeinbase(modulus, 2));

    mpz_set_si(value, 1);
    unsigned long int exponent = 1UL << window;
    int max_mask = 1 << (window - 1);
    int j;
    for (j = (int)window_count - 1; j >= 0; j--)
    {
        mpz_powm_ui(value, value, exponent, modulus);

        mpz_set_si(product, 1);
        int mask;
        for (mask = max_mask; mask > 0; mask = mask >> 1)
        {
            mpz_powm_ui(product, product, 2, modulus);
            int i;
            for (i = 0; i < bases_count; i++)
            {
                if ((windowed_exponents[i][j] & mask) > 0)
                {
                    mpz_mul(product, product, bases[i]);
                    mpz_mod(product, product, modulus);
                }
            }
        }

        mpz_mul(product, value, product);
        mpz_mod(value, product, modulus);
    }

    mpz_clear(product);
}

/*
Clears windowed exponents

windowed_exponents - the exponents
count - the count.
*/
static void clear_windowed_exponents(int **windowed_exponents, size_t count)
{
    int i;
    for (i = 0; i < count; i++)
    {
         free(windowed_exponents[i]);
    }
    free(windowed_exponents);
}
