#pragma once

#include "representation.h"
#include "../protobuf/midi.pb.h"
#include "../enum/encoder_config.h"
#include "../enum/train_config.h"
#include "../midi_io.h"
#include "../feature/feature.h"

class ENCODER {
public:

  virtual vector<int> encode(midi::Piece *p) {}
  virtual void decode(vector<int> &tokens, midi::Piece *p) {}
  virtual EncoderConfig *get_encoder_config() {} // don't do virtual just yet

  string midi_to_json(string &filepath) {
    midi::Piece p;
    parse_new(filepath, &p, config);
    string json_string;
    google::protobuf::util::MessageToJsonString(p, &json_string);
    return json_string;
  }

  vector<double> midi_to_feature(string &filepath) {
    midi::Piece p;
    parse_new(filepath, &p, config);
    return piece_to_quantize_info(&p);
  }
  
  #ifdef PYBIND
  py::bytes midi_to_json_bytes(string &filepath, TrainConfig *tc) {
    string x;
    midi::Piece p;
    parse_new(filepath, &p, config, NULL);
    update_valid_segments(&p, tc->num_bars, tc->min_tracks, config->te);
    if (!p.valid_segments_size()) {
      return py::bytes(x); // empty bytes
    }
    p.SerializeToString(&x);
    return py::bytes(x);
  }
  #endif
  

  vector<int> midi_to_tokens(string &filepath) {
    midi::Piece p;
    parse_new(filepath, &p, config);
    return encode(&p);
  }

  void json_to_midi(string &json_string, string &filepath) {
    midi::Piece p;
    google::protobuf::util::JsonStringToMessage(json_string.c_str(), &p);
    write_midi(&p, filepath);
  }

  vector<int> json_to_tokens(string &json_string) {
    midi::Piece p;
    google::protobuf::util::JsonStringToMessage(json_string.c_str(), &p);
    return encode(&p); 
  }

  string tokens_to_json(vector<int> &tokens) {
    midi::Piece p;
    decode(tokens, &p);
    string json_string;
    google::protobuf::util::MessageToJsonString(p, &json_string);
    return json_string;
  }

  void tokens_to_json_array(vector<vector<int>> &seqs, vector<midi::Piece> &output) {
    for (int i=0; i<seqs.size(); i++) {
      decode(seqs[i], &(output[i]));
    }
  }

  void tokens_to_midi(vector<int> &tokens, string &filepath) {
    midi::Piece p;
    decode(tokens, &p);
    write_midi(&p, filepath);
  }

  EncoderConfig *config;
  REPRESENTATION *rep;
};


vector<int> resolve_bar_infill_tokens(vector<int> &raw_tokens, REPRESENTATION *rep) {
  int fill_pholder = rep->encode(FILL_IN, 0);
  int fill_start = rep->encode(FILL_IN, 1);
  int fill_end = rep->encode(FILL_IN, 2);

  vector<int> tokens;

  auto start_pholder = raw_tokens.begin();
  auto start_fill = raw_tokens.begin();
  auto end_fill = raw_tokens.begin();

  while (start_pholder != raw_tokens.end()) {
    start_pholder = next(start_pholder); // FIRST TOKEN IS PIECE_START ANYWAYS
    auto last_start_pholder = start_pholder;
    start_pholder = find(start_pholder, raw_tokens.end(), fill_pholder);
    if (start_pholder != raw_tokens.end()) {
      start_fill = find(next(start_fill), raw_tokens.end(), fill_start);
      end_fill = find(next(end_fill), raw_tokens.end(), fill_end);

      // insert from last_start_pholder --> start_pholder
      tokens.insert(tokens.end(), last_start_pholder, start_pholder);
      tokens.insert(tokens.end(), next(start_fill), end_fill);
    }
    else {
      // insert from last_start_pholder --> end of sequence (excluding fill)
      start_fill = find(raw_tokens.begin(), raw_tokens.end(), fill_start);
      tokens.insert(tokens.end(), last_start_pholder, start_fill);
    }
  }
  return tokens;
}