#include "hyper.h"

namespace hr {

EX namespace wfc {

typedef map<vector<eWall>, short> wfc_data;

wfc_data probs;

typedef pair<const vector<eWall>, short> probdata;

wfc_data gen_decompressed(const wfc_data& data) {
  wfc_data res;
  for(auto& d: data) {
    auto cur = d.first;
    for(int i=0; i<isize(cur)-1; i++) {
      for(int k=2; k<isize(cur); k++) swap(cur[k-1], cur[k]);
      res[cur] += d.second / (isize(cur)-1);
      }
    }
  return res;
  }

wfc_data gen_compressed(const wfc_data& data) {
  wfc_data res;
  for(auto& d: data) {
    vector<eWall> cur = d.first;
    vector<eWall> best = d.first;
    for(int i=0; i<isize(best)-1; i++) {
      for(int k=2; k<isize(cur); k++) swap(cur[k-1], cur[k]);
      if(cur < best) best = cur;
      }
    res[best] += d.second;
    }
  return res;
  }

void hwrite(hstream& hs, wfc_data& data) {
  int siz = isize(data);
  hr::hwrite(hs, siz);
  for(auto p: data) {
    hr::hwrite<char>(hs, isize(p.first));
    for(auto c: p.first) hr::hwrite<char>(hs, c);
    hr::hwrite<short>(hs, p.second);
    }
  }

void hread(hstream& hs, wfc_data& data) {
  data.clear();
  int qty = hs.get<int>();
  for(int i=0; i<qty; i++) {
    vector<eWall> w;
    w.resize(hs.get<char>());
    for(auto& c: w) c = eWall(hs.get<char>());
    data[w] = hs.get<short>();
    }
  }

EX void load_probs() {
  start_game();
  manual_celllister cl;
  cl.add(cwt.at);
  for(int i=0; i<isize(cl.lst); i++) {
    cell *c = cl.lst[i];
    if(c->wall == waChasm || c->mpdist > 7) goto next;
    forCellEx(c1, c) if(c1->wall == waChasm || c1->mpdist > 7) goto next;
    
    for(int i=0; i<c->type; i++) {
      vector<eWall> s;
      if(c->type == 6 && (i&1)) continue;
      s.push_back(c->wall);
      for(int j=0; j<c->type; j++) s.push_back(c->modmove(i+j)->wall);
      bool interesting = false;
      for(eWall w: s) if(w != waNone) interesting = true;
      if(interesting) probs[s] += 42 / c->type;
      }
    
    forCellEx(c1, c) cl.add(c1);
    next: ;
    }
  }

bool agree(cell *c, probdata& p) {
  // println(hlog, p);
  if(isize(p.first) != c->type + 1) return false;
  int idx = 1;
  if(c->wall != waChasm && eWall(c->wparam) != p.first[0]) return false;
  forCellEx(c1, c) {
    eWall found = p.first[idx++];
    if(c1->wparam != found && c1->wall != waChasm) return false;
    }
  return true;
  }

vector<probdata*> gen_picks(cell *c, int& total, wfc_data& data) {
  vector<probdata*> picks;
  total = 0;
  
  for(auto& wp: data) {
    if(!agree(c, wp)) continue;
    picks.push_back(&wp);
    total += wp.second;
    }
  
  return picks;
  }

EX vector<cell*> centers;

EX void schedule(cell *c) {
  centers.push_back(c);
  }

void generate_all() {
  start_game();
  manual_celllister cl;
  cl.add(cwt.at);
  vector<cell*> centers;
  for(int i=0; i<isize(cl.lst); i++) {
    cell *c = cl.lst[i];
    bool incenter = true;
    forCellEx(c1, c) {
      if(c1->mpdist <= 7) 
        cl.add(c1); 
      else 
        incenter = false;
      }
    if(incenter) 
      centers.push_back(c);
    }
  
  for(cell *c: cl.lst) c->wall = waChasm;
  }

string eclectic_c = string("\x78\xda\x8d\x99\x3b\x72\xa4\x30\x10\x86\x51\x24\xd2\x0d\x60\x12\xc5\x1b\x08\x12\x09\x88\x89\x18\xa5\x5b\x2e\x5f\xc2\xc9\x1c\xc1\x97\xf0\x25\x7c\x85\x0d\xf6\x12\x3e\x8f\x83\x9d\x31\x30\xe8\xd1\xfd\x37\x53\xa5\x82\x5a\x7f\xfb\x77\xeb\xd1\x4d\xd3\xfc\x52\x55\x55\x57\xf8\x77\x47\xaa\x4f\x91\x6a\xef\xc3\x8a\xd4\xe5\x14\x65\x4e\x51\xf6\x3e\xbe\x45\xaa\x3b\xa5\xe5\x4e\xcd\x71\x3a\x45\x2d\xa7\x2c\x5e\x4f\x69\x85\xfb\xf8\x12\xa9\xdb\x0f\xa5\x85\x3d\xfc\x07\x75\xd4\x36\x2c\xd4\x69\x45\xaf\xdb\x6d\x7c\x43\x9d\xc7\x39\x78\x83\x84\xf9\x21\x6a\x81\x30\xa2\xc7\x56\xd4\xb1\xcf\x53\xa2\x85\x73\xf4\x01\x75\xba\x4d\xcb\x8a\xd4\x63\xbc\x8a\x54\xff\x43\x21\x9f\x7a\xd1\xa7\xfe\x54\x04\xf4\xa7\xac\x3d\x62\xe4\xaf\xaa\x05\xe2\x31\xfe\x28\xa4\xe3\xef\xa3\x81\xfe\xf8\x6d\xe0\x1d\x19\x45\x8f\x1f\xf1\x3a\x43\x62\x11\xbd\xbd\x8a\x27\x55\x8e\xd1\xb0\x0d\xec\xed\xcb\xcf\x6e\x6a\x21\xce\x71\xf4\xdd\xb6\xc1\xaf\xdc\x1e\xe7\xbf\xd9\x9d\x54\xa7\xb2\xc1\x1e\xe7\xef\xac\x3f\x6d\x34\xf8\x99\x5f\x9e\xcf\x0e\x2d\xc4\x39\x1f\x31\xe6\x54\x36\x30\xdb\x89\x68\x58\x42\xce\x04\x7b\x8c\x37\x90\xe8\xe0\xb9\xdb\x63\x9b\x8f\x81\x2e\x8a\x49\xde\x97\x3d\xb6\x79\x9d\x3e\xc9\x4a\x5a\x8c\xfe\x5a\x20\x70\x26\xe9\x13\xaf\x6b\x90\x23\x8e\x4a\x40\x0b\x99\xe4\x03\xea\xb8\x6d\x47\xb1\x8e\x87\xfb\xe5\xb6\x4c\x82\x35\xf0\xb9\xf1\x1b\x85\x09\x9c\xf7\xc6\xe4\x0c\x23\x6a\x9f\x39\xef\xf1\x28\x7a\x3c\x3d\x3d\xae\x05\x02\x5b\x9a\x36\x6b\xbc\xa5\x65\xcb\x7d\x33\x6b\xe9\x1a\x65\x49\x2b\x52\x38\xbb\x85\xcd\xe2\x0c\x89\x00\x67\x1e\x22\x1d\x2b\x52\x01\x46\xce\x4d\xac\x09\x55\x51\x39\x6a\x40\x29\xd6\x73\x15\x51\x7c\xbc\xc7\x39\x99\xd6\x31\xa7\xaa\xab\x38\xe3\xd2\x7b\x6f\x93\x4c\xa7\x41\xa6\xb4\x22\xc1\x67\xb9\xfe\x54\x96\x8b\x33\x13\x6d\xc9\x45\x91\xc3\x13\xfb\xa0\x9f\x66\x3e\x22\x68\x8d\x21\x21\x6a\x86\x18\xc4\xc8\x1b\x23\x9f\x68\x4b\x4b\x74\x4a\x31\x71\x63\x7d\x59\x8a\x9a\xa2\x06\x71\x59\x41\x2a\x64\xff\x83\x9e\x57\x10\x33\x41\xc8\xa2\x8f\xd7\x59\x60\x9c\x87\x6c\xd0\x3a\x71\xf4\x36\x90\xb8\x92\x84\x8a\xf6\x69\x60\x89\x38\xb6\xcb\x2c\x99\x13\x74\x55\x16\xc7\xf5\x85\xb4\xb4\xff\xa5\x62\xcf\x9f\x21\xae\xa5\x25\x39\x3f\x98\x8c\x6a\x20\xc1\xfb\x62\x44\x5f\x6c\xb6\x2e\xa5\xce\x9e\x85\x70\x86\xcd\x89\xd2\x63\x17\xe5\x08\x43\x5a\x8a\xf3\xc3\xf1\xfe\x82\xa8\x3d\xd6\x79\xca\x47\x4f\xd9\x92\x9a\x88\xbc\x53\xce\x6d\x8a\x3c\x57\xe4\xdc\x96\x2c\x2e\x68\x22\x7d\x2a\x6a\x90\x05\x68\x8d\x10\x45\xef\x95\x5c\xc1\x20\xc6\xa5\x22\xae\xe9\xba\x50\x84\x22\x75\x72\xe2\x9d\x24\x14\x9b\x69\xa9\x68\xe1\x89\x7d\xcc\x90\xa0\xad\xe4\x54\x9a\x23\x0c\x11\x4f\xae\x58\x97\x9d\x88\xeb\xc7\x32\xa6\x5c\x76\xd2\xcb\x1a\x32\x3f\x99\x63\xb1\x93\x79\x0c\x54\x64\x8d\x1e\x57\xef\x74\x1d\xe1\x33\x5f\x46\x48\xd0\x35\xba\xcf\x3c\xc5\xc4\x44\xce\xc7\x13\xbe\x94\xf3\x99\x88\x99\xa7\x3a\x4b\x16\x03\xe5\x09\x57\xcc\x38\x76\xb2\x05\xe3\xd0\x71\x99\xd7\xa5\x3f\x01\x8c\xd5\x9a\x12\x3a\x85\xab\x35\x05\x88\x86\xd5\x89\x23\xef\x8b\xa4\x54\x76\xff\x45\x5a\x8b\x75\x3e\xa0\x8e\x62\xbd\xce\xef\xdf\x48\xa2\xac\xbe\x5b\xa6\xb6\xae\x92\xdd\x68\x01\xf1\x4a\xea\xb4\xd9\x4e\x5b\x92\xba\x14\x7d\x77\x0d\xa8\xf6\x79\xea\x5b\xc2\xd2\x7e\x3f\xb3\x44\x0b\x35\x5a\xd6\x0a\x75\x5d\x57\xef\x22\x7e\x47\xd0\x80\x3a\x74\x8c\xf8\x0d\x41\x03\xea\x78\xd3\x36\xe0\x3d\x68\x82\x3a\x26\xa9\x07\x0c\xa3\xb1\xdf\x7f\x92\x3a\x71\x7e\x9f\x0b\xc2\x88\x1e\x1b\xb6\x27\x6e\xc4\xaa\xcc\xc0\x1a\x47\x33\xcf\xb9\x34\xf3\xdb\x53\xdd\x34\x0b\x7a\x5d\x6f\x2c\x91\xbe\x1f\x76\xa0\xf3\xb4\xfa\xd2\x09\xbd\x29\x4b\x52\xb6\xb0\xa8\x01\xc5\x5b\xdb\xdf\x77\xd3\x77\xf0\x4e\x7c\x1b\xed\x80\x2f\xa5\x25\x9b\x55\xad\xb4\xc7\x1d\x9c\x93\xcd\xe6\x73\x58\xea\x85\xef\x04\xab\x4e\x0f\x3a\x96\x0d\x49\xf4\x89\x25\x44\xd0\x1a\x5d\xb1\xdf\x14\xd1\x89\x44\x6a\x25\x9f\x71\x4f\x9e\x19\xcd\xfc\xa5\x12\x09\x7a\xc6\x1d\x39\xe3\x3a\xe2\x7b\x66\xac\xe7\xd7\x09\xdf\x60\xd6\x5a\xd6\x01\x62\x66\x89\xf4\x4d\xc2\xc1\x6e\x27\xe5\x8b\x89\xfe\x8d\xd7\x49\x89\x9a\x21\x86\xa4\x3f\x4b\x51\x53\x56\xf3\xf0\x94\x8b\x3a\x2c\xa5\x4f\xf9\x3d\x4f\xf8\xe7\x0a\xd6\x02\x41\x7b\x9d\x53\x13\xb4\x36\xb0\x6b\xed\x88\x15\xa0\x75\x4c\xf1\x0e\xe3\x98\x7a\x9d\x27\x7c\x72\x76\x6a\x40\x78\xe8\x4b\xd9\x23\x73\x44\xa7\x9d\xf6\x96\xba\xd2\xc4\x58\x74\xfe\x39\x62\x60\x89\x29\xd1\xa8\x33\xc2\x11\x6f\x4c\xe9\x8c\x3d\xd1\x13\x58\x75\x3c\x11\x4f\x65\xe4\x79\xe6\xc9\x7c\x9c\x72\x0f\xde\xdc\x3c\xa9\x93\xdf\x37\x2c\xe1\x92\xbd\x46\xc4\xe1\xf1\xc0\xc4\xd3\xb1\xc6\x03\xe3\xed\xf0\x5c\x9b\x11\x64\xb6\xe9\xd9\x79\x1d\x41\xde\x6a\x0a\x82\xae\x6b\x26\x90\x1d\x31\x71\x7c\x2d\x9a\xc4\xfc\x39\x31\xb5\xde\x11\xd3\x4b\x85\x7f\x56\xa4\x8e\x2f\xcb\x0b\x20\x66\x56\x27\xed\xba\xea\xac\x9b\x93\xff\xd6\x27\x47\x00\x44\x03\x89\xc3\x97\x20\xf4\xad\x2d\x4b\x85\xa2\x9b\x17\xc4\xfe\x77\x00\x1a\x98\x38\xd6\x37\x80\xae\x59\xea\xc7\x42\x68\xa4\x56\xa8\xeb\xba\x2e\x2f\xe2\x89\xd0\x51\xe7\x99\xdf\xa3\x1b\xf8\x26\x65\xab\xff\xa2\xf7\xca\x99", 1190);

void extend(wfc_data& res, vector<eWall> v, int val, int sides) {
  if(isize(v) == sides+1)
    res[v] += val;
  else if(isize(v) < sides + 1) {
    for(int i=1; i<isize(v); i++) {
      auto w = v;
      w.insert(w.begin()+i, v[i]);
      extend(res, w, val, sides);
      }
    v.insert(v.begin()+1, v.back());
    extend(res, v, val, sides);
    }
  else {
    for(int i=1; i<isize(v); i++) {      
      auto w = v;
      auto l = v[i==1 ? isize(v)-1 : i-1];
      auto r = v[i == isize(v)-1 ? 1 : i+1];
      if(l == waCharged && r == waGrounded) continue;
      if(l == waGrounded && r == waCharged) continue;
      w.erase(w.begin()+i);
      extend(res, w, val, sides);
      }
    }
  }

wfc_data renumerate(const wfc_data& d, int sides) {
  wfc_data res;
  for(auto& p: d) {
    extend(res, p.first, p.second, sides);
    }
  return res;
  }

wfc_data& eclectic_data() {
  static wfc_data d = gen_decompressed(deserialize<wfc_data>(decompress_string(eclectic_c)));
  if(geometry == gNormal && !PURE) return d;
  if(geometry == gNormal && PURE) {
    static wfc_data dpure = renumerate(d, 7);
    return dpure;
    }
  if(geometry == gEuclid) {
    static wfc_data deuc = renumerate(d, 6);
    return deuc;
    }
  return d;
  }

EX bool use_eclectic = true;

EX void invoke() {
  wfc_data& d = use_eclectic ? eclectic_data() : probs;

  while(isize(centers)) {
    int pos = -1;
    ld best_entropy = 1e9;
    for(int p=0; p<isize(centers); p++) {
      cell *c = centers[p];
      int total;
      
      auto picks = gen_picks(c, total, d);
      
      ld entropy = 0;
      for(auto p: picks) entropy += p->second * log(total * 1. / p->second) / total;
      
      if(entropy < best_entropy) best_entropy = entropy, pos = p;
      }
    
    cell *c = centers[pos];
    centers[pos] = centers.back();
    centers.pop_back();

    // println(hlog, "chosen ", c, " at entropy ", best_entropy, " in distance ", c->mpdist);

    int total;
    auto picks = gen_picks(c, total, d);

    if(total) total = hrand(total);
    for(auto pp: picks) {
      auto& p = *pp;
      total -= p.second;
      if(total < 0) {
        int idx = 1;
        c->wall = p.first[0];
        c->wparam = p.first[0];
        forCellEx(c1, c) c1->wparam = c1->wall = p.first[idx++];
        break;
        }
      }

    }

  }

void use_probs() {
  stop_game();
  use_eclectic = false;
  firstland = specialland = laEclectic;
  start_game();
  }

EX void wfc_menu() {
  cmode = sm::SIDE | sm::MAYDARK;
  gamescreen(1);
  dialog::init(XLAT("Wave Function Collapse"));
  
  dialog::addSelItem(XLAT("import the current map"), XLAT("rules: %1", its(isize(probs))), 'i');
  dialog::add_action(load_probs);

  dialog::addBoolItem(XLAT("generate a map using WFC"), !use_eclectic, 'g');
  if(probs.empty()) dialog::lastItem().value = "(no rules)";
  
  dialog::add_action([] {
    if(probs.empty()) return;
    if(use_eclectic) {
      use_probs();
      }
    else {
      stop_game();
      use_eclectic = true;
      start_game();
      }
    });
  
  dialog::addBack();
  dialog::display();
  }

auto wfc_hook = 
  // addHook(hooks_handleKey, 100, wfc_handleKey) +
  addHook(hooks_args, 100, [] {
    using namespace arg;

    if(0) ;
    else if(argis("-wfc-load")) {
      load_probs();
      }
    else if(argis("-wfc-c")) {
      println(hlog, as_cstring(compress_string(serialize(gen_compressed(probs)))));
      }
    else if(argis("-wfc-use")) {
      use_probs();
      }
      
    else return 1;
    return 0;
    });

EX }

}