#include "htio2/Kmer.h"
#include "htio2/Error.h"

#include <cmath>

#include <cstdio>
#include <limits>

using namespace std;

namespace htio2
{

void encode_nt5(const std::string& seq, EncodedSeq& encoded_seq)
{
    const size_t len = seq.length();
    encoded_seq.resize(len);

    for (size_t i = 0; i < len; i++)
    {
        EncodedType curr = encode_nt5(seq[i]);
        encoded_seq[i] = curr;
    }
}

void encode_nt4(const string &seq, EncodedSeq &encoded_seq)
{
    const size_t len = seq.length();
    encoded_seq.resize(len);

    for (size_t i = 0; i < len; i++)
    {
        EncodedType curr = encode_nt4(seq[i]);
        encoded_seq[i] = curr;
    }
}

void encode_aa(const std::string& seq, EncodedSeq& encoded_seq)
{
    const size_t len = seq.length();
    encoded_seq.resize(len);

    for (size_t i = 0; i < len; i++)
    {
        EncodedType curr = encode_aa(seq[i]);
        encoded_seq[i] = curr;
    }
}

void decode_nt5(const EncodedSeq& encoded_seq, std::string& seq)
{
    const size_t len = encoded_seq.size();
    seq.resize(len);

    for (size_t i = 0; i < len; i++)
    {
        char curr = decode_nt5(encoded_seq[i]);
        seq[i] = curr;
    }
}

void decode_nt4(const EncodedSeq &encoded_seq, string &seq)
{
    const size_t len = encoded_seq.size();
    seq.resize(len);

    for (size_t i = 0; i < len; i++)
    {
        char curr = decode_nt4(encoded_seq[i]);
        seq[i] = curr;
    }
}

void decode_aa(const EncodedSeq& encoded_seq, std::string& seq)
{
    const size_t len = encoded_seq.size();
    seq.resize(len);

    for (size_t i = 0; i < len; i++)
    {
        char curr = decode_aa(encoded_seq[i]);
        seq[i] = curr;
    }
}

template <EncodedType nary>
inline void _revcom_(const EncodedSeq& in, EncodedSeq& out)
{
    const size_t len = in.size();
    out.resize(len);
    for (size_t i = 0; i < len; i++)
    {
        const EncodedType curr = in[i];
        size_t i_out = len - i - 1;
        if (curr == ENCODED_ERROR)
            out[i_out] = ENCODED_ERROR;
        else
            out[i_out] = nary - curr;
    }
}

void revcom_nt5(const EncodedSeq& in, EncodedSeq& out)
{
    _revcom_<5>(in, out);
}

void revcom_nt4(const EncodedSeq& in, EncodedSeq& out)
{
    _revcom_<4>(in, out);
}

bool operator==(const EncodedSeq& a, const EncodedSeq& b)
{
    const size_t len = a.size();
    if (len == b.size())
        return false;

    for (size_t i = 0; i < len; i++)
    {
        if (a[i] != b[i]) return false;
    }

    return true;
}

template <typename kmer_type, int nary>
inline bool gen_kmer_impl(const EncodedSeq& encoded_seq, vector<kmer_type>& result, size_t kmer_size)
{
    if (kmer_size > KmerTrait<nary, kmer_type>::max_size())
        return false;

    size_t len = encoded_seq.size();
    if (len < kmer_size)
        return false;

    const size_t num_kmer = len - kmer_size + 1;
    result.resize(num_kmer);

    const size_t digit_base = std::pow(nary, kmer_size-1);
    size_t i_begin = 0;
    const size_t i_end = num_kmer;

    for (;;)
    {
        if (i_begin >= i_end) break;

        // build the first kmer
        bool init_got_unusable = false;
        size_t kmer = 0;

        for (size_t k = 0; k < kmer_size; ++k)
        {
            size_t i = i_begin + k;
            EncodedType curr = encoded_seq[i];

            if (curr != ENCODED_ERROR)
                kmer += curr * std::pow(nary, k);
            else
            {
                // mark affected position as error
                for (size_t i_error = i_begin;
                     i_error <= i && i_error < i_end;
                     ++i_error)
                {
                    result[i_error] = numeric_limits<kmer_type>::max();
                }

                // skip to the base next to unusable
                i_begin += k + 1;
                init_got_unusable = true;
                break;
            }
        }

        // if got unusable, build first kmer again again
        if (init_got_unusable)
            continue;

        result[i_begin] = kmer;

        // elongate kmer
        // each time, read one base, and remove previoius base
        bool elong_got_error = false;
        for (size_t i = i_begin + 1; i < i_end; ++i)
        {
            size_t i_kmer_tail = i + kmer_size - 1;
            EncodedType prev_enc = encoded_seq[i - 1];
            EncodedType tail_enc = encoded_seq[i_kmer_tail];

            if (tail_enc != ENCODED_ERROR)
            {
                kmer -= prev_enc;
                kmer /= nary;
                kmer += tail_enc * digit_base;
                result[i] = kmer;
            }
            else
            {
                // mark affected position as error
                for (size_t i_error = i;
                     i_error <= i_kmer_tail && i_error < i_end;
                     ++i_error)
                {
                    result[i_error] = numeric_limits<kmer_type>::max();
                }

                // skip to the base next to unusable
                i_begin = i_kmer_tail + 1;
                elong_got_error = true;
                break;
            }
        }

        // if we don't get any unusable base, we've finished everything
        if (elong_got_error)
            continue;
        else
            break;
    }

    return true;
}

typedef void (*encode_func) (const std::string& seq, EncodedSeq& encoded_seq);

template<typename kmer_type, int nary>
inline kmer_type gen_kmer_simple_impl(encode_func func, const std::string& seq)
{
    if (seq.length() > KmerTrait<nary, kmer_type>::max_size())
        return numeric_limits<kmer_type>::max();

    EncodedSeq enc;
    func(seq, enc);

    const size_t size = seq.length();
    kmer_type kmer = 0;
    kmer_type digit_base = 1;
    for (size_t i = 0; i < size; i++)
    {
        if (enc[i] == ENCODED_ERROR)
        {
            kmer = numeric_limits<kmer_type>::max();
            break;
        }
        kmer += enc[i] * digit_base;
        digit_base *= nary;
    }
    return kmer;
}

template<>
bool gen_kmer_nt5<uint64_t>(const EncodedSeq& encoded_seq, std::vector<uint64_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint64_t, 5>(encoded_seq, result, kmer_size);
}

template<>
bool gen_kmer_nt5<uint32_t>(const EncodedSeq& encoded_seq, std::vector<uint32_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint32_t, 5>(encoded_seq, result, kmer_size);
}

template<>
bool gen_kmer_nt5<uint16_t>(const EncodedSeq& encoded_seq, std::vector<uint16_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint16_t, 5>(encoded_seq, result, kmer_size);
}

template<>
uint64_t gen_kmer_nt5<uint64_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint64_t, 5>(encode_nt5, seq);
}

template<>
uint32_t gen_kmer_nt5<uint32_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint32_t, 5>(encode_nt5, seq);
}

template<>
uint16_t gen_kmer_nt5<uint16_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint16_t, 5>(encode_nt5, seq);
}

template<>
bool gen_kmer_nt4<uint64_t>(const EncodedSeq& encoded_seq, std::vector<uint64_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint64_t, 4>(encoded_seq, result, kmer_size);
}

template<>
bool gen_kmer_nt4<uint32_t>(const EncodedSeq& encoded_seq, std::vector<uint32_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint32_t, 4>(encoded_seq, result, kmer_size);
}

template<>
bool gen_kmer_nt4<uint16_t>(const EncodedSeq& encoded_seq, std::vector<uint16_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint16_t, 4>(encoded_seq, result, kmer_size);
}

template<>
uint64_t gen_kmer_nt4<uint64_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint64_t, 4>(encode_nt4, seq);
}

template<>
uint32_t gen_kmer_nt4<uint32_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint32_t, 4>(encode_nt4, seq);
}

template<>
uint16_t gen_kmer_nt4<uint16_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint16_t, 4>(encode_nt4, seq);
}

template<>
bool gen_kmer_aa<uint64_t>(const EncodedSeq& encoded_seq, std::vector<uint64_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint64_t, 26>(encoded_seq, result, kmer_size);
}

template<>
bool gen_kmer_aa<uint32_t>(const EncodedSeq& encoded_seq, std::vector<uint32_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint32_t, 26>(encoded_seq, result, kmer_size);
}

template<>
bool gen_kmer_aa<uint16_t>(const EncodedSeq& encoded_seq, std::vector<uint16_t>& result, size_t kmer_size)
{
    return gen_kmer_impl<uint16_t, 26>(encoded_seq, result, kmer_size);
}

template<>
uint64_t gen_kmer_aa<uint64_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint64_t, 26>(encode_aa, seq);
}

template<>
uint32_t gen_kmer_aa<uint32_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint32_t, 26>(encode_aa, seq);
}

template<>
uint16_t gen_kmer_aa<uint16_t>(const std::string& seq)
{
    return gen_kmer_simple_impl<uint16_t, 26>(encode_aa, seq);
}

typedef char (*decode_func_type)(EncodedType input);

template<typename kmer_type, int nary>
inline string decode_kmer_impl(decode_func_type decode_func, kmer_type kmer, size_t kmer_size)
{
    string result;
    if (kmer == numeric_limits<kmer_type>::max())
        return result;
//    printf("decode %d-ary %lu-mer %lu\n", nary, kmer_size, uint64_t(kmer));

    for (size_t i=0; i<kmer_size; i++)
    {
        EncodedType enc = kmer % nary;
        char dec = decode_func(enc);
//        printf("  %lu: residue %c (%d) from %lu\n", i, dec, int(enc), uint64_t(kmer));
        result.push_back(dec);
        kmer -= enc;
        kmer /= nary;
    }

    return result;
}

template<>
std::string decode_kmer_nt5<uint64_t>(uint64_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint64_t, 5>(decode_nt5, kmer, kmer_size);
}

template<>
std::string decode_kmer_nt5<uint32_t>(uint32_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint32_t, 5>(decode_nt5, kmer, kmer_size);
}

template<>
std::string decode_kmer_nt5<uint16_t>(uint16_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint16_t, 5>(decode_nt5, kmer, kmer_size);
}

template<>
std::string decode_kmer_nt4<uint64_t>(uint64_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint64_t, 4>(decode_nt4, kmer, kmer_size);
}

template<>
std::string decode_kmer_nt4<uint32_t>(uint32_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint32_t, 4>(decode_nt4, kmer, kmer_size);
}

template<>
std::string decode_kmer_nt4<uint16_t>(uint16_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint16_t, 4>(decode_nt4, kmer, kmer_size);
}

template<>
std::string decode_kmer_aa<uint64_t>(uint64_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint64_t, 26>(decode_aa, kmer, kmer_size);
}

template<>
std::string decode_kmer_aa<uint32_t>(uint32_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint32_t, 26>(decode_aa, kmer, kmer_size);
}

template<>
std::string decode_kmer_aa<uint16_t>(uint16_t kmer, size_t kmer_size)
{
    return decode_kmer_impl<uint16_t, 26>(decode_aa, kmer, kmer_size);
}

} // namespace htio2
