#include "general_library.hpp"

// https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Modular_integers
uint32_t multiplicative_inverse(int a, int n) {
    int t = 0, newt = 1, r = n, newr = a;
    while (newr != 0) {
        int quotient = r/newr;

        int tmp = newt;
        newt = t-quotient*newt;
        t = tmp;

        tmp = newr;
        newr = r-quotient*newr;
        r = tmp;
    }
    if (r > 1)
        throw std::runtime_error("Tried to inverse a noninvertible element!");
    if (t < 0)
        t = t + n;

    return t;
}

// https://stackoverflow.com/questions/18620942/find-the-smallest-period-of-input-string-in-on
std::vector<int> calculateLPS(const char * pat, int m) {
    /* int[] lps = new int[pat.length()]; */
    int len = 0;
    int i = 1;
    std::vector<int> lps = {0};
    lps.resize(m);

    while (i < m) {
        if (pat[i] == pat[len]) {
            len++;
            lps[i] = len;
            i++;
        }
        else {
            if (len != 0) {
                len = lps[len - 1];
            }
            else {
                lps[i] = len;
                i++;
            }
        }
    }
    return lps;
}

// calculates the length of the shortest period
int len_of_shortest_period (const char * pattern, int m) {
    std::vector<int> lps = calculateLPS(pattern, m);
    //start at the end of the string
    int i = lps.size()-1;
    while (lps[i] != 0) {
        //shift back
        i -= lps[i];
    }
    return i+1;
}

// SECTION: Functions related to calculating with polynomials in Z2[x]

uint32_t irreducible_polynomials[] {
    0b11,                              // irreducible polynomial of degree 1
    0b111,                             // irreducible polynomial of degree 2
    0b1011,                            // irreducible polynomial of degree 3
    0b10011,                           // irreducible polynomial of degree 4
    0b100101,                          // irreducible polynomial of degree 5
    0b1011011,                         // irreducible polynomial of degree 6
    0b10000011,                        // irreducible polynomial of degree 7
    0b100011101,                       // irreducible polynomial of degree 8
    0b1000010001,                      // irreducible polynomial of degree 9
    0b10001101111,                     // irreducible polynomial of degree 10
    0b100000000101,                    // irreducible polynomial of degree 11
    0b1000011101011,                   // irreducible polynomial of degree 12
    0b10000000011011,                  // irreducible polynomial of degree 13
    0b100000010101001,                 // irreducible polynomial of degree 14
    0b1000000000110101,                // irreducible polynomial of degree 15
    0b10000000000101101,               // irreducible polynomial of degree 16
    0b100000000000001001,              // irreducible polynomial of degree 17
    0b1000001010000000011,             // irreducible polynomial of degree 18
    0b10000000000000100111,            // irreducible polynomial of degree 19
    0b100000000011011110011,           // irreducible polynomial of degree 20
    0b1000000000000001100101,          // irreducible polynomial of degree 21
    0b10000000001111101100001,         // irreducible polynomial of degree 22
    0b100000000000000000100001,        // irreducible polynomial of degree 23
    0b1000000011110011010101001,       // irreducible polynomial of degree 24
    0b10000000000000000101000101,      // irreducible polynomial of degree 25
    0b100000000000100010111010011,     // irreducible polynomial of degree 26
    0b1000000000000001011010101101,    // irreducible polynomial of degree 27
    0b10000000000000010000011100101,   // irreducible polynomial of degree 28
    0b100000000000000000000000000101,  // irreducible polynomial of degree 29
    0b1000000000000110010100010101111, // irreducible polynomial of degree 30
    0b10000000000000000000000000001001 // irreducible polynomial of degree 31
};

size_t degree (uint32_t polynomial) {
    return 31 - __builtin_clz(polynomial);
}

uint32_t multiply_modulo_polynomials_in_Z2 (uint32_t q1, uint32_t q2, uint32_t p) {
    size_t p_degree = degree(p);
    size_t q1_degree = degree(q1);
    size_t q2_degree = degree(q2);

    if (q1_degree > p_degree)
        throw std::logic_error("We only support multiply-modulo whenever the multiplied polynomials are initially at most the same degree as the modulo polynomial.");

    if (q1_degree == p_degree)
        q1 ^= p;

    uint32_t arr[q2_degree+1];
    arr[0] = q1;
    for (size_t i = 1; i <= q2_degree; i++) {
        arr[i] = arr[i-1]<<1;
        if (degree(arr[i]) == p_degree)
            arr[i] ^= p;
    }

    uint32_t sum = 0;

    std::bitset<32> b(q2);
    for (size_t i = 0; i <= q2_degree; i++) {
        if (b[i] == 1)
            sum ^= arr[i];
    }

    return sum;
}

uint32_t polynomial_power_in_Z2(uint32_t q, uint32_t exp, uint32_t p) {
    if (exp == 0)
        return 1;

    if (degree(q) == degree(p))
        q ^= p;

    size_t log2exp = degree(exp) + 1;

    uint32_t arr[log2exp];
    arr[0] = q;
    for (size_t i = 1; i < log2exp; i++)
        arr[i] = multiply_modulo_polynomials_in_Z2(arr[i-1], arr[i-1], p);

    uint32_t res = 1;
    std::bitset<32> b(exp);
    for (size_t i = 0; i <= log2exp; i++) {
        if (b[i] == 1)
            res = multiply_modulo_polynomials_in_Z2(res, arr[i], p);
    }

    return res;
}

void rref_Z2(std::vector<std::bitset<32>> &A, size_t m, size_t n) {
    if (m > 31 || n > 32)
        throw std::logic_error("We only support rref_Z2 for up to a (31x32)-matrix.");

    // calculate row echelon form
    // https://en.wikipedia.org/wiki/Gaussian_elimination#Pseudocode
    size_t h = 0; // Initialization of the pivot row
    size_t k = 0; // Initialization of the pivot column

    while (h < m && k < n) {
        // Find the k-th pivot:
        int i_max = -1;
        for (size_t i = h; i < m; i++) {
            if (A[i][k]) {
                i_max = i;
                break;
            }
        }

        if (i_max == -1) {
            // No pivot in this column, pass to next column
            k++;
        } else {
            auto tmp = A[i_max];
            A[i_max] = A[h];
            A[h] = tmp;

            // Do for all rows below pivot:
            for (size_t i = h+1; i < m; i++){
                bool f = A[i][k];
                if (f == 0)
                    continue;
                // Fill with zeros the lower part of pivot column:
                A[i][k] = 0;
                // Do for all remaining elements in current row:
                for (size_t j = k+1; j<n; j++) {
                    A[i][j] = A[i][j] ^ A[h][j];
                }
            }
            // Increase pivot row and column
            h++;
            k++;
        }
    }

    // Perform back substitution
    for (size_t i = 1; i < m; i++) {
        for (size_t j = 0; j < i; j++) {
            if (A[j][i] == 1) {
                for (size_t l = i; l < n; l++) // check this
                    A[j][l] = A[j][l] ^ A[i][l];
            }
        }
    }
}

uint32_t get_random_irreducible_polynomial_in_Z2 (size_t k) {
    if (k > 31)
        throw std::logic_error("We only support polynomials of degree 31 or less.");

    uint32_t init_p = irreducible_polynomials[k-1];

    srand(time(0));
    uint32_t gamma = (rand() % (((uint32_t)1 << (k+1)) - 2)) + 2; // gamma in GF(2^k)\GF(2)

    uint32_t gamma_pow[k+1];
    for (size_t i = 0; i <= k; i++)
        gamma_pow[i] = polynomial_power_in_Z2(gamma, i, init_p);

    // Create an array of the polynomials in gamma_pow, and transpose it at the same time
    std::vector<std::bitset<32>> A(k); // NOTE: That 32 > 31 >= k
    for (size_t i = 0; i < k; i++) {
        for (size_t j = 0; j <= k; j++) {
            A[i][k-j] = gamma_pow[j] & ((uint32_t)1<<i);
        }
    }

    rref_Z2(A, k, k+1);

    std::bitset<32> bp;
    for (size_t i = 0; i < k; i++) {
        bp[i] = A[i][k];
    }
    bp[k] = 1;
    uint32_t p = (uint32_t)bp.to_ulong();

    return p;
}