#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
namespace py = pybind11;
using namespace std;

#include "midi_io.h"
#include "dataset/jagged.h"
#include "encoder/encoder_all.h"
#include "enum/model_type.h"
#include "version.h"
#include "feature/feature_internal.h"

#include "sample_fix.h"
#include "feature/feature.h"

#ifndef NOTORCH
#include "sample_internal.h"

// function for python use
vector<string> generate_py(string &status_str, string &piece_str, float temp, int batch_size, bool verbose, map<tuple<int,MODEL_TYPE>,tuple<ENCODER_TYPE,string>> &ckpt_map, int max_steps) {
  midi::Piece piece;
  google::protobuf::util::JsonStringToMessage(piece_str.c_str(), &piece);
  midi::Status status;
  google::protobuf::util::JsonStringToMessage(status_str.c_str(), &status);
  vector<midi::Piece> output = generate(&status, &piece, temp, batch_size, verbose, ckpt_map, max_steps);
  vector<string> output_str;
  for (const auto x : output) {
    string json_string;
    google::protobuf::util::MessageToJsonString(x, &json_string);
    output_str.push_back( json_string );
  }
  return output_str;
}
#else 
void generate_py() {

}
#endif

// test function that modifies in place
void modify(py::list &a) {
  a.append(2);
}

PYBIND11_MODULE(dataset_builder_2,m) {

  m.def("modify", &modify);

  m.def("update_note_density", &update_note_density_py);
  m.def("update_valid_segments", &update_valid_segments_py);
  m.def("select_random_segment", &select_random_segment_py);
  m.def("reorder_tracks", &reorder_tracks_py);
  m.def("prune_tracks", &prune_tracks_py);
  m.def("prune_empty_tracks", &prune_empty_tracks_py);

  m.def("version", &version);
  m.def("getEncoderSize", &getEncoderSize);
  m.def("getEncoderType", &getEncoderType);
  m.def("getEncoder", &getEncoder);

  m.def("gm_inst_to_string", &gm_inst_to_string);
  m.def("prepare_generate", &prepare_generate_py);

  m.def("generate", &generate_py);


  m.def("to_pianoroll", &to_pianoroll_py);
  m.def("piece_to_barrolls", &piece_to_barrolls_py);

  py::class_<Control>(m, "Control")
    .def(py::init<ENCODER*>())
    .def("start", &Control::start)
    .def("is_finished", &Control::is_finished)
    .def("all_finished", &Control::all_finished)
    .def("get_mask", &Control::get_mask)
    .def("increment", &Control::increment)
    .def("check_trigger", &Control::check_trigger);

  //py::class_<Generator>(m, "Generator")
  //  .def(py::init<map<tuple<int,MODEL_TYPE,ENCODER_TYPE>,string> &>())
  //  .def("generate", &Generator::generate_py)
  //  .def("generate_w_midi", &Generator::generate_w_midi_py);

  py::enum_<MODEL_TYPE>(m, "MODEL_TYPE", py::arithmetic())
  .value("TRACK_MODEL", MODEL_TYPE::TRACK_MODEL)
  .value("BAR_INFILL_MODEL", MODEL_TYPE::BAR_INFILL_MODEL)
  .export_values();

  py::class_<Jagged>(m, "Jagged")
    .def(py::init<string &>())
    .def("set_seed", &Jagged::set_seed)
    .def("set_num_bars", &Jagged::set_num_bars)
    .def("set_min_tracks", &Jagged::set_min_tracks)
    .def("set_max_tracks", &Jagged::set_max_tracks)
    .def("set_max_seq_len", &Jagged::set_max_seq_len)
    .def("enable_write", &Jagged::enable_write)
    .def("enable_read", &Jagged::enable_read)
    .def("append", &Jagged::append)
    .def("read", &Jagged::read)
    .def("read_bytes", &Jagged::read_bytes)
    .def("read_json", &Jagged::read_json)
    .def("read_batch", &Jagged::read_batch)
    //.def("read_batch_w_continue", &Jagged::read_batch_w_continue)
    .def("close", &Jagged::close)
    .def("get_size", &Jagged::get_size)
    .def("get_split_size", &Jagged::get_split_size);

  /*
  py::class_<TOKEN>(m, "TOKEN")
    .def(py::init<vector<pair<TOKEN_TYPE,int>>>())
    .def("encode", &TOKEN::encode)
    .def("decode", &TOKEN::decode)
    .def("shift", &TOKEN::shift)
    .def("max_token", &TOKEN::max_token)
    .def("get_domain", &TOKEN::get_domain);
  */

  py::class_<TrainConfig>(m, "TrainConfig")
    .def(py::init<>())
    .def_readwrite("num_bars", &TrainConfig::num_bars)
    .def_readwrite("min_tracks", &TrainConfig::min_tracks)
    .def_readwrite("max_tracks", &TrainConfig::max_tracks)
    .def_readwrite("max_mask_percentage", &TrainConfig::max_mask_percentage)
    .def_readwrite("opz", &TrainConfig::opz)
    .def("to_json", &TrainConfig::to_json)
    .def("from_json", &TrainConfig::from_json);

  py::class_<REPRESENTATION>(m, "REPRESENTATION")
    .def(py::init<vector<pair<TOKEN_TYPE,TOKEN_DOMAIN>>>())
    .def("decode", &REPRESENTATION::decode)
    .def("is_token_type", &REPRESENTATION::is_token_type)
    .def("encode", &REPRESENTATION::encode)
    .def("encode_to_one_hot", &REPRESENTATION::encode_to_one_hot)
    //.def("shift", &REPRESENTATION::shift)
    //.def("get_domain", &REPRESENTATION::get_domain)
    //.def("show", &REPRESENTATION::show)
    .def("pretty", &REPRESENTATION::pretty)
    //.def("where", &REPRESENTATION::where)
    //.def("where_values", &REPRESENTATION::where_values)
    .def_readonly("vocab_size", &REPRESENTATION::vocab_size)
    .def("get_type_mask", &REPRESENTATION::get_type_mask)
    .def("max_token", &REPRESENTATION::max_token);

py::class_<EncoderConfig>(m, "EncoderConfig")
  .def(py::init<>())
  .def_readwrite("both_in_one", &EncoderConfig::both_in_one)
  .def_readwrite("unquantized", &EncoderConfig::unquantized)
  .def_readwrite("interleaved", &EncoderConfig::interleaved)
  .def_readwrite("multi_segment", &EncoderConfig::multi_segment)
  .def_readwrite("te", &EncoderConfig::te)
  .def_readwrite("do_fill", &EncoderConfig::do_fill)
  .def_readwrite("do_multi_fill", &EncoderConfig::do_multi_fill)
  .def_readwrite("do_track_shuffle", &EncoderConfig::do_track_shuffle)
  .def_readwrite("force_instrument", &EncoderConfig::force_instrument)
  .def_readwrite("mark_polyphony", &EncoderConfig::mark_polyphony)
  .def_readwrite("mark_density", &EncoderConfig::mark_density)
  .def_readwrite("mark_time_sigs", &EncoderConfig::mark_time_sigs)
  .def_readwrite("instrument_header", &EncoderConfig::instrument_header)
  .def_readwrite("use_velocity_levels", &EncoderConfig::use_velocity_levels)
  .def_readwrite("genre_header", &EncoderConfig::genre_header)
  .def_readwrite("piece_header", &EncoderConfig::piece_header)
  .def_readwrite("bar_major", &EncoderConfig::bar_major)
  .def_readwrite("force_four_four", &EncoderConfig::force_four_four)
  .def_readwrite("segment_mode", &EncoderConfig::segment_mode)
  .def_readwrite("force_valid", &EncoderConfig::force_valid)
  .def_readwrite("use_drum_offsets", &EncoderConfig::use_drum_offsets)
  .def_readwrite("transpose", &EncoderConfig::transpose)
  .def_readwrite("seed", &EncoderConfig::seed)
  .def_readwrite("segment_idx", &EncoderConfig::segment_idx)
  .def_readwrite("fill_track", &EncoderConfig::fill_track)
  .def_readwrite("fill_bar", &EncoderConfig::fill_bar)
  .def_readwrite("max_tracks", &EncoderConfig::max_tracks)
  .def_readwrite("resolution", &EncoderConfig::resolution)
  .def_readwrite("default_tempo", &EncoderConfig::default_tempo)
  .def_readwrite("num_bars", &EncoderConfig::num_bars)
  .def_readwrite("min_tracks", &EncoderConfig::min_tracks)
  .def_readwrite("fill_percentage", &EncoderConfig::fill_percentage)
  .def_readwrite("multi_fill", &EncoderConfig::multi_fill)
  .def_readwrite("genre_tags", &EncoderConfig::genre_tags);

py::enum_<TOKEN_TYPE>(m, "TOKEN_TYPE", py::arithmetic())
  .value("PIECE_START", TOKEN_TYPE::PIECE_START)
  .value("NOTE_ONSET", TOKEN_TYPE::NOTE_ONSET)
  .value("NOTE_OFFSET", TOKEN_TYPE::NOTE_OFFSET)
  .value("PITCH", TOKEN_TYPE::PITCH)
  .value("NON_PITCH", TOKEN_TYPE::NON_PITCH)
  .value("VELOCITY", TOKEN_TYPE::VELOCITY)
  .value("TIME_DELTA", TOKEN_TYPE::TIME_DELTA)
  .value("INSTRUMENT", TOKEN_TYPE::INSTRUMENT)
  .value("BAR", TOKEN_TYPE::BAR)
  .value("BAR_END", TOKEN_TYPE::BAR_END)
  .value("TRACK", TOKEN_TYPE::TRACK)
  .value("TRACK_END", TOKEN_TYPE::TRACK_END)
  .value("DRUM_TRACK", TOKEN_TYPE::DRUM_TRACK)
  .value("FILL_IN", TOKEN_TYPE::FILL_IN)
  .value("HEADER", TOKEN_TYPE::HEADER)
  .value("VELOCITY_LEVEL", TOKEN_TYPE::VELOCITY_LEVEL)
  .value("GENRE", TOKEN_TYPE::GENRE)
  .value("DENSITY_LEVEL", TOKEN_TYPE::DENSITY_LEVEL)
  .value("TIME_SIGNATURE", TOKEN_TYPE::TIME_SIGNATURE)
  .value("SEGMENT", TOKEN_TYPE::SEGMENT)
  .value("SEGMENT_END", TOKEN_TYPE::SEGMENT_END)
  .value("SEGMENT_FILL_IN", TOKEN_TYPE::SEGMENT_FILL_IN)
  .value("NONE", TOKEN_TYPE::NONE)
  .export_values();

py::enum_<ENCODER_TYPE>(m, "ENCODER_TYPE", py::arithmetic())
  .value("TE_TRACK_DENSITY_ENCODER", ENCODER_TYPE::TE_TRACK_DENSITY_ENCODER)
  .value("TRACK_DENSITY_ENCODER_V2", ENCODER_TYPE::TRACK_DENSITY_ENCODER_V2)
  .value("TRACK_SEGMENT_ENCODER", ENCODER_TYPE::TRACK_SEGMENT_ENCODER)
  .value("TRACK_INTERLEAVED_ENCODER", ENCODER_TYPE::TRACK_INTERLEAVED_ENCODER)
  .value("TRACK_INTERLEAVED_W_HEADER_ENCODER", ENCODER_TYPE::TRACK_INTERLEAVED_W_HEADER_ENCODER)
  .value("TRACK_ENCODER", ENCODER_TYPE::TRACK_ENCODER)
  .value("TRACK_UNQUANTIZED_ENCODER", ENCODER_TYPE::TRACK_UNQUANTIZED_ENCODER)
  .value("TRACK_DENSITY_ENCODER", ENCODER_TYPE::TRACK_DENSITY_ENCODER)
  .value("TRACK_BAR_FILL_DENSITY_ENCODER", ENCODER_TYPE::TRACK_BAR_FILL_DENSITY_ENCODER)
  .value("NO_ENCODER", ENCODER_TYPE::NO_ENCODER)
  .export_values();

py::class_<TeTrackDensityEncoder>(m, "TeTrackDensityEncoder")
  .def(py::init<>())
  .def("encode", &TeTrackDensityEncoder::encode)
  .def("decode", &TeTrackDensityEncoder::decode)
  .def("midi_to_json", &TeTrackDensityEncoder::midi_to_json)
  .def("midi_to_feature", &TeTrackDensityEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TeTrackDensityEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TeTrackDensityEncoder::midi_to_tokens)
  .def("json_to_midi", &TeTrackDensityEncoder::json_to_midi)
  .def("json_to_tokens", &TeTrackDensityEncoder::json_to_tokens)
  .def("tokens_to_json", &TeTrackDensityEncoder::tokens_to_json)
  .def("tokens_to_midi", &TeTrackDensityEncoder::tokens_to_midi)
  .def_readwrite("config", &TeTrackDensityEncoder::config)
  .def_readwrite("rep", &TeTrackDensityEncoder::rep);

py::class_<TrackDensityEncoderV2>(m, "TrackDensityEncoderV2")
  .def(py::init<>())
  .def("encode", &TrackDensityEncoderV2::encode)
  .def("decode", &TrackDensityEncoderV2::decode)
  .def("midi_to_json", &TrackDensityEncoderV2::midi_to_json)
  .def("midi_to_feature", &TrackDensityEncoderV2::midi_to_feature)
  .def("midi_to_json_bytes", &TrackDensityEncoderV2::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackDensityEncoderV2::midi_to_tokens)
  .def("json_to_midi", &TrackDensityEncoderV2::json_to_midi)
  .def("json_to_tokens", &TrackDensityEncoderV2::json_to_tokens)
  .def("tokens_to_json", &TrackDensityEncoderV2::tokens_to_json)
  .def("tokens_to_midi", &TrackDensityEncoderV2::tokens_to_midi)
  .def_readwrite("config", &TrackDensityEncoderV2::config)
  .def_readwrite("rep", &TrackDensityEncoderV2::rep);

py::class_<TrackSegmentEncoder>(m, "TrackSegmentEncoder")
  .def(py::init<>())
  .def("encode", &TrackSegmentEncoder::encode)
  .def("decode", &TrackSegmentEncoder::decode)
  .def("midi_to_json", &TrackSegmentEncoder::midi_to_json)
  .def("midi_to_feature", &TrackSegmentEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TrackSegmentEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackSegmentEncoder::midi_to_tokens)
  .def("json_to_midi", &TrackSegmentEncoder::json_to_midi)
  .def("json_to_tokens", &TrackSegmentEncoder::json_to_tokens)
  .def("tokens_to_json", &TrackSegmentEncoder::tokens_to_json)
  .def("tokens_to_midi", &TrackSegmentEncoder::tokens_to_midi)
  .def_readwrite("config", &TrackSegmentEncoder::config)
  .def_readwrite("rep", &TrackSegmentEncoder::rep);

py::class_<TrackInterleavedEncoder>(m, "TrackInterleavedEncoder")
  .def(py::init<>())
  .def("encode", &TrackInterleavedEncoder::encode)
  .def("decode", &TrackInterleavedEncoder::decode)
  .def("midi_to_json", &TrackInterleavedEncoder::midi_to_json)
  .def("midi_to_feature", &TrackInterleavedEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TrackInterleavedEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackInterleavedEncoder::midi_to_tokens)
  .def("json_to_midi", &TrackInterleavedEncoder::json_to_midi)
  .def("json_to_tokens", &TrackInterleavedEncoder::json_to_tokens)
  .def("tokens_to_json", &TrackInterleavedEncoder::tokens_to_json)
  .def("tokens_to_midi", &TrackInterleavedEncoder::tokens_to_midi)
  .def_readwrite("config", &TrackInterleavedEncoder::config)
  .def_readwrite("rep", &TrackInterleavedEncoder::rep);

py::class_<TrackInterleavedWHeaderEncoder>(m, "TrackInterleavedWHeaderEncoder")
  .def(py::init<>())
  .def("encode", &TrackInterleavedWHeaderEncoder::encode)
  .def("decode", &TrackInterleavedWHeaderEncoder::decode)
  .def("midi_to_json", &TrackInterleavedWHeaderEncoder::midi_to_json)
  .def("midi_to_feature", &TrackInterleavedWHeaderEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TrackInterleavedWHeaderEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackInterleavedWHeaderEncoder::midi_to_tokens)
  .def("json_to_midi", &TrackInterleavedWHeaderEncoder::json_to_midi)
  .def("json_to_tokens", &TrackInterleavedWHeaderEncoder::json_to_tokens)
  .def("tokens_to_json", &TrackInterleavedWHeaderEncoder::tokens_to_json)
  .def("tokens_to_midi", &TrackInterleavedWHeaderEncoder::tokens_to_midi)
  .def_readwrite("config", &TrackInterleavedWHeaderEncoder::config)
  .def_readwrite("rep", &TrackInterleavedWHeaderEncoder::rep);

py::class_<TrackEncoder>(m, "TrackEncoder")
  .def(py::init<>())
  .def("encode", &TrackEncoder::encode)
  .def("decode", &TrackEncoder::decode)
  .def("midi_to_json", &TrackEncoder::midi_to_json)
  .def("midi_to_feature", &TrackEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TrackEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackEncoder::midi_to_tokens)
  .def("json_to_midi", &TrackEncoder::json_to_midi)
  .def("json_to_tokens", &TrackEncoder::json_to_tokens)
  .def("tokens_to_json", &TrackEncoder::tokens_to_json)
  .def("tokens_to_midi", &TrackEncoder::tokens_to_midi)
  .def_readwrite("config", &TrackEncoder::config)
  .def_readwrite("rep", &TrackEncoder::rep);

py::class_<TrackUnquantizedEncoder>(m, "TrackUnquantizedEncoder")
  .def(py::init<>())
  .def("encode", &TrackUnquantizedEncoder::encode)
  .def("decode", &TrackUnquantizedEncoder::decode)
  .def("midi_to_json", &TrackUnquantizedEncoder::midi_to_json)
  .def("midi_to_feature", &TrackUnquantizedEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TrackUnquantizedEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackUnquantizedEncoder::midi_to_tokens)
  .def("json_to_midi", &TrackUnquantizedEncoder::json_to_midi)
  .def("json_to_tokens", &TrackUnquantizedEncoder::json_to_tokens)
  .def("tokens_to_json", &TrackUnquantizedEncoder::tokens_to_json)
  .def("tokens_to_midi", &TrackUnquantizedEncoder::tokens_to_midi)
  .def_readwrite("config", &TrackUnquantizedEncoder::config)
  .def_readwrite("rep", &TrackUnquantizedEncoder::rep);

py::class_<TrackDensityEncoder>(m, "TrackDensityEncoder")
  .def(py::init<>())
  .def("encode", &TrackDensityEncoder::encode)
  .def("decode", &TrackDensityEncoder::decode)
  .def("midi_to_json", &TrackDensityEncoder::midi_to_json)
  .def("midi_to_feature", &TrackDensityEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TrackDensityEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackDensityEncoder::midi_to_tokens)
  .def("json_to_midi", &TrackDensityEncoder::json_to_midi)
  .def("json_to_tokens", &TrackDensityEncoder::json_to_tokens)
  .def("tokens_to_json", &TrackDensityEncoder::tokens_to_json)
  .def("tokens_to_midi", &TrackDensityEncoder::tokens_to_midi)
  .def_readwrite("config", &TrackDensityEncoder::config)
  .def_readwrite("rep", &TrackDensityEncoder::rep);

py::class_<TrackBarFillDensityEncoder>(m, "TrackBarFillDensityEncoder")
  .def(py::init<>())
  .def("encode", &TrackBarFillDensityEncoder::encode)
  .def("decode", &TrackBarFillDensityEncoder::decode)
  .def("midi_to_json", &TrackBarFillDensityEncoder::midi_to_json)
  .def("midi_to_feature", &TrackBarFillDensityEncoder::midi_to_feature)
  .def("midi_to_json_bytes", &TrackBarFillDensityEncoder::midi_to_json_bytes)
  .def("midi_to_tokens", &TrackBarFillDensityEncoder::midi_to_tokens)
  .def("json_to_midi", &TrackBarFillDensityEncoder::json_to_midi)
  .def("json_to_tokens", &TrackBarFillDensityEncoder::json_to_tokens)
  .def("tokens_to_json", &TrackBarFillDensityEncoder::tokens_to_json)
  .def("tokens_to_midi", &TrackBarFillDensityEncoder::tokens_to_midi)
  .def_readwrite("config", &TrackBarFillDensityEncoder::config)
  .def_readwrite("rep", &TrackBarFillDensityEncoder::rep);


  py::class_<feat_density>(m, "feat_density")
    .def(py::init<>());

  m.def("get_feature", static_cast<vector<tuple<int,int,int>>(*)(string&,EncoderConfig*,feat_density)>(&get_feature));


  py::class_<feat_test>(m, "feat_test")
    .def(py::init<>());

  m.def("get_feature", static_cast<vector<double>(*)(string&,EncoderConfig*,feat_test)>(&get_feature));
}