#ifndef __KMER_INTERVAL_MAP_HPP__
#define __KMER_INTERVAL_MAP_HPP__

extern "C" {
#include "bwt.h"
}

#include <fstream>
#include <unordered_map>

#include <jellyfish/mer_dna.hpp>

#include "cereal/archives/binary.hpp"
#include "cereal/types/unordered_map.hpp"

#include "xxhash.h"

using JFMer = jellyfish::mer_dna_ns::mer_base_static<uint64_t, 1>;

// What will be the keys in our k-mer has map
struct KmerKey {
  KmerKey() { mer_.polyT(); }

  KmerKey(uint8_t* seq, uint32_t len) : mer_(len) {
    mer_.polyT();
    for (size_t i = 0; i < len; ++i) {
      mer_.shift_left(seq[i]);
    }
  }

  bool operator==(const KmerKey& ok) const { return mer_ == ok.mer_; }

  // Is there a smarter way to do save / load here?
  template <typename Archive> void save(Archive& archive) const {
    auto key = mer_.get_bits(0, 2 * mer_.k());
    archive(key);
  }

  template <typename Archive> void load(Archive& archive) {
    mer_.polyT();
    uint64_t bits;
    archive(bits);
    mer_.set_bits(0, 2 * mer_.k(), bits);
  }

  JFMer mer_;
};

template <typename Archive> void load(Archive& archive, bwtintv_t& interval) {
  archive(interval.x[0], interval.x[1], interval.x[2], interval.info);
}

template <typename Archive>
void save(Archive& archive, const bwtintv_t& interval) {
  archive(interval.x[0], interval.x[1], interval.x[2], interval.info);
}

/**
 *  This class provides an efficent hash-map from
 *  k-mers to BWT intervals.
 */
class KmerIntervalMap {
public:
  // How we hash the keys
  struct KmerHasher {
    std::size_t operator()(const KmerKey& k) const {
      void* data = static_cast<void*>(const_cast<KmerKey&>(k).mer_.data__());
      return XXH64(data, sizeof(uint64_t), 0);
    }
  };

private:
  std::unordered_map<KmerKey, bwtintv_t, KmerHasher> map_;

public:
  void setK(unsigned int k) { JFMer::k(k); }
  uint32_t k() { return JFMer::k(); }

  bool hasKmer(KmerKey& k) { return map_.find(k) != map_.end(); }

  decltype(map_)::iterator find(const KmerKey& k) { return map_.find(k); }
  decltype(map_)::iterator find(KmerKey&& k) { return map_.find(k); }

  decltype(map_)::iterator end() { return map_.end(); }

  bwtintv_t& operator[](const KmerKey& k) { return map_[k]; }
  bwtintv_t& operator[](KmerKey&& k) { return map_[k]; }

  decltype(map_)::size_type size() { return map_.size(); }

  void save(boost::filesystem::path indexPath) {
    std::ofstream ofs(indexPath.string(), std::ios::binary);
    {
      cereal::BinaryOutputArchive oa(ofs);
      oa(map_);
    }
    ofs.close();
  }

  void load(boost::filesystem::path indexPath) {
    std::ifstream ifs(indexPath.string(), std::ios::binary);
    {
      cereal::BinaryInputArchive ia(ifs);
      ia(map_);
    }
    ifs.close();
  }
};

#endif // __KMER_INTERVAL_MAP_HPP__