// Implementation of ViterbiCodec.
//
// Author: Min Xu <xukmin@gmail.com>
// Date: 01/30/2015

#include "viterbi.h"

#include <algorithm>
#include <cassert>
#include <iostream>
#include <limits>
#include <string>
#include <utility>
#include <vector>

namespace {

int HammingDistance(const std::string& x, const std::string& y)
{
    assert(x.size() == y.size());
    int distance = 0;
    for (int i = 0; i < (int)x.size(); i++) {
        distance += x[i] != y[i];
    }
    return distance;
}

} // namespace

std::ostream& operator<<(std::ostream& os, const ViterbiCodec& codec)
{
    os << "ViterbiCodec(" << codec.constraint() << ", {";
    const std::vector<int>& polynomials = codec.polynomials();
    assert(!polynomials.empty());
    os << polynomials.front();
    for (int i = 1; i < (int)polynomials.size(); i++) {
        os << ", " << polynomials[i];
    }
    return os << "})";
}

int ReverseBits(int num_bits, int input)
{
    assert(input < (1 << num_bits));
    int output = 0;
    while (num_bits-- > 0) {
        output = (output << 1) + (input & 1);
        input >>= 1;
    }
    return output;
}

ViterbiCodec::ViterbiCodec(int constraint, const std::vector<int>& polynomials)
    : constraint_(constraint), polynomials_(polynomials)
{
    assert(!polynomials_.empty());
    for (int i = 0; i < (int)polynomials_.size(); i++) {
        assert(polynomials_[i] > 0);
        assert(polynomials_[i] < (1 << constraint_));
    }
    InitializeOutputs();
}

int ViterbiCodec::num_parity_bits() const { return polynomials_.size(); }

int ViterbiCodec::NextState(int current_state, int input) const
{
    return (current_state >> 1) | (input << (constraint_ - 2));
}

std::string ViterbiCodec::Output(int current_state, int input) const
{
    return outputs_.at(current_state | (input << (constraint_ - 1)));
}

std::string ViterbiCodec::Encode(const std::string& bits) const
{
    std::string encoded;
    int state = 0;

    // Encode the message bits.
    for (int i = 0; i < (int)bits.size(); i++) {
        char c = bits[i];
        assert(c == '0' || c == '1');
        int input = c - '0';
        encoded += Output(state, input);
        state = NextState(state, input);
    }

    // Encode (constaint_ - 1) flushing bits.
    for (int i = 0; i < constraint_ - 1; i++) {
        encoded += Output(state, 0);
        state = NextState(state, 0);
    }

    return encoded;
}

void ViterbiCodec::InitializeOutputs()
{
    outputs_.resize(1 << constraint_);
    for (int i = 0; i < (int)outputs_.size(); i++) {
        for (int j = 0; j < num_parity_bits(); j++) {
            // Reverse polynomial bits to make the convolution code simpler.
            int polynomial = ReverseBits(constraint_, polynomials_[j]);
            int input = i;
            int output = 0;
            for (int k = 0; k < constraint_; k++) {
                output ^= (input & 1) & (polynomial & 1);
                polynomial >>= 1;
                input >>= 1;
            }
            outputs_[i] += output ? "1" : "0";
        }
    }
}

int ViterbiCodec::BranchMetric(const std::string& bits,
                               int source_state,
                               int target_state) const
{
    assert(bits.size() == num_parity_bits());
    assert((target_state & ((1 << (constraint_ - 2)) - 1)) == source_state >> 1);
    const std::string output = Output(source_state, target_state >> (constraint_ - 2));

    return HammingDistance(bits, output);
}

std::pair<int, int> ViterbiCodec::PathMetric(const std::string& bits,
                                             const std::vector<int>& prev_path_metrics,
                                             int state) const
{
    int s = (state & ((1 << (constraint_ - 2)) - 1)) << 1;
    int source_state1 = s | 0;
    int source_state2 = s | 1;

    int pm1 = prev_path_metrics[source_state1];
    if (pm1 < std::numeric_limits<int>::max()) {
        pm1 += BranchMetric(bits, source_state1, state);
    }
    int pm2 = prev_path_metrics[source_state2];
    if (pm2 < std::numeric_limits<int>::max()) {
        pm2 += BranchMetric(bits, source_state2, state);
    }

    if (pm1 <= pm2) {
        return std::make_pair(pm1, source_state1);
    } else {
        return std::make_pair(pm2, source_state2);
    }
}

void ViterbiCodec::UpdatePathMetrics(const std::string& bits,
                                     std::vector<int>* path_metrics,
                                     Trellis* trellis) const
{
    std::vector<int> new_path_metrics(path_metrics->size());
    std::vector<int> new_trellis_column(1 << (constraint_ - 1));
    for (int i = 0; i < (int)path_metrics->size(); i++) {
        std::pair<int, int> p = PathMetric(bits, *path_metrics, i);
        new_path_metrics[i] = p.first;
        new_trellis_column[i] = p.second;
    }

    *path_metrics = new_path_metrics;
    trellis->push_back(new_trellis_column);
}

std::string ViterbiCodec::Decode(const std::string& bits) const
{
    // Compute path metrics and generate trellis.
    Trellis trellis;
    std::vector<int> path_metrics(1 << (constraint_ - 1),
                                  std::numeric_limits<int>::max());
    path_metrics.front() = 0;
    for (int i = 0; i < (int)bits.size(); i += num_parity_bits()) {
        std::string current_bits(bits, i, num_parity_bits());
        // If some bits are missing, fill with trailing zeros.
        // This is not ideal but it is the best we can do.
        if ((int)current_bits.size() < num_parity_bits()) {
            current_bits.append(
                std::string(num_parity_bits() - current_bits.size(), '0'));
        }
        UpdatePathMetrics(current_bits, &path_metrics, &trellis);
    }

    // Traceback.
    std::string decoded;
    int state =
        std::min_element(path_metrics.begin(), path_metrics.end()) - path_metrics.begin();
    for (int i = trellis.size() - 1; i >= 0; i--) {
        decoded += state >> (constraint_ - 2) ? "1" : "0";
        state = trellis[i][state];
    }
    std::reverse(decoded.begin(), decoded.end());

    // Remove (constraint_ - 1) flushing bits.
    return decoded.substr(0, decoded.size() - constraint_ + 1);
}
