import prune_h # this will run it

TOKEN_TYPE = [
  "PIECE_START",
  "NOTE_ONSET",
  "NOTE_OFFSET",
  "PITCH",
  "NON_PITCH",
  "VELOCITY",
  "TIME_DELTA",
  "INSTRUMENT",
  "BAR",
  "BAR_END",
  "TRACK",
  "TRACK_END",
  "DRUM_TRACK",
  "FILL_IN",
  "HEADER",
  "VELOCITY_LEVEL",
  "GENRE",
  "DENSITY_LEVEL",
  "TIME_SIGNATURE",
  "SEGMENT",
  "SEGMENT_END",
  "SEGMENT_FILL_IN",
  "NONE",
]

ENCODER_TYPE = [
  #"TRACK_ENCODER",
  #"TRACK_BAR_MAJOR_ENCODER",
  #"TRACK_GENRE_ENCODER",
  #"TRACK_VELOCITY_ENCODER",
  #"TRACK_VELOCITY_LEVEL_ENCODER",
  #"TRACK_ONE_BAR_FILL_ENCODER",
  #"TRACK_MONO_POLY_ENCODER",
  #"TRACK_INST_HEADER_ENCODER",
  #"TRACK_ONE_TWO_THREE_BAR_FILL_ENCODER",
  #"TRACK_BAR_FILL_ENCODER",
  #"TRACK_MONO_POLY_DENSITY_ENCODER",
  #"SEGMENT_ENCODER",
  #"TRACK_BAR_FILL_SIXTEEN_ENCODER",
  #"TRACK_DENSITY_ENCODER",
  #"TRACK_DENSITY_VELOCITY_ENCODER",
  #"TRACK_BAR_FILL_DENSITY_VELOCITY_ENCODER",

  "TE_TRACK_DENSITY_ENCODER",
  "TRACK_DENSITY_ENCODER_V2",
  #"TE_TRACK_BAR_FILL_DENSITY_ENCODER",
  "TRACK_SEGMENT_ENCODER",

  "TRACK_INTERLEAVED_ENCODER",
  "TRACK_INTERLEAVED_W_HEADER_ENCODER",

  "TRACK_ENCODER",
  "TRACK_UNQUANTIZED_ENCODER",
  "TRACK_DENSITY_ENCODER",
  "TRACK_BAR_FILL_DENSITY_ENCODER",
  "NO_ENCODER"
]

# tuples of type, name, default
ENCODER_CONFIG = [
  ("bool", "both_in_one", "false"),
  ("bool", "unquantized", "false"),
  ("bool", "interleaved", "false"),
  ("bool", "multi_segment", "false"),
  ("bool", "te", "false"),
  ("bool", "do_fill", "false"),
  ("bool", "do_multi_fill", "false"),
  ("bool", "do_track_shuffle", "true"),
  ("bool", "force_instrument", "false"),
  ("bool", "mark_polyphony", "false"),
  ("bool", "mark_density", "false"),
  ("bool", "mark_time_sigs", "false"),
  ("bool", "instrument_header", "false"),
  ("bool", "use_velocity_levels", "false"),
  ("bool", "genre_header", "false"),
  ("bool", "piece_header", "true"),
  ("bool", "bar_major", "false"),
  ("bool", "force_four_four", "false"),
  ("bool", "segment_mode", "false"),
  ("bool", "force_valid", "false"),
  ("bool", "use_drum_offsets", "true"),
  ("int", "transpose", "0"),
  ("int", "seed", "-1"),
  ("int", "segment_idx", "-1"),
  ("int", "fill_track", "-1"),
  ("int", "fill_bar", "-1"),
  ("int", "max_tracks", "12"),
  ("int", "resolution", "12"),
  ("int", "default_tempo", "104"),
  ("int", "num_bars", "4"),
  ("int", "min_tracks", "1"),
  ("float", "fill_percentage", "0"),
  ("set<tuple<int,int>>", "multi_fill", None),
  ("vector<string>", "genre_tags", None)
]

def build_enum(name, values):
  template = """enum {name} {{\n  {values}\n}};\n\n"""
  return template.format(name=name, values=",\n  ".join(values))

def build_py_enum(name, values):
  header = """py::enum_<{name}>(m, "{name}", py::arithmetic())
"""
  footer = """
  .export_values();\n\n"""
  content = []
  for value in values:
    content.append( """  .value("{val}", {name}::{val})""".format(val=value,name=name) )
  return header.format(name=name) + "\n".join(content) + footer

def build_get_encoder(values):
  header = """ENCODER* getEncoder(ENCODER_TYPE et) {
  switch (et) {
"""
  footer = """
    case NO_ENCODER: return NULL;
  }
}\n\n"""
  content = []
  for value in values:
    cname = "".join([w.capitalize() for w in value.split("_")])
    content.append( """    case {}: return new {}();""".format(value, cname) )
  return header + "\n".join(content) + footer

def build_get_encoder_type(values):
  header = """ENCODER_TYPE getEncoderType(string &s) {
"""
  footer = """
  return NO_ENCODER;
}\n\n"""
  content = []
  for value in values:
    content.append( """  if (s == "{}") return {};""".format(value,value) )
  return header + "\n".join(content) + footer

def build_toString(values):
  header = """inline const char* toString(TOKEN_TYPE tt) {
  switch (tt) {
"""
  footer = """
  }
}\n\n"""
  content = []
  for value in values:
    content.append( """    case {}: return "{}";""".format(value,value) )
  return header + "\n".join(content) + footer

def build_py_encoder_class(name):
  template = """py::class_<{name}>(m, "{name}")
  .def(py::init<>())
  .def("encode", &{name}::encode)
  .def("decode", &{name}::decode)
  .def("midi_to_json", &{name}::midi_to_json)
  .def("midi_to_feature", &{name}::midi_to_feature)
  .def("midi_to_json_bytes", &{name}::midi_to_json_bytes)
  .def("midi_to_tokens", &{name}::midi_to_tokens)
  .def("json_to_midi", &{name}::json_to_midi)
  .def("json_to_tokens", &{name}::json_to_tokens)
  .def("tokens_to_json", &{name}::tokens_to_json)
  .def("tokens_to_midi", &{name}::tokens_to_midi)
  .def_readwrite("config", &{name}::config)
  .def_readwrite("rep", &{name}::rep);\n\n"""
  cname = "".join([w.capitalize() for w in name.split("_")])
  return template.format(name=cname)

def build_encoder_config_class():
  template = """class EncoderConfig {{
public:
  EncoderConfig() {{
{init}
  }}
{declare}
}};"""
  init_content = []
  declare_content = []
  for ftype,fname,fdefault in ENCODER_CONFIG:
    if fdefault is not None:
      init_content.append( "    {} = {};".format(fname, fdefault))
    declare_content.append( "  {} {};".format(ftype, fname))
  return template.format(
    init="\n".join(init_content), 
    declare="\n".join(declare_content))

def build_py_encoder_config_class():
  template = """py::class_<EncoderConfig>(m, "EncoderConfig")
  .def(py::init<>())
{content};\n
"""
  content = []
  for ftype,fname,fdefault in ENCODER_CONFIG:
    content.append("""  .def_readwrite("{name}", &EncoderConfig::{name})""".format(name=fname))
  return template.format(content="\n".join(content))

# token_types.h
with open("enum/token_types.h", "w") as f:
  f.write("""#pragma once\n\n""")
  f.write(build_enum("TOKEN_TYPE", TOKEN_TYPE))
  f.write(build_toString(TOKEN_TYPE))

# encoder_types.h
with open("enum/encoder_types.h", "w") as f:
  f.write("""#pragma once\n
#include "../encoder/encoder_all.h"
#include <string>\n\n""")
  f.write(build_enum("ENCODER_TYPE", ENCODER_TYPE))
  f.write(build_get_encoder(ENCODER_TYPE[:-1]))
  f.write(build_get_encoder_type(ENCODER_TYPE[:-1]))
  f.write("""int getEncoderSize(ENCODER_TYPE et) {
  ENCODER *encoder = getEncoder(et);
  if (!encoder) {
    return 0;
  }
  int size = encoder->rep->max_token();
  delete encoder;
  return size;
}""")

# encoder_config.h
with open("enum/encoder_config.h", "w") as f:
  f.write("""#pragma once
  
#include <vector>
#include <tuple>
#include <map>
""")
  f.write(build_encoder_config_class())

# stuff for lib.cpp
header = """#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);

"""
footer = """}"""

from feature import build_features

with open("lib.cpp", "w") as f:
  f.write(header)
  f.write(build_py_encoder_config_class())
  f.write(build_py_enum("TOKEN_TYPE", TOKEN_TYPE))
  f.write(build_py_enum("ENCODER_TYPE", ENCODER_TYPE))
  for enc in ENCODER_TYPE[:-1]:
    f.write(build_py_encoder_class(enc))
  f.write(build_features())
  f.write(footer)