//
// LDPC decoder for FT8.
//
// given a 174-bit codeword as an array of log-likelihood of zero,
// return a 174-bit corrected codeword, or zero-length array.
// last 87 bits are the (systematic) plain-text.
// this is an implementation of the sum-product algorithm
// from Sarah Johnson's Iterative Error Correction book.
// codeword[i] = log ( P(x=0) / P(x=1) )
//

#include "ldpc.h"
#include "constants.h"

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <stdbool.h>

static int ldpc_check(uint8_t codeword[]);
static float fast_tanh(float x);
static float fast_atanh(float x);

// Packs a string of bits each represented as a zero/non-zero byte in plain[],
// as a string of packed bits starting from the MSB of the first byte of packed[]
void pack_bits(const uint8_t plain[], int num_bits, uint8_t packed[])
{
    int num_bytes = (num_bits + 7) / 8;
    for (int i = 0; i < num_bytes; ++i)
    {
        packed[i] = 0;
    }

    uint8_t mask = 0x80;
    int byte_idx = 0;
    for (int i = 0; i < num_bits; ++i)
    {
        if (plain[i])
        {
            packed[byte_idx] |= mask;
        }
        mask >>= 1;
        if (!mask)
        {
            mask = 0x80;
            ++byte_idx;
        }
    }
}

// codeword is 174 log-likelihoods.
// plain is a return value, 174 ints, to be 0 or 1.
// max_iters is how hard to try.
// ok == 87 means success.
void ldpc_decode(float codeword[], int max_iters, uint8_t plain[], int *ok)
{
    float m[FT8_M][FT8_N]; // ~60 kB
    float e[FT8_M][FT8_N]; // ~60 kB
    int min_errors = FT8_M;

    for (int j = 0; j < FT8_M; j++)
    {
        for (int i = 0; i < FT8_N; i++)
        {
            m[j][i] = codeword[i];
            e[j][i] = 0.0f;
        }
    }

    for (int iter = 0; iter < max_iters; iter++)
    {
        for (int j = 0; j < FT8_M; j++)
        {
            for (int ii1 = 0; ii1 < kFT8_LDPC_num_rows[j]; ii1++)
            {
                int i1 = kFT8_LDPC_Nm[j][ii1] - 1;
                float a = 1.0f;
                for (int ii2 = 0; ii2 < kFT8_LDPC_num_rows[j]; ii2++)
                {
                    int i2 = kFT8_LDPC_Nm[j][ii2] - 1;
                    if (i2 != i1)
                    {
                        a *= fast_tanh(-m[j][i2] / 2.0f);
                    }
                }
                e[j][i1] = logf((1 - a) / (1 + a));
            }
        }

        for (int i = 0; i < FT8_N; i++)
        {
            float l = codeword[i];
            for (int j = 0; j < 3; j++)
                l += e[kFT8_LDPC_Mn[i][j] - 1][i];
            plain[i] = (l > 0) ? 1 : 0;
        }

        int errors = ldpc_check(plain);

        if (errors < min_errors)
        {
            // Update the current best result
            min_errors = errors;

            if (errors == 0)
            {
                break; // Found a perfect answer
            }
        }

        for (int i = 0; i < FT8_N; i++)
        {
            for (int ji1 = 0; ji1 < 3; ji1++)
            {
                int j1 = kFT8_LDPC_Mn[i][ji1] - 1;
                float l = codeword[i];
                for (int ji2 = 0; ji2 < 3; ji2++)
                {
                    if (ji1 != ji2)
                    {
                        int j2 = kFT8_LDPC_Mn[i][ji2] - 1;
                        l += e[j2][i];
                    }
                }
                m[j1][i] = l;
            }
        }
    }

    *ok = min_errors;
}

//
// does a 174-bit codeword pass the FT8's LDPC parity checks?
// returns the number of parity errors.
// 0 means total success.
//
static int ldpc_check(uint8_t codeword[])
{
    int errors = 0;

    for (int j = 0; j < FT8_M; ++j)
    {
        uint8_t x = 0;
        for (int i = 0; i < kFT8_LDPC_num_rows[j]; ++i)
        {
            x ^= codeword[kFT8_LDPC_Nm[j][i] - 1];
        }
        if (x != 0)
        {
            ++errors;
        }
    }
    return errors;
}

void bp_decode(float codeword[], int max_iters, uint8_t plain[], int *ok)
{
    float tov[FT8_N][3];
    float toc[FT8_M][7];

    int min_errors = FT8_M;

    int nclast = 0;
    int ncnt = 0;

    // initialize message data
    for (int i = 0; i < FT8_N; ++i)
    {
        for (int j = 0; j < 3; ++j)
        {
            tov[i][j] = 0;
        }
    }

    for (int iter = 0; iter < max_iters; ++iter)
    {
        float zn[FT8_N];
        int plain_sum = 0;

        // Update bit log likelihood ratios (tov=0 in iter 0)
        for (int i = 0; i < FT8_N; ++i)
        {
            zn[i] = codeword[i] + tov[i][0] + tov[i][1] + tov[i][2];
            plain[i] = (zn[i] > 0) ? 1 : 0;
            plain_sum += plain[i];
        }

        if (plain_sum == 0)
        {
            // message converged to all-zeros, which is prohibited
            break;
        }

        // Check to see if we have a codeword (check before we do any iter)
        int errors = ldpc_check(plain);

        if (errors < min_errors)
        {
            // we have a better guess - update the result
            min_errors = errors;

            if (errors == 0)
            {
                break; // Found a perfect answer
            }
        }

        // Send messages from bits to check nodes
        for (int i = 0; i < FT8_M; ++i)
        {
            for (int j = 0; j < kFT8_LDPC_num_rows[i]; ++j)
            {
                int ibj = kFT8_LDPC_Nm[i][j] - 1;
                float Tnm = zn[ibj];
                for (int kk = 0; kk < 3; ++kk)
                {
                    // subtract off what the bit had received from the check
                    if ((kFT8_LDPC_Mn[ibj][kk] - 1) == i)
                    {
                        Tnm -= tov[ibj][kk];
                        break;
                    }
                }
                toc[i][j] = fast_tanh(-Tnm / 2);
            }
        }

        // send messages from check nodes to variable nodes
        for (int i = 0; i < FT8_N; ++i)
        {
            for (int j = 0; j < 3; ++j)
            {
                int ichk = kFT8_LDPC_Mn[i][j] - 1;
                float Tmn = 1.0f;
                for (int k = 0; k < kFT8_LDPC_num_rows[ichk]; ++k)
                {
                    if ((kFT8_LDPC_Nm[ichk][k] - 1) != i)
                    {
                        Tmn *= toc[ichk][k];
                    }
                }
                tov[i][j] = 2 * fast_atanh(-Tmn);
            }
        }
    }

    *ok = min_errors;
}

// https://varietyofsound.wordpress.com/2011/02/14/efficient-tanh-computation-using-lamberts-continued-fraction/
// http://functions.wolfram.com/ElementaryFunctions/ArcTanh/10/0001/
// https://mathr.co.uk/blog/2017-09-06_approximating_hyperbolic_tangent.html

// thank you Douglas Bagnall
// https://math.stackexchange.com/a/446411
static float fast_tanh(float x)
{
    if (x < -4.97f)
    {
        return -1.0f;
    }
    if (x > 4.97f)
    {
        return 1.0f;
    }
    float x2 = x * x;
    //float a = x * (135135.0f + x2 * (17325.0f + x2 * (378.0f + x2)));
    //float b = 135135.0f + x2 * (62370.0f + x2 * (3150.0f + x2 * 28.0f));
    //float a = x * (10395.0f + x2 * (1260.0f + x2 * 21.0f));
    //float b = 10395.0f + x2 * (4725.0f + x2 * (210.0f + x2));
    float a = x * (945.0f + x2 * (105.0f + x2));
    float b = 945.0f + x2 * (420.0f + x2 * 15.0f);
    return a / b;
}

static float fast_atanh(float x)
{
    float x2 = x * x;
    //float a = x * (-15015.0f + x2 * (19250.0f + x2 * (-5943.0f + x2 * 256.0f)));
    //float b = (-15015.0f + x2 * (24255.0f + x2 * (-11025.0f + x2 * 1225.0f)));
    //float a = x * (-1155.0f + x2 * (1190.0f + x2 * -231.0f));
    //float b = (-1155.0f + x2 * (1575.0f + x2 * (-525.0f + x2 * 25.0f)));
    float a = x * (945.0f + x2 * (-735.0f + x2 * 64.0f));
    float b = (945.0f + x2 * (-1050.0f + x2 * 225.0f));
    return a / b;
}