/*
 * Copyright (C) 2015 Craig Shelley (craig@microtron.org.uk)
 * Copyright (C) 2016 Furrtek
 *
 * BCH Encoder/Decoder - Adapted from GNURadio
 *
 * This file is part of PortaPack.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#include <math.h>
#include <stdlib.h>
#include <ch.h>

#include "bch_code.hpp"

void BCHCode::generate_gf() {
    /*
     * generate GF(2**m) from the irreducible polynomial p(X) in p[0]..p[m]
     * lookup tables:  index->polynomial form   alpha_to[] contains j=alpha**i;
     * polynomial form -> index form  index_of[j=alpha**i] = i alpha=2 is the
     * primitive element of GF(2**m)
     */

    int i, mask;
    mask = 1;
    alpha_to[m] = 0;

    for (i = 0; i < m; i++) {
        alpha_to[i] = mask;

        index_of[alpha_to[i]] = i;

        if (p[i] != 0)
            alpha_to[m] ^= mask;

        mask <<= 1;
    }

    index_of[alpha_to[m]] = m;

    mask >>= 1;

    for (i = m + 1; i < n; i++) {
        if (alpha_to[i - 1] >= mask)
            alpha_to[i] = alpha_to[m] ^ ((alpha_to[i - 1] ^ mask) << 1);
        else
            alpha_to[i] = alpha_to[i - 1] << 1;

        index_of[alpha_to[i]] = i;
    }

    index_of[0] = -1;
}

void BCHCode::gen_poly() {
    /*
     * Compute generator polynomial of BCH code of length = 31, redundancy = 10
     * (OK, this is not very efficient, but we only do it once, right? :)
     */

    int ii, jj, ll, kaux;
    int test, aux, nocycles, root, noterms, rdncy;
    int cycle[15][6], size[15], min[11], zeros[11];

    // Generate cycle sets modulo 31
    cycle[0][0] = 0;
    size[0] = 1;
    cycle[1][0] = 1;
    size[1] = 1;
    jj = 1;  // cycle set index

    do {
        // Generate the jj-th cycle set
        ii = 0;
        do {
            ii++;
            cycle[jj][ii] = (cycle[jj][ii - 1] * 2) % n;
            size[jj]++;
            aux = (cycle[jj][ii] * 2) % n;

        } while (aux != cycle[jj][0]);

        // Next cycle set representative
        ll = 0;
        do {
            ll++;
            test = 0;
            for (ii = 1; ((ii <= jj) && (!test)); ii++)
                // Examine previous cycle sets
                for (kaux = 0; ((kaux < size[ii]) && (!test)); kaux++)
                    if (ll == cycle[ii][kaux])
                        test = 1;
        } while ((test) && (ll < (n - 1)));

        if (!(test)) {
            jj++;  // next cycle set index
            cycle[jj][0] = ll;
            size[jj] = 1;
        }

    } while (ll < (n - 1));

    nocycles = jj;  // number of cycle sets modulo n
    // Search for roots 1, 2, ..., d-1 in cycle sets

    kaux = 0;
    rdncy = 0;

    for (ii = 1; ii <= nocycles; ii++) {
        min[kaux] = 0;

        for (jj = 0; jj < size[ii]; jj++)
            for (root = 1; root < d; root++)
                if (root == cycle[ii][jj])
                    min[kaux] = ii;

        if (min[kaux]) {
            rdncy += size[min[kaux]];
            kaux++;
        }
    }

    noterms = kaux;
    kaux = 1;

    for (ii = 0; ii < noterms; ii++)
        for (jj = 0; jj < size[min[ii]]; jj++) {
            zeros[kaux] = cycle[min[ii]][jj];
            kaux++;
        }

    // Compute generator polynomial
    g[0] = alpha_to[zeros[1]];
    g[1] = 1;  // g(x) = (X + zeros[1]) initially

    for (ii = 2; ii <= rdncy; ii++) {
        g[ii] = 1;
        for (jj = ii - 1; jj > 0; jj--)
            if (g[jj] != 0)
                g[jj] = g[jj - 1] ^ alpha_to[(index_of[g[jj]] + zeros[ii]) % n];
            else
                g[jj] = g[jj - 1];

        g[0] = alpha_to[(index_of[g[0]] + zeros[ii]) % n];
    }
}

int* BCHCode::encode(int data[]) {
    // Calculate redundant bits bb[]

    int h, i, j = 0, start = 0, end = (n - k);  // 10
    int Mr[31];

    if (!valid) return nullptr;

    for (i = 0; i < n; i++) {
        Mr[i] = 0;
    }

    for (h = 0; h < k; ++h)
        Mr[h] = data[h];

    while (end < n) {
        for (i = end; i > start - 2; --i) {
            if (Mr[start] != 0) {
                Mr[i] ^= g[j];
                ++j;
            } else {
                ++start;
                j = 0;
                end = start + (n - k);
                break;
            }
        }
    }

    j = 0;
    for (i = start; i < end; ++i) {
        bb[j] = Mr[i];
        ++j;
    }

    return bb;
};

int BCHCode::decode(int recd[]) {
    // We do not need the Berlekamp algorithm to decode.
    // We solve before hand two equations in two variables.

    int i, j, q;
    int elp[3], s[5], s3;
    int count = 0, syn_error = 0;
    int loc[3], reg[3];
    int aux;
    int retval = 0;

    if (!valid) return 0;

    for (i = 1; i <= 4; i++) {
        s[i] = 0;
        for (j = 0; j < n; j++) {
            if (recd[j] != 0) {
                s[i] ^= alpha_to[(i * j) % n];
            }
        }
        if (s[i] != 0) {
            syn_error = 1;  // set flag if non-zero syndrome
        }
        /* NOTE: If only error detection is needed,
         * then exit the program here...
         */
        // Convert syndrome from polynomial form to index form
        s[i] = index_of[s[i]];
    };

    if (syn_error) {  // If there are errors, try to correct them
        if (s[1] != -1) {
            s3 = (s[1] * 3) % n;
            if (s[3] == s3) {  // Was it a single error ?
                // printf("One error at %d\n", s[1]);
                recd[s[1]] ^= 1;  // Yes: Correct it
            } else {
                /* Assume two errors occurred and solve
                 * for the coefficients of sigma(x), the
                 * error locator polynomial
                 */
                if (s[3] != -1) {
                    aux = alpha_to[s3] ^ alpha_to[s[3]];
                } else {
                    aux = alpha_to[s3];
                }
                elp[0] = 0;
                elp[1] = (s[2] - index_of[aux] + n) % n;
                elp[2] = (s[1] - index_of[aux] + n) % n;
                // printf("sigma(x) = ");
                // for (i = 0; i <= 2; i++) {
                //	printf("%3d ", elp[i]);
                // }
                // printf("\n");
                // printf("Roots: ");

                // Find roots of the error location polynomial
                for (i = 1; i <= 2; i++) {
                    reg[i] = elp[i];
                }
                count = 0;
                for (i = 1; i <= n; i++) {  // Chien search
                    q = 1;
                    for (j = 1; j <= 2; j++) {
                        if (reg[j] != -1) {
                            reg[j] = (reg[j] + j) % n;
                            q ^= alpha_to[reg[j]];
                        }
                    }
                    if (!q) {  // store error location number indices
                        loc[count] = i % n;
                        count++;
                    }
                }

                if (count == 2) {
                    // no. roots = degree of elp hence 2 errors
                    for (i = 0; i < 2; i++)
                        recd[loc[i]] ^= 1;
                } else {  // Cannot solve: Error detection
                    retval = 1;
                }
            }
        } else if (s[2] != -1) {  // Error detection
            retval = 1;
        }
    }

    return retval;
}

/*
 * Example usage BCH(31,21,5)
 *
 * p[] = coefficients of primitive polynomial used to generate GF(2**5)
 * m = order of the field GF(2**5) = 5
 * n = 2**5 - 1 = 31
 * t = 2 = error correcting capability
 * d = 2*t + 1 = 5 = designed minimum distance
 * k = n - deg(g(x)) = 21 = dimension
 * g[] = coefficients of generator polynomial, g(x) [n - k + 1]=[11]
 * alpha_to [] = log table of GF(2**5)
 * index_of[] = antilog table of GF(2**5)
 * data[] = coefficients of data polynomial, i(x)
 * bb[] = coefficients of redundancy polynomial ( x**(10) i(x) ) modulo g(x)
 */
BCHCode::BCHCode(
    std::vector<int> p_init,
    int m,
    int n,
    int k,
    int t)
    : m{m},
      n{n},
      k{k},
      t{t} {
    size_t i;

    d = 5;

    alpha_to = (int*)chHeapAlloc(NULL, sizeof(int) * (n + 1));
    index_of = (int*)chHeapAlloc(0, sizeof(int) * (n + 1));
    p = (int*)chHeapAlloc(0, sizeof(int) * (m + 1));
    g = (int*)chHeapAlloc(0, sizeof(int) * (n - k + 1));
    bb = (int*)chHeapAlloc(0, sizeof(int) * (n - k + 1));

    if (alpha_to == NULL ||
        index_of == NULL ||
        p == NULL ||
        g == NULL ||
        bb == NULL)
        valid = false;
    else
        valid = true;

    if (valid) {
        for (i = 0; i < (size_t)(m + 1); i++) {
            p[i] = p_init[i];
        }

        generate_gf(); /* generate the Galois Field GF(2**m) */
        gen_poly();    /* Compute the generator polynomial of BCH code */
    }
}

BCHCode::~BCHCode() {
    if (alpha_to != NULL) chHeapFree(alpha_to);
    if (index_of != NULL) chHeapFree(index_of);
    if (p != NULL) chHeapFree(p);
    if (g != NULL) chHeapFree(g);
    if (bb != NULL) chHeapFree(bb);
}