#pragma once

#include <vector>
#include <map>
#include <tuple>

#include "../enum/token_types.h"
#include "../enum/constants.h"
// here we implement two class TOKEN and REPRESENTATION

// for code clarity we need to have values for tokens
// rep is simply 
// 1 map<pair<TOKEN_TYPE,int>,int> 
// 2 map<int,pair<TOKEN_TYPE,int>>

class TOKEN_DOMAIN {
public:
  TOKEN_DOMAIN(int n) {
    for (int i=0; i<n; i++) {
      domain.push_back( i );
    }
  }
  TOKEN_DOMAIN(vector<int> values) {
    for (const auto value : values) {
      domain.push_back( value );
    }
  }
  vector<int> domain;
};

class REPRESENTATION {
public:
  REPRESENTATION(vector<pair<TOKEN_TYPE,TOKEN_DOMAIN>> spec) {
    vocab_size = 0;
    for (const auto token_domain : spec) {
      TOKEN_TYPE tt = get<0>(token_domain);
      TOKEN_DOMAIN domain = get<1>(token_domain);
      for (const auto value : domain.domain) {
        forward[make_tuple(tt,value)] = vocab_size;
        backward[vocab_size] = make_tuple(tt,value);
        vocab_size++;
      }
      domains[tt] = domain.domain.size();
      full_domains[tt] = domain.domain;
    }
  }
  int encode(TOKEN_TYPE tt, int value) {
    tuple<TOKEN_TYPE,int> key = make_tuple(tt,value);
    if (forward.find(key) == forward.end()) {
      ostringstream buffer;
      buffer << "token value out of range " << toString(tt) << " = " << value;
      throw std::runtime_error(buffer.str());
    }
    return forward[key];
  }
  int decode(int token) {
    return get<1>(backward[token]);
  }
  int max_token() {
    return vocab_size;
  }
  int get_domain_size(TOKEN_TYPE tt) {
    auto it = domains.find(tt);
    if (it == domains.end()) {
      return 0;
    }
    return it->second;
  }
  bool is_token_type(int token, TOKEN_TYPE tt) {
    return get<0>(backward[token]) == tt;
  }
  vector<int> encode_to_one_hot(TOKEN_TYPE tt, vector<int> values) {
    vector<int> x(vocab_size,0);
    for (const auto value : values) {
      if (value == -1) {
        int domain_size = get_domain_size(tt);
        for (const auto v : full_domains[tt]) {
          x[encode(tt, v)] = 1;
        }
      }
      else {
        x[encode(tt, value)] = 1;
      }
    }
    return x;
  }
  vector<int> decode_to_one_hot(vector<int> tokens) {
    vector<int> one_hot(vocab_size,0);
    for (const auto token : tokens) {
      one_hot[decode(token)] = 1;
    }
    return one_hot;
  }
  vector<int> get_type_mask(vector<TOKEN_TYPE> tts) {
    vector<int> mask(vocab_size,0);
    for (int i=0; i<vocab_size; i++) {
      for (const auto tt : tts) {
        if (is_token_type(i,tt)) {
          mask[i] = 1;
          break;
        }
      }
    }
    return mask;
  }
  string pretty(int token) {
    auto token_value = backward[token];
    return toString(get<0>(token_value)) + string(" = ") + to_string(get<1>(token_value));
  }

  int vocab_size;
  map<tuple<TOKEN_TYPE,int>,int> forward;
  map<int,tuple<TOKEN_TYPE,int>> backward;
  map<TOKEN_TYPE,int> domains;
  map<TOKEN_TYPE,vector<int>> full_domains;
};



/*

class TOKEN {
public:
  TOKEN(vector<pair<TOKEN_TYPE,int>> spec) {
    cprod.push_back( 1 );
    for (size_t i=0; i<spec.size(); i++) {
      tt_2_index[get<0>(spec[i])] = i;
      domain.push_back( get<1>(spec[i]) );
      cprod.push_back( cprod.back() * get<1>(spec[i]) );
    }
  }
  int encode(map<TOKEN_TYPE,int> values) {
    // notifies if you try to add an unknown token
    int value;
    TOKEN_TYPE tt;
    int token = 0;
    for (const auto kv : values) {
      tt = kv.first;
      value = kv.second;
      auto it = tt_2_index.find(tt);
      if (it == tt_2_index.end()) {
        //cout << "WARNING : TRYING TO ADD UNKNOWN TOKEN TYPE" << endl;
      }
      else {
        if ((value < 0) || (value >= domain[it->second])) {
          cout << "ERROR : TRYING TO ENCODE VALUE (" << value << ") OUTSIDE OF DOMAIN (" << toString(tt) << ")" << endl;
          throw(2);
        }
        token += cprod[it->second] * value;
      }
    }
    return token;
  }
  int decode(int token, TOKEN_TYPE tt) {
    auto it = tt_2_index.find(tt);
    assert(it != tt_2_index.end());
    return (token / cprod[it->second]) % domain[it->second];
  }
  int shift(int token, TOKEN_TYPE tt, int shift) {
    auto it = tt_2_index.find(tt);
    assert(it != tt_2_index.end());
    int shifted = token + (shift * cprod[it->second]);
    int dec = decode(token, tt) + shift;
    if ((dec<0)||(dec>=domain[it->second])||(shifted<0)||(shifted>=cprod.back())) {
      throw 2;
    }
    return shifted;
  }
  int max_token() {
    return cprod.back();
  }
  int get_domain(TOKEN_TYPE tt) {
    auto it = tt_2_index.find(tt);
    assert(it != tt_2_index.end());
    return domain[it->second];
  }
  vector<int> domain;
  vector<int> cprod;
  map<TOKEN_TYPE,int> tt_2_index;
};

class REPRESENTATION {
public:
  REPRESENTATION(vector<vector<pair<TOKEN_TYPE,int>>> spec, const char *vname) {
    int count = 0;
    for (size_t i=0; i<spec.size(); i++) {
      for (size_t j=0; j<spec[i].size(); j++) {
        tt_2_index[get<0>(spec[i][j])] = i;
      }
      TOKEN tok(spec[i]);
      toks.push_back(tok);
      starts.push_back(count);
      count += tok.max_token();
      ends.push_back(count);
    }
    velocity_map_name = vname;
  }
  int decode(int token, TOKEN_TYPE tt) {
    auto it = tt_2_index.find(tt);
    assert(it != tt_2_index.end());
    if ((token >= starts[it->second]) && (token < ends[it->second])) {
      int value = toks[it->second].decode(token - starts[it->second], tt);
      if (tt == VELOCITY_LEVEL) {
        //cout << value << endl;
        return velocity_rev_maps[velocity_map_name][value];
      }
      else if (tt == VELOCITY) {
        return velocity_rev_maps["no_velocity"][value];
      }
      return value;
    }
    return -1;
  }
  bool is_token_type(int token, TOKEN_TYPE tt) {
    auto it = tt_2_index.find(tt);
    if ((it != tt_2_index.end()) && (token >= starts[it->second]) && (token < ends[it->second])) {
      return true;
    }
    return false;
  }
  int encode(map<TOKEN_TYPE,int> values) {
    // make sure that tts all belong to same token
    auto it = tt_2_index.find(values.begin()->first);
    assert(it != tt_2_index.end());
    // what is going on here ???
    //auto vit = values.find(VELOCITY);
    //if (vit != values.end()) {
    //  values[VELOCITY] = velocity_maps[velocity_map_name][vit->second];
    //}
    return toks[it->second].encode(values) + starts[it->second];
  }
  vector<int> encode_to_one_hot(map<TOKEN_TYPE,vector<int>> values) {
    auto it = tt_2_index.find(values.begin()->first);
    if (it == tt_2_index.end()) {
      cout << "NO TOKEN TYPE : " << toString(values.begin()->first) << endl;
      assert(false);
    }
    
    vector<int> oh(max_token(),0);
    if (values.begin()->second[0] == -1) {
      for (int i = starts[it->second]; i<ends[it->second]; i++) {
        oh[i] = 1;
      }
    }
    else {
      for (const auto x : values.begin()->second) {
        oh[toks[it->second].encode({{values.begin()->first,x}}) + starts[it->second]] = 1;
      }
    }
    return oh;
  }
  int shift(int token, TOKEN_TYPE tt, int shift) {
    auto it = tt_2_index.find(tt);
    assert(it != tt_2_index.end());
    if ((token >= starts[it->second]) && (token < ends[it->second])) {
      int shifted = toks[it->second].shift(token-starts[it->second], tt, shift);
      if (shifted == -1) { return -1; } // if it fails make sure to return -1
      return shifted + starts[it->second];
    }
    return token; // don't modify
  }
  int get_domain(TOKEN_TYPE tt) {
    auto it = tt_2_index.find(tt);
    assert(it != tt_2_index.end());
    return toks[it->second].get_domain(tt);
  }
  void show(vector<int> tokens) {
    int value;
    for (const auto token : tokens) {
      for (const auto kv : tt_2_index) {
        value = decode(token, kv.first);
        if (value >= 0) {
          cout << toString(kv.first) << "=" << value << " ";
        }
      }
      cout << endl;
    }
  }
  vector<string> pretty(vector<int> tokens) {
    int value;
    vector<string> output;
    for (const auto token : tokens) {
      string ts;
      for (const auto kv : tt_2_index) {
        value = decode(token, kv.first);
        if (value >= 0) {
          string tmp(toString(kv.first));
          ts += tmp + "=" + to_string(value) + " ";
        }
      }
      output.push_back(ts);
    }
    return output;
  }
  vector<int> where(vector<int> tokens, TOKEN_TYPE tt) {
    vector<int> indices;
    for (int i=0; i<tokens.size(); i++) {
      if (is_token_type(tokens[i], tt)) {
        indices.push_back(i);
      }
    }
    return indices;
  }
  vector<int> where_values(vector<int> tokens, TOKEN_TYPE tt) {
    vector<int> values;
    for (const auto i : where(tokens, tt)) {
      values.push_back(decode(tokens[i], tt));
    }
    return values;
  }
  int max_token() {
    return ends.back();
  }
  vector<TOKEN> toks;
  vector<int> starts;
  vector<int> ends;
  map<TOKEN_TYPE,int> tt_2_index;
  string velocity_map_name;
};

*/