#include <algorithm>
#include <cstring>
#include <iostream>
#include <vector>

#define INTERLEAVER_BLOCK_SIZE 8


// Forward interleaver dimensions:
//  PPM   == number of bits per symbol OUT of interleaver        AND number of codewords IN to interleaver
//  RDD+4 == number of bits per codeword IN to interleaver       AND number of interleaved codewords OUT of interleaver
//
// bit width in:  (4+rdd)   block length: ppm
// bit width out: ppm       block length: (4+rdd)

// Reverse interleaver (de-interleaver) dimensions:
//  PPM   == number of bits per symbol IN to deinterleaver       AND number of codewords OUT of deinterleaver
//  RDD+4 == number of bits per codeword OUT of deinterleaver    AND number of interleaved codewords IN to deinterleaver
//
// bit width in:  ppm       block length: (4+rdd)
// bit width out: (4+rdd)   block length: ppm
void
deinterleave(std::vector <unsigned short> &symbols,
                              std::vector <unsigned char> &codewords,
                              unsigned char ppm,
                              unsigned char rdd)
    {
      int symbol_offset = 0;
      int bit_offset    = 0;
      int bit_idx       = 0;
      unsigned char block[INTERLEAVER_BLOCK_SIZE];    // maximum bit-width is 8, should RDD==4

      // Swap MSBs of each symbol within buffer (one of LoRa's quirks)
      for (int symbol_idx = 0; symbol_idx < symbols.size(); symbol_idx++)
      {
        symbols[symbol_idx] = ( (symbols[symbol_idx] &  (0x1 << (ppm-1))) >> 1 |
                                (symbols[symbol_idx] &  (0x1 << (ppm-2))) << 1 |
                                (symbols[symbol_idx] & ((0x1 << (ppm-2)) - 1))
                              );
      }

      // Block interleaver: de-interleave RDD+4 symbols at a time into PPM codewords
      for (int block_count = 0; block_count < symbols.size()/(4+rdd); block_count++)
      {
        memset(block, 0, INTERLEAVER_BLOCK_SIZE*sizeof(unsigned char));
        bit_idx = 0;
        bit_offset = 0;

        // Iterate through each bit in the interleaver block
        for (int bitcount = 0; bitcount < ppm*(4+rdd); bitcount++)
        {
              // Symbol indexing                                     // Diagonal pattern mask
          if (symbols[(bitcount % (4+rdd)) + (4+rdd)*block_count] & ((0x1 << (ppm-1)) >> ((bit_idx + bit_offset) % ppm)))
          {
            block[bitcount / (4+rdd)] |= 0x1 << (bitcount % (4+rdd));   // integer divison in C++ is defined to floor
          }

          // bit_idx walks through diagonal interleaving pattern, bit_offset adjusts offset starting point for each codeword
          if (bitcount % (4+rdd) == (4+rdd-1))
          {
            bit_idx = 0;
            bit_offset++;
          }
          else
          {
            bit_idx++;
          }
        }

        // Post-process de-interleaved codewords
        for (int cw_idx = 0; cw_idx < ppm; cw_idx++)
        {
          // Put bits into traditional Hamming order
          switch (rdd)
          {
            case 4:
              block[cw_idx] = (block[cw_idx] & 128) | (block[cw_idx] & 64) | (block[cw_idx] & 32) >> 5 | (block[cw_idx] & 16) | (block[cw_idx] & 8) << 2 | (block[cw_idx] & 4) << 1 | (block[cw_idx] & 2) << 1 | (block[cw_idx] & 1) << 1;
              break;

            case 3:
              block[cw_idx] = (block[cw_idx] & 64) | (block[cw_idx] & 32) | (block[cw_idx] & 16) >> 1 | (block[cw_idx] & 8) << 1 | (block[cw_idx] & 4) | (block[cw_idx] & 2) | (block[cw_idx] & 1);
              break;

            default:
              break;
          }

          // Mask
          block[cw_idx] = block[cw_idx] & ((1 << (4+rdd)) - 1);
        }

        // Append deinterleaved codewords to codeword buffer, rearranging into proper order
        if (ppm == 8)
        {
          codewords.push_back(block[6]);
          codewords.push_back(block[7]);
          codewords.push_back(block[4]);
          codewords.push_back(block[5]);
        }
        else if (ppm == 7)
        {
          codewords.push_back(block[6]);
          codewords.push_back(block[4]);
          codewords.push_back(block[5]);
        }
        else if (ppm == 6)
        {
          codewords.push_back(block[4]);
          codewords.push_back(block[5]);
        }
        else if (ppm == 5)
        {
          codewords.push_back(block[4]);
        }
        codewords.push_back(block[2]);
        codewords.push_back(block[3]);
        codewords.push_back(block[0]);
        codewords.push_back(block[1]);
      }
    }

// Forward interleaver dimensions:
//  PPM   == number of bits per symbol OUT of interleaver        AND number of codewords IN to interleaver
//  RDD+4 == number of bits per codeword IN to interleaver       AND number of interleaved codewords OUT of interleaver
//
// bit width in:  (4+rdd)   block length: ppm
// bit width out: ppm       block length: (4+rdd)
void
interleave(const std::vector <unsigned char> &codewords,
           std::vector <unsigned short> &symbols,
           unsigned char ppm,
           unsigned char rdd)
{
    int bit_offset     = 0;
    int bit_idx        = 0;
    unsigned char block[INTERLEAVER_BLOCK_SIZE];    // maximum bit-width is 8, should RDD==4
    unsigned char reordered[INTERLEAVER_BLOCK_SIZE];

    // Block interleaver: interleave PPM codewords at a time into 4+RDD codewords
    for (int block_count = 0; block_count < codewords.size()/ppm; block_count++)
    {

        memset(block, 0, INTERLEAVER_BLOCK_SIZE*sizeof(unsigned char));
        bit_idx = 0;
        bit_offset = 0;

        if (ppm == 6)
        {
            reordered[4] = codewords[0 + block_count*ppm];
            reordered[5] = codewords[1 + block_count*ppm];
            reordered[2] = codewords[2 + block_count*ppm];
            reordered[3] = codewords[3 + block_count*ppm];
            reordered[0] = codewords[4 + block_count*ppm];
            reordered[1] = codewords[5 + block_count*ppm];
        }
        else if (ppm == 8)
        {
            reordered[0] = codewords[0 + block_count*ppm];
            reordered[7] = codewords[1 + block_count*ppm];
            reordered[2] = codewords[2 + block_count*ppm];
            reordered[1] = codewords[3 + block_count*ppm];
            reordered[4] = codewords[4 + block_count*ppm];
            reordered[3] = codewords[5 + block_count*ppm];
            reordered[6] = codewords[6 + block_count*ppm];
            reordered[5] = codewords[7 + block_count*ppm];
        }

        // Iterate through each bit in the interleaver block
        for (int bitcount = 0; bitcount < ppm*(4+rdd); bitcount++)
        {

            if (reordered[(bitcount / (4+rdd)) /*+ ppm*block_count*/] & 0x1 << (bitcount % (4+rdd)))
            {
                block[bitcount % (4+rdd)] |= ((0x1 << (ppm-1)) >> ((bit_idx + bit_offset) % ppm));   // integer divison in C++ is defined to floor
            }

            // bit idx walks through diagonal interleaving pattern
            if (bitcount % (4+rdd) == (4+rdd-1))
            {
                bit_idx = 0;
                bit_offset++;
            }
            else
            {
                bit_idx++;
            }
        }

        for (int block_idx = 0; block_idx < (4+rdd); block_idx++)
        {
            symbols.push_back(block[block_idx]);
        }
    }

    // Swap MSBs of each symbol within buffer (one of LoRa's quirks)
    for (int symbol_idx = 0; symbol_idx < symbols.size(); symbol_idx++)
    {
        symbols[symbol_idx] = ( (symbols[symbol_idx] &  (0x1 << (ppm-1))) >> 1 |
                                (symbols[symbol_idx] &  (0x1 << (ppm-2))) << 1 |
                                (symbols[symbol_idx] & ((0x1 << (ppm-2)) - 1))
                                );
    }
}

void
hamming_encode(const std::vector<unsigned char> &nybbles,
               std::vector<unsigned char> &codewords,
               unsigned char rdd)
{
    unsigned char p1, p2, p4, p8;
    unsigned char mask;

    for (int i = 0; i < nybbles.size(); i++)
    {
//     p1 = parity((unsigned char)nybbles[i], mask = (unsigned char)HAMMING_P1_BITMASK);
//     p2 = parity((unsigned char)nybbles[i], mask = (unsigned char)HAMMING_P2_BITMASK);
//     p4 = parity((unsigned char)nybbles[i], mask = (unsigned char)HAMMING_P4_BITMASK);
//     p8 = parity((unsigned char)nybbles[i] | p1 << 7 | p2 << 6 | p4 << 4,
//         mask = (unsigned char)HAMMING_P8_BITMASK);
       p1 = p2 = p4 = p8 = 0;


    codewords.push_back((   (p1 << 7) |
                            (p2 << 6) |
                            (p8 << 5) |
                            (p4 << 4) |
                            (nybbles[i] & 0x08) |
                            (nybbles[i] & 0x04) |
                            (nybbles[i] & 0x02) |
                            (nybbles[i] & 0x01)   ));

    }
}


void
hamming_decode(std::vector<unsigned char> &codewords,
               std::vector<unsigned char> &bytes,
               unsigned char rdd)
{
    unsigned char p1, p2, p4, p8;
    unsigned char mask;
    unsigned int num_set_bits;
    unsigned int num_set_flags;
    int error_pos = 0;

    for (int i = 0; i < codewords.size(); i++)
    {
//         switch (rdd) {
//             case 4:
//             p8 = parity(codewords[i], mask = (unsigned char)HAMMING_P8_BITMASK);
//             case 3:
//             p4 = parity(codewords[i], mask = (unsigned char)HAMMING_P4_BITMASK >> (4 - rdd));
//             case 2:
//             p2 = parity(codewords[i], mask = (unsigned char)HAMMING_P2_BITMASK >> (4 - rdd));
//             case 1:
//             p1 = parity(codewords[i], mask = (unsigned char)HAMMING_P1_BITMASK >> (4 - rdd));
//             break;
//         }
        p1 = p2 = p4 = p8 = 0;

        error_pos = -1;
        if (p1 != 0) error_pos += 1;
        if (p2 != 0) error_pos += 2;
        if (p4 != 0) error_pos += 4;

        num_set_flags = p1 + p2 + p4;

        // Hamming(4+rdd,4) is only corrective if rdd >= 3
        if (rdd > 2)
        {
            num_set_bits = 0;
            for (int bit_idx = 0; bit_idx < 8; bit_idx++)
            {
            if (codewords[i] & (0x01 << bit_idx))
            {
                num_set_bits++;
            }
            }

            if (error_pos >= 0 && num_set_bits < 6 && num_set_bits > 2)
            {
            codewords[i] ^= (0x80 >> (4-rdd)) >> error_pos;
            }

            num_set_bits = 0;
            for (int bit_idx = 0; bit_idx < 8; bit_idx++)
            {
            if (codewords[i] & (0x01 << bit_idx))
            {
                num_set_bits++;
            }
            }

            // if (rdd == 4)
            // {
            //   if (num_set_bits < 3)
            //   {
            //     codewords[i] = 0;
            //   }
            //   else if (num_set_bits > 5)
            //   {
            //     codewords[i] = 0xFF;
            //   }
            // }
        }

        switch (rdd)
        {
            case 1:
            case 2:
            codewords[i] = codewords[i] & 0x0F;
            break;
            case 3:
            codewords[i] = (((codewords[i] & 0x10) >> 1) | \
                            ((codewords[i] & 0x04))      | \
                            ((codewords[i] & 0x02))      | \
                            ((codewords[i] & 0x01))) & 0x0F;
            break;
            case 4:
            codewords[i] = (((codewords[i] & 0x20) >> 2) | \
                            ((codewords[i] & 0x08) >> 1) | \
                            ((codewords[i] & 0x04) >> 1) | \
                            ((codewords[i] & 0x02) >> 1)) & 0x0F;
            break;
        }

        bytes.push_back(codewords[i] & 0x0F);
    }
}



template <typename T>
void printiterable(T && iterable)
{
    std::cout << std::hex;
    for(auto && e : iterable)
        std::cout << static_cast<unsigned int>(e) << ' ';
    std::cout << std::dec << "\n";
}


int main()
{
    const std::vector<unsigned char> payload_nybbles{0xA, 0xB, 0xC, 0xD, 0xE, 0xF, 0x1, 0x5};
    const unsigned char sf{8};
    const unsigned char cr{1};

    std::cout << "Payload nybbles:\n";
    printiterable(payload_nybbles);

    std::vector<unsigned char> payload_codewords;
    hamming_encode(payload_nybbles, payload_codewords, cr);
    // NOTE: By disabling the hamming parity bits, the input is equal to the output

    std::cout << "Post hamming_encode:\n";
    printiterable(payload_codewords);

    std::vector<unsigned short> payload_symbols;
    interleave(payload_codewords, payload_symbols, sf, cr);

    std::cout << "Post interleave:\n";
    printiterable(payload_symbols);





    std::vector<unsigned char> deinter_codewords;
    deinterleave(payload_symbols, deinter_codewords, sf, cr);

    std::cout << "Post deinterleave:\n";
    printiterable(deinter_codewords);

    std::vector<unsigned char> hamdecode_bytes;
    hamming_decode(deinter_codewords, hamdecode_bytes, cr);

    std::cout << "Post hamming_decode:\n";
    printiterable(hamdecode_bytes);

    return 0;
}