#include <iostream>
#include <cstdint>
#include <cstring>
#include <cassert>

// (7,4) Hamming Code example from chapter 3
class Hamming_coder
{
public:

    static const size_t BLOCK_SIZE = 7;
    static const size_t INFO_SIZE  = 4;

    Hamming_coder()
    {
        init_syndrome_table();
    }
    ~Hamming_coder(){}

    // input size must be 4, output size must be 7
    void encode(uint8_t *input, uint8_t *output)
    {
        for (size_t generator_col=0; generator_col<BLOCK_SIZE; ++generator_col)
        {
            output[generator_col] = 0;
            for (size_t generator_row=0; generator_row<INFO_SIZE; ++generator_row)
            {
                output[generator_col] ^= input[generator_row] * Hamming_generator[generator_row][generator_col];
            }
        }
    } 

    // input size must be 7, output size must be 4
    bool decode(uint8_t *input, uint8_t *output)
    {
      // 1. calculate syndrome
      uint8_t syndrome[(BLOCK_SIZE-INFO_SIZE)];
      calc_syndrome(input, syndrome);
      
      // 2. Look up coset leader based off syndrome
      uint8_t coset_leader_index = (syndrome[0]<<2) | (syndrome[1]<<1) | (syndrome[2] << 0);
      uint8_t *coset_leader = syndrome_table[coset_leader_index];


      // 3. Extract information bits
      std::memcpy(output, &input[(BLOCK_SIZE-INFO_SIZE)], INFO_SIZE*sizeof(uint8_t));

      // 4. Use coset leader as error vector to correct bits
      for (size_t i=0; i<INFO_SIZE; ++i)
      {
        // modulo 2 addition for GF(2)
        output[i] ^= coset_leader[i+(BLOCK_SIZE-INFO_SIZE)];
      }

      return true;
    }

private:
  static const uint8_t Hamming_generator              [INFO_SIZE]  [BLOCK_SIZE];
  static const uint8_t Hamming_parity_check_transpose [BLOCK_SIZE] [BLOCK_SIZE-INFO_SIZE];

  // standard array coset leaders
  static const uint8_t coset_leaders[(1<<(BLOCK_SIZE-INFO_SIZE))][BLOCK_SIZE];
  
  // maps syndromes to coset leaders
  // the number of syndromes is 2^(k-n)
  uint8_t syndrome_table[(1<<(BLOCK_SIZE-INFO_SIZE))][BLOCK_SIZE];

  // each syndrome in the table is (error vec i)*transpose(H)
  void init_syndrome_table()
  {
    for (size_t i=0; i<(1<<(BLOCK_SIZE-INFO_SIZE)); ++i)
    {
      uint8_t syndrome_bits[(BLOCK_SIZE-INFO_SIZE)];

      // calculate the syndrome
      calc_syndrome(coset_leaders[i], syndrome_bits);

      // convert into a number for use as an array index
      uint8_t index = (syndrome_bits[0]<<2) | (syndrome_bits[1]<<1) | (syndrome_bits[2]<<0);

      // store the corresponding coset leader in the map
      std::memcpy(&syndrome_table[index], coset_leaders[i], BLOCK_SIZE*sizeof(uint8_t));
    }
  }

  // compute r * transpose(H)
  void calc_syndrome(const uint8_t *input, uint8_t *output)
  {
      for (size_t parity_col=0; parity_col<(BLOCK_SIZE-INFO_SIZE); ++parity_col)
      {
          output[parity_col] = 0;
          for (size_t parity_row=0; parity_row<BLOCK_SIZE; ++parity_row)
          {
              // binary Modulo 2 addition for GF(2)
              output[parity_col] ^= input[parity_row] * Hamming_parity_check_transpose[parity_row][parity_col];
          }
      }
  }

}; // end class Hamming_coder

const uint8_t Hamming_coder::coset_leaders[8][7] =
{
  {0,0,0,0,0,0,0},
  {1,0,0,0,0,0,0},
  {0,1,0,0,0,0,0},
  {0,0,1,0,0,0,0},
  {0,0,0,1,0,0,0},
  {0,0,0,0,1,0,0},
  {0,0,0,0,0,1,0},
  {0,0,0,0,0,0,1},
};


// Generator matrix
const uint8_t Hamming_coder::Hamming_generator[4][7] = {
  {1,1,0,1,0,0,0},
  {0,1,1,0,1,0,0},
  {1,1,1,0,0,1,0},
  {1,0,1,0,0,0,1}
};

// Parity check matrix transposed
const uint8_t Hamming_coder::Hamming_parity_check_transpose[7][3] = {
  {1,0,0},
  {0,1,0},
  {0,0,1},
  {1,1,0},
  {0,1,1},
  {1,1,1},
  {1,0,1}
};

// example correction from book
void test_book_example()
{
  uint8_t input[Hamming_coder::INFO_SIZE] = {1,0,1,1};
  uint8_t encoded_input[Hamming_coder::BLOCK_SIZE];
  uint8_t decoded_output[Hamming_coder::INFO_SIZE];

  Hamming_coder hamming_coder;

  // encode
  hamming_coder.encode(input, encoded_input);

  std::cout << "Encoded input: ";
  for (size_t i=0; i<Hamming_coder::BLOCK_SIZE; ++i)
  {
    std::cout << (int)encoded_input[i];
  }
  std::cout << std::endl;

  // introduce the error
  encoded_input[Hamming_coder::BLOCK_SIZE-3] = !encoded_input[Hamming_coder::BLOCK_SIZE-3];

  std::cout << "Error Encoded input: ";
  for (size_t i=0; i<Hamming_coder::BLOCK_SIZE; ++i)
  {
    std::cout << (int)encoded_input[i];
  }
  std::cout << std::endl;

  // decode
  hamming_coder.decode(encoded_input, decoded_output);
  std::cout << "Decoded output: ";
  for (size_t i=0; i<Hamming_coder::INFO_SIZE; ++i)
  {
    std::cout << (int)decoded_output[i];
  }
  std::cout << std::endl;
}

// test all possible codewords and errors
void test_coder()
{
  Hamming_coder hamming_coder;

  uint8_t decoded_output[Hamming_coder::INFO_SIZE];

  // test all codewords
  // 2^(INFO SIZE) - 1 = 15
  for (size_t input=0; input<(1<<Hamming_coder::INFO_SIZE); ++input)
  {
    // fill in info bits
    uint8_t input_bits[Hamming_coder::INFO_SIZE];
    for (size_t i=0; i<Hamming_coder::INFO_SIZE; ++i)
    {
      input_bits[i] = (input>>(Hamming_coder::INFO_SIZE-i-1)) & 1;
    }

    // encode the codeword
    uint8_t encoded_input[Hamming_coder::BLOCK_SIZE];
    hamming_coder.encode(input_bits, encoded_input);   

    // introduce all possible errors
    for (size_t j=0; j<Hamming_coder::BLOCK_SIZE; ++j)
    {
      uint8_t error_codeword[Hamming_coder::BLOCK_SIZE];
      std::memcpy(error_codeword, encoded_input, Hamming_coder::BLOCK_SIZE*sizeof(uint8_t));

      error_codeword[j] = !error_codeword[j];

      // decode the codeword
      hamming_coder.decode(error_codeword, decoded_output);

      // verify the error was corrected
      for (size_t k=0; k<Hamming_coder::INFO_SIZE; ++k)
      {
        assert(decoded_output[k] == input_bits[k]);
      }
    }
 
  }

  std::cout << "All errors corrected!" << std::endl;
}

int main(void)
{
//  test_book_example();
  test_coder();

  return 0;
}