#include <iostream>
#include <vector>
#include <tuple>
#include <map>
#include <set>

#include <google/protobuf/util/json_util.h>

#include "lz4.h"
#include "../protobuf/midi.pb.h"
#include "../encoder/encoder_all.h"
#include "../enum/encoder_types.h"
#include "../enum/train_config.h"
#include "../random.h"

using namespace std;

class Jagged {
public:
  Jagged(string filepath_) {
    filepath = filepath_;
    header_filepath = filepath_ + ".header";
    can_write = false;
    can_read = false;
    flush_count = 0;
    num_bars = 4;
    min_tracks = 2;
    max_tracks = 12;
    max_seq_len = 2048;

    engine.seed(time(NULL));

    encoder = NULL;
  }

  void set_seed(int seed) {
    srand(seed); // set the seed
    engine.seed(seed);
  }

  void set_num_bars(int x) {
    num_bars = x;
  }

  void set_min_tracks(int x) {
    min_tracks = x;
  }

  void set_max_tracks(int x) {
    max_tracks = x;
  }

  void set_max_seq_len(int x) {
    max_seq_len = x;
  }

  void enable_write() {
    assert(can_read == false);
    if (can_write) { return; }
    // check that the current file is empty unless force flag is present ?
    fs.open(filepath, ios::out | ios::binary);
    can_write = true;
  }
  
  void enable_read() {
    assert(can_write == false);
    if (can_read) { return; }
    fs.open(filepath, ios::in | ios::binary);
    if (!fs.is_open()) {
      throw runtime_error("COULD NOT OPEN FILE!");
    }
    header_fs.open(header_filepath, ios::in | ios::binary);
    header.ParseFromIstream(&header_fs);
    can_read = true;
  }

  void append(string &s, size_t split_id) {
    enable_write();

    size_t start = fs.tellp();
    // begin compress ===============================
    size_t src_size = sizeof(char)*s.size();
    size_t dst_capacity = LZ4_compressBound(src_size);
    char* dst = new char[dst_capacity];
    size_t dst_size = LZ4_compress_default(
      (char*)s.c_str(), dst, src_size, dst_capacity);
    fs.write(dst, dst_size);
    delete[] dst;
    // end compress =================================
    size_t end = fs.tellp();
    midi::Dataset::Item *item;
    switch (split_id) {
      case 0: item = header.add_train(); break;
      case 1: item = header.add_valid(); break;
      case 2: item = header.add_test(); break;
    }
    item->set_start(start);
    item->set_end(end);
    item->set_src_size(src_size);
    flush_count++;

    if (flush_count >= 1000) {
      flush();
      flush_count = 0;
    }
  }

  string read(size_t index, size_t split_id) {
    enable_read();

    midi::Dataset::Item item;
    switch (split_id) {
      case 0: item = header.train(index); break;
      case 1: item = header.valid(index); break;
      case 2: item = header.test(index); break;
    }
    size_t csize = item.end() - item.start();
    char* src = new char[csize/sizeof(char)];
    fs.seekg(item.start());
    fs.read(src, csize);
    string x(item.src_size(), ' ');
    LZ4_decompress_safe(src,(char*)x.c_str(),csize,item.src_size());
    delete[] src;
    return x;
  }

  py::bytes read_bytes(size_t index, size_t split_id) {
    return py::bytes(read(index, split_id));
  }

  string read_json(size_t index, size_t split_id) {
    midi::Piece p;
    string serialized_data = read(index, split_id);
    p.ParseFromString(serialized_data);
    string json_string;
    google::protobuf::util::MessageToJsonString(p, &json_string);
    return json_string;
  }

  tuple<vector<vector<int>>,vector<vector<int>>> read_batch(int batch_size, size_t split_id, ENCODER_TYPE et, TrainConfig *tc) {
    enable_read();
    midi::Piece x;
    int index;
    int nitems = get_split_size(split_id);
    int maxlen = 0;

    ENCODER* enc = getEncoder(et);
    if (!enc) {
      throw runtime_error("ENCODER TYPE DOES NOT EXIST");
    }

    vector<vector<int>> mask;
    vector<vector<int>> batch;

    while (batch.size() < batch_size) {

      try {

        //index = rand() % nitems;
        index = random_on_range(nitems, &engine);
        string serialized_data = read(index, split_id);
        x.ParseFromString(serialized_data);

        // pick random segment
        select_random_segment(&x, tc->num_bars, tc->min_tracks, tc->max_tracks, enc->config->te, &engine);

        // pick random transpose
        tuple<int,int> pitch_ext = get_pitch_extents(&x);
        //cout << get<0>(pitch_ext) << " -- " << get<1>(pitch_ext) << endl;
        vector<int> choices;
        for (int tr=-6; tr<6; tr++) {
          if ((get<0>(pitch_ext)+tr >= 0) && (get<1>(pitch_ext)+tr < 128)) {
            choices.push_back( tr );
          }
        }
        //enc->config->transpose = choices[rand() % choices.size()];
        enc->config->transpose = choices[random_on_range(choices.size(),&engine)];

        // 75 % of the time we do bar infill
        // the old code actually only 25 %
        if (enc->config->both_in_one) {
          //enc->config->do_multi_fill = (rand() % 1000 > 750);
          enc->config->do_multi_fill = random_on_unit(&engine) < .75;
        }

        // pick bars for infilling
        if (enc->config->do_multi_fill) {
          enc->config->multi_fill = make_bar_mask(
            &x, tc->max_mask_percentage, &engine);
        }

        // encode tokens
        vector<int> raw_tokens = enc->encode(&x);

        // reject partially full sequences
        if (et == TRACK_SEGMENT_ENCODER) {
          if (raw_tokens.size() < max_seq_len) {
            throw runtime_error("TOKEN SEQ NOT LONG ENOUGH");
          }
          raw_tokens.resize(max_seq_len);
        }

        vector<int> tokens;
        if (raw_tokens.size() > max_seq_len) {
          // pick a random section
          //int offset = rand() % (raw_tokens.size() - max_seq_len + 1);
          int offset = random_on_range(
            (int)raw_tokens.size() - max_seq_len + 1, &engine);
          copy(
            raw_tokens.begin() + offset, 
            raw_tokens.begin() + offset + max_seq_len,
            back_inserter(tokens));
        }
        else {
          copy(raw_tokens.begin(), raw_tokens.end(), back_inserter(tokens));
        }
        batch.push_back( tokens );
        mask.push_back( vector<int>(tokens.size(),1) );
        maxlen = max((int)tokens.size(), maxlen);
      }
      //catch(...) { /*do nothing*/ }
      catch (const std::exception &exc)
      {
          //std::cerr << exc.what() << std::endl;
      }

    }
    // right pad the sequences
    for (int i=0; i<batch_size; i++) {
      batch[i].insert(batch[i].end(), maxlen-batch[i].size(), 0);
      mask[i].insert(mask[i].end(), maxlen-mask[i].size(), 0);
      assert(mask[i].size() == maxlen);
      assert(batch[i].size() == maxlen);
    }
    return make_pair(batch,mask);
  }

  int get_size() {
    enable_read();
    return header.train_size() + header.valid_size() + header.test_size();
  }

  int get_split_size(int split_id) {
    enable_read();
    switch (split_id) {
      case 0 : return header.train_size();
      case 1 : return header.valid_size();
      case 2 : return header.test_size();
    }
    return 0; // invalid split id
  }

  void flush() {
    fs.flush();
    header_fs.open(header_filepath, ios::out | ios::binary);
    if (!header.SerializeToOstream(&header_fs)) {
      cerr << "ERROR : Failed to write header file" << endl;
    }
    header_fs.close();
  }

  void close() {
    flush();
    fs.close();
    header_fs.close();
    can_read = false;
    can_write = false;
  }
  
private:
  string filepath;
  string header_filepath;
  fstream fs;
  fstream header_fs;
  bool can_write;
  bool can_read;
  midi::Dataset header;
  int flush_count;

  int num_bars;
  int min_tracks;
  int max_tracks;
  int max_seq_len;

  std::mt19937 engine;

  vector<vector<int>> bstore;
  ENCODER *encoder;
};



