#include <vector>
#include <algorithm>

#include "bwt.hh"

namespace
{
    typedef unsigned char byte;

    class BWTComparator
    {
        const std::vector<byte>& input;
        const size_t size;

     public:
        BWTComparator(const std::vector<byte>& i): input(i), size(i.size()) {}

        bool operator()(size_t li, size_t ri) const
        {
            if(li == ri) return false;
            size_t amount = size;
            while(input[li] == input[ri])
            {
                if(++li == size) li = 0;
                if(++ri == size) ri = 0;
                if(--amount == 0) return false;
            }
            return input[li] < input[ri];
        }
    };
}

const std::vector<byte>
    BWT_encode(const std::vector<byte>& input,
               size_t& primaryIndex)
{
    const size_t size = input.size();

    std::vector<size_t> indices(size);
    for(size_t i = 0; i < size; ++i) indices[i] = i;
    std::sort(indices.begin(), indices.end(), BWTComparator(input));
    
    std::vector<unsigned char> encoded ( size );

    primaryIndex = 0;
    for(size_t i = 0; i < size; ++i)
    {
        encoded[i] = input[(indices[i] + size - 1) % size];
        if(indices[i] == 1) primaryIndex = i;
    }
    return encoded;
}

const std::vector<byte>
    BWT_decode(const std::vector<byte>& encoded,
               size_t primaryIndex)
{
    const size_t size = encoded.size();

    size_t buckets[256] = { 0 };
    
    std::vector<size_t> indices(size);

    std::vector<byte> F(size), decoded(size);

    for(size_t i = 0; i < size; ++i)
        ++buckets[encoded[i]];

    for(size_t i = 0, k = 0; i < 256; ++i)
        for(size_t j = 0; j < buckets[i]; ++j)
            F[k++] = byte(i);

    for(size_t i = 0, j = 0; i < 256; ++i)
    {
        while(i > F[j] && j < size) ++j;
        buckets[i] = j;
    }

    for(size_t i = 0; i < size; ++i)
        indices[buckets[encoded[i]]++] = i;

    for(size_t i = 0, j = primaryIndex; i < size; ++i, j = indices[j])
        decoded[i] = encoded[j];
    
    return decoded;
}