#ifndef HTIO2_KMER_H
#define	HTIO2_KMER_H

#include <vector>
#include <map>

#include <stdint.h>

#include "htio2/Enum.h"

namespace htio2
{

typedef enum
{
    SEQ_TYPE_NT4 = 0,
    SEQ_TYPE_NT5 = 1,
    SEQ_TYPE_AA  = 2,
} SeqType;

typedef uint8_t EncodedType;
#define ENCODED_ERROR 255
//#define RESIDUE_ERROR -1
typedef std::vector<EncodedType> EncodedSeq;

#define KMER_NT4_LIMIT_64 31
#define KMER_NT4_LIMIT_32 15
#define KMER_NT4_LIMIT_16 7
#define KMER_NT5_LIMIT_64 27
#define KMER_NT5_LIMIT_32 13
#define KMER_NT5_LIMIT_16 6
#define KMER_AA_LIMIT_64 13
#define KMER_AA_LIMIT_32 6
#define KMER_AA_LIMIT_16 3

/**
 * Properties for a N-ary kmer stored in a primitive data type of specific size.
 */
template<int nary, typename KmerType>
struct KmerTrait;

template<>
struct KmerTrait<4, uint64_t>
{
    static size_t max_size() { return KMER_NT4_LIMIT_64; }
};

template<>
struct KmerTrait<4, uint32_t>
{
    static size_t max_size() { return KMER_NT4_LIMIT_32; }
};

template<>
struct KmerTrait<4, uint16_t>
{
    static size_t max_size() { return KMER_NT4_LIMIT_16; }
};

template<>
struct KmerTrait<5, uint64_t>
{
    static size_t max_size() { return KMER_NT5_LIMIT_64; }
};

template<>
struct KmerTrait<5, uint32_t>
{
    static size_t max_size() { return KMER_NT5_LIMIT_32; }
};

template<>
struct KmerTrait<5, uint16_t>
{
    static size_t max_size() { return KMER_NT5_LIMIT_16; }
};

template<>
struct KmerTrait<26, uint64_t>
{
    static size_t max_size() { return KMER_AA_LIMIT_64; }
};

template<>
struct KmerTrait<26, uint32_t>
{
    static size_t max_size() { return KMER_AA_LIMIT_32; }
};

template<>
struct KmerTrait<26, uint16_t>
{
    static size_t max_size() { return KMER_AA_LIMIT_16; }
};


#define ENCODED_NT5_A 0
#define ENCODED_NT5_T 4
#define ENCODED_NT5_G 1
#define ENCODED_NT5_C 3
#define ENCODED_NT5_N 2

/**
 * map ATGCN to numbers 0-4, and map all other characters to 255
 * @param input a nucleotide in ATGCN
 * @return a number from 0 to 4
 */
inline EncodedType encode_nt5(char input)
{
    switch (input)
    {
    case 'A':
    case 'a':
        return ENCODED_NT5_A;
    case 'T':
    case 't':
        return ENCODED_NT5_T;
    case 'G':
    case 'g':
        return ENCODED_NT5_G;
    case 'C':
    case 'c':
        return ENCODED_NT5_C;
    case 'N':
    case 'n':
        return ENCODED_NT5_N;
    default:
        return ENCODED_ERROR;
    }
}


#define ENCODED_NT4_A 0
#define ENCODED_NT4_T 3
#define ENCODED_NT4_G 1
#define ENCODED_NT4_C 2

/**
 * map ATGC to numbers 0-3, and map all other characters to 255
 * @param input a nucleotide in ATGC
 * @return a number from 0 to 3
 */
inline EncodedType encode_nt4(char input)
{
    switch (input)
    {
    case 'A':
    case 'a':
        return ENCODED_NT4_A;
    case 'T':
    case 't':
        return ENCODED_NT4_T;
    case 'G':
    case 'g':
        return ENCODED_NT4_G;
    case 'C':
    case 'c':
        return ENCODED_NT4_C;
    default:
        return ENCODED_ERROR;
    }
}

/**
 * map amino acids residues to 0-25, "*" to 'O' minus 'A', and all other characters to 255
 * @param input a amino acids residue
 * @return numbers from 0 to 25.
 */
inline EncodedType encode_aa(char input)
{
    if (input == '*') input = 'O';

    if ('a' <= input && input <= 'z') return input - 'a';
    else if ('A' <= input && input <= 'Z') return input - 'A';
    else return ENCODED_ERROR;
}

/**
 * Map encoded nucleotides back to characters: 0-4 will be mapped to ATGCN;
 * character '-' and '.' will be specially treated, which will be returned
 * as-is; and all other characters will be mapped to 'N'.
 *
 * @param input encoded bases in 0-4
 * @return bases in ATGCN.
 *
 * @see encode_nt5(char);
 */
inline char decode_nt5(EncodedType input)
{
    switch (input)
    {
    case ENCODED_NT5_A:
        return 'A';
    case ENCODED_NT5_T:
        return 'T';
    case ENCODED_NT5_G:
        return 'G';
    case 3:
        return 'C';
    case '-':
    case '.':
        return input;
    case 2:
    default:
        return 'N';
    }
}

/**
 * Map encoded nucleotides back to characters: 0-3 will be mapped to ATGC;
 * character '-' and '.' will be specially treated, which will be returned
 * as-is; and all other characters will be mapped to 'N'.
 *
 * @param input encoded bases in 0-4
 * @return bases in ATGCN.
 *
 * @see encode_nt4(char)
 */
inline char decode_nt4(EncodedType input)
{
    switch (input)
    {
    case 0:
        return 'A';
    case 3:
        return 'T';
    case 1:
        return 'G';
    case 2:
        return 'C';
    case '-':
    case '.':
        return input;
    default:
        return 'N';
    }
}

/**
 * Map encoded amino acids residues back to characters: nomal residues will
 * be mapped back to characters; 'O' minus 'A' will be mapped to '*'; other
 * things will be mapped to 'X'.
 *
 * @param input residues in 0-25
 * @return residues in characters.
 *
 * @see encode_aa(char);
 */
inline char decode_aa(EncodedType input)
{
    char result;
    if (0 <= input && input <= 25)
        result = input + 'A';
    else if (input == '-' || input == '.')
        result = input;
    else
        result = 'X';

    if (result == 'O') result = '*';

    return result;
}

/**
 * Encode a string of nucleotides, ATGCN are allowed.
 * @param seq input sequence
 * @param encoded_seq object to store encoded result
 *
 * @see encode_nt5(char)
 */
void encode_nt5(const std::string& seq, EncodedSeq& encoded_seq);

/**
 * Encode a string of nucleotides, ATGC are allowed.
 * @param seq input sequence
 * @param encoded_seq object to store encoded result
 *
 * @see encode_nt4(char)
 */
void encode_nt4(const std::string& seq, EncodedSeq& encoded_seq);

/**
 * Encode a string of amino acid residues.
 * @param seq input sequence
 * @param encoded_seq object to store encoded result
 *
 * @see encode_aa(char)
 */
void encode_aa(const std::string& seq, EncodedSeq& encoded_seq);

/**
 * Decode an array of encoded nucleotides.
 *
 * @param encoded_seq sequence to decode
 * @param seq string object to store decoded result
 *
 * @see decode_nt5(EncodedType)
 */
void decode_nt5(const EncodedSeq& encoded_seq, std::string& seq);

/**
 * Decode an array of encoded nucleotides.
 *
 * @param encoded_seq sequence to decode
 * @param seq string object to store decoded result
 *
 * @see decode_nt4(EncodedType)
 */
void decode_nt4(const EncodedSeq& encoded_seq, std::string& seq);

/**
 * Decode an array of encoded nucleotides.
 *
 * @param encoded_seq sequence to decode
 * @param seq string object to store decoded result
 *
 * @see decode_aa(EncodedType)
 */
void decode_aa(const EncodedSeq& encoded_seq, std::string& seq);

/**
 * do reverse-complement for an encoded nucleotide sequence
 * @param in input sequence containing encoded ATGCN
 * @param out object to store result
 */
void revcom_nt5(const EncodedSeq& in, EncodedSeq& out);

/**
 * do reverse-complement for an encoded nucleotide sequence
 * @param in input sequence containing encoded ATGC
 * @param out object to store result
 */
void revcom_nt4(const EncodedSeq& in, EncodedSeq& out);

/**
 * test whether two encoded sequences are identical
 * @param a
 * @param b
 * @return 
 */
bool operator==(const EncodedSeq& a, const EncodedSeq& b);

/**
 * Generate a series of kmers along encoded nucleotide sequence. Each kmer is
 * generated as a 5-ary number of K digits, and is stored in an unsigned
 * integer of specified size. On positions where kmer has invalid base,
 * TYPE_MAX will be set.
 *
 * @param encoded_seq encoded nucleotide sequence
 * @param result series of generated kmers
 * @param kmer_size
 * @return false if kmer size is larger than input sequence, or kmer size is
 * larger than allowed kmer size limit for specified data type.
 */
template<typename KmerType>
bool gen_kmer_nt5(const EncodedSeq& encoded_seq, std::vector<KmerType>& result, size_t kmer_size);

template<>
bool gen_kmer_nt5<uint64_t>(const EncodedSeq& encoded_seq, std::vector<uint64_t>& result, size_t kmer_size);
template<>
bool gen_kmer_nt5<uint32_t>(const EncodedSeq& encoded_seq, std::vector<uint32_t>& result, size_t kmer_size);
template<>
bool gen_kmer_nt5<uint16_t>(const EncodedSeq& encoded_seq, std::vector<uint16_t>& result, size_t kmer_size);

/**
 * Generate one kmer from nucleotide sequence in ATGCN. The kmer is generated
 * as a 5-ary number of K digits, and is stored in an unsigned integer of
 * specified size. The whole sequence is generated to one kmer, and sequence
 * length is used as kmer size.
 *
 * @param seq input sequence
 * @return generated kmer
 */
template<typename KmerType>
KmerType gen_kmer_nt5(const std::string& seq);

template<>
uint64_t gen_kmer_nt5<uint64_t>(const std::string& seq);
template<>
uint32_t gen_kmer_nt5<uint32_t>(const std::string& seq);
template<>
uint16_t gen_kmer_nt5<uint16_t>(const std::string& seq);

/**
 * Generate a series of kmers along encoded nucleotide sequence. Each kmer is
 * generated as a 4-ary number of K digits, and is stored in an unsigned
 * integer of specified size. On positions where kmer has invalid base,
 * TYPE_MAX will be set.
 *
 * @param encoded_seq encoded nucleotide sequence
 * @param result series of generated kmers
 * @param kmer_size
 * @return false if kmer size is larger than input sequence, or kmer size is
 * larger than allowed kmer size limit for specified data type.
 */
template<typename KmerType>
bool gen_kmer_nt4(const EncodedSeq& encoded_seq, std::vector<KmerType>& result, size_t kmer_size);

template<>
bool gen_kmer_nt4<uint64_t>(const EncodedSeq& encoded_seq, std::vector<uint64_t>& result, size_t kmer_size);
template<>
bool gen_kmer_nt4<uint32_t>(const EncodedSeq& encoded_seq, std::vector<uint32_t>& result, size_t kmer_size);
template<>
bool gen_kmer_nt4<uint16_t>(const EncodedSeq& encoded_seq, std::vector<uint16_t>& result, size_t kmer_size);


/**
 * Generate one kmer from nucleotide sequence in ATGC. The kmer is generated
 * as a 4-ary number of K digits, and is stored in an unsigned integer of
 * specified size. The whole sequence is generated to one kmer, and sequence
 * length is used as kmer size.
 *
 * @param seq input sequence
 * @return generated kmer
 */
template<typename KmerType>
KmerType gen_kmer_nt4(const std::string& seq);

template<>
uint64_t gen_kmer_nt4<uint64_t>(const std::string& seq);
template<>
uint32_t gen_kmer_nt4<uint32_t>(const std::string& seq);
template<>
uint16_t gen_kmer_nt4<uint16_t>(const std::string& seq);

/**
 * Generate a series of kmers along encoded amino acids sequence. Each kmer is
 * generated as a 26-ary number of K digits, and is stored in an unsigned
 * integer of specified size. On positions where kmer has invalid base,
 * TYPE_MAX will be set.
 *
 * @param encoded_seq encoded amino acids sequence
 * @param result series of generated kmers
 * @param kmer_size
 * @return false if kmer size is larger than input sequence, or kmer size is
 * larger than allowed kmer size limit for specified data type.
 */
template<typename KmerType>
bool gen_kmer_aa(const EncodedSeq& encoded_seq, std::vector<KmerType>& result, size_t kmer_size);

template<>
bool gen_kmer_aa<uint64_t>(const EncodedSeq& encoded_seq, std::vector<uint64_t>& result, size_t kmer_size);
template<>
bool gen_kmer_aa<uint32_t>(const EncodedSeq& encoded_seq, std::vector<uint32_t>& result, size_t kmer_size);
template<>
bool gen_kmer_aa<uint16_t>(const EncodedSeq& encoded_seq, std::vector<uint16_t>& result, size_t kmer_size);

/**
 * Generate one kmer from amino acids sequence. The kmer is generated as a
 * 26-ary number of K digits, and is stored in an unsigned integer of specified
 * size. The whole sequence is generated to one kmer, and sequence length is
 * used as kmer size.
 *
 * @param seq input sequence
 * @return generated kmer
 */
template<typename KmerType>
KmerType gen_kmer_aa(const std::string& seq);

template<>
uint64_t gen_kmer_aa<uint64_t>(const std::string& seq);
template<>
uint32_t gen_kmer_aa<uint32_t>(const std::string& seq);
template<>
uint16_t gen_kmer_aa<uint16_t>(const std::string& seq);

/**
 * test whether two series of kmers are identical
 * @param a
 * @param b
 * @return 
 */
template<typename KmerType>
bool operator==(const std::vector<KmerType>& a, const std::vector<KmerType>& 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;
}

/**
 * Generate mapping from kmer to position. Position is 0-based.
 * @param kmer_list series of kmers
 * @param kmer_pos a multimap with kmers as key.
 */
template<typename K, typename P>
void summ_kmer_pos(const std::vector<K>& kmer_list, std::multimap<K, P>& result)
{
    for (P i = 0; i < kmer_list.size(); i++)
        result.insert(std::make_pair(kmer_list[i], i));
}

/**
 * decode nucleotide kmer to human-readable string
 * @param kmer
 * @param kmer_size
 * @return 
 */
template<typename KmerType>
std::string decode_kmer_nt5(KmerType kmer, size_t kmer_size);

template<>
std::string decode_kmer_nt5<uint64_t>(uint64_t kmer, size_t kmer_size);
template<>
std::string decode_kmer_nt5<uint32_t>(uint32_t kmer, size_t kmer_size);
template<>
std::string decode_kmer_nt5<uint16_t>(uint16_t kmer, size_t kmer_size);

/**
 * decode nucleotide kmer to human-readable string
 * @param kmer
 * @param kmer_size
 * @return
 */
template<typename KmerType>
std::string decode_kmer_nt4(KmerType kmer, size_t kmer_size);

template<>
std::string decode_kmer_nt4<uint64_t>(uint64_t kmer, size_t kmer_size);
template<>
std::string decode_kmer_nt4<uint32_t>(uint32_t kmer, size_t kmer_size);
template<>
std::string decode_kmer_nt4<uint16_t>(uint16_t kmer, size_t kmer_size);

/**
 * decode amino acids kmer to human-readable string
 * @param kmer
 * @param kmer_size
 * @return 
 */
template<typename KmerType>
std::string decode_kmer_aa(KmerType kmer, size_t kmer_size);

template<>
std::string decode_kmer_aa<uint64_t>(uint64_t kmer, size_t kmer_size);
template<>
std::string decode_kmer_aa<uint32_t>(uint32_t kmer, size_t kmer_size);
template<>
std::string decode_kmer_aa<uint16_t>(uint16_t kmer, size_t kmer_size);


} // namespace htio2

#endif	/* HTIO2_KMER_H */

