#include "aoc.h"
#include <map>

namespace aoc2022 {

static int64_t pow(size_t n) {
  int64_t x{1};
  for (size_t i = 0; i < n; i++) {
    x *= 5;
  }
  return x;
}

int64_t pow5(size_t n) {
  static std::vector<int64_t> ns;
  if (n >= ns.size()) {
    auto s = ns.size();
    ns.resize(n + 1);
    for (size_t i = s; i < ns.size(); i++) {
      ns[i] = pow(i);
    }
  }
  return ns[n];
}

int64_t base10(snafu s) {
  int64_t x{0};
  for (size_t i = 0; i < s.length; i++) {
    switch (*(s.digits + i)) {
    case '=':
      x = x * 5 - 2;
      break;
    case '-':
      x = x * 5 - 1;
      break;
    case '0':
      x = x * 5;
      break;
    case '1':
      x = x * 5 + 1;
      break;
    case '2':
      x = x * 5 + 2;
      break;
    default:
      break;
    }
  }
  return x;
}

struct snafu_digit {
  int64_t l;
  int64_t m1;
  int64_t d1;
  int64_t d2;
  int64_t m2;
  int64_t h;
};

snafu_digit range(size_t n) {
  static std::map<size_t, snafu_digit> cache;
  auto p = cache.insert({n, {}});
  if (p.second) {
    int64_t m1 = pow5(n);
    int64_t m2 = 2 * m1;
    int64_t d1 = (m1 + m2) / 2;
    int64_t d2 = d1 + 1;

    int64_t l = m1;
    for (size_t i = n; i > 0; i--) {
      l -= 2 * pow5(i - 1);
    }

    int64_t h = m2;
    for (size_t i = n; i > 0; i--) {
      h += 2 * pow5(i - 1);
    }
    p.first->second = {l, m1, d1, d2, m2, h};
  }
  return p.first->second;
}

void decode(int64_t x, size_t n, std::vector<int>& is) {
  if (n == 0) {
    is.push_back(x);
  } else {
    int64_t sign = x == 0 ? 0 : x > 0 ? 1 : -1;
    auto r = range(n);
    auto ax = std::abs(x);
    if (ax < r.l) {
      is.push_back(0);
      decode(x, n - 1, is);
    }
    if (ax >= r.l && ax < r.d2) {
      is.push_back(sign * 1);
      decode(x - sign * r.m1, n - 1, is);
    }
    if (ax >= r.d2) {
      is.push_back(sign * 2);
      decode(x - sign * r.m2, n - 1, is);
    }
  }
}

std::string baseSNAFU(int64_t x) {
  size_t n{0};
  while (true) {
    auto p = range(n);
    if (x >= p.l && x <= p.h) {
      break;
    }
    n++;
  }

  std::vector<int> is;
  decode(x, n, is);
  std::string s;
  char encodings[] = {'0', '1', '2', '=', '-'};
  for (auto& i : is) {
    s.push_back(i >= 0 ? encodings[i] : encodings[i + 5]);
  }
  return s;
}

std::pair<std::string, int64_t> day25(line_view file) {
  std::vector<snafu> ss;
  per_line(file, [&ss](line_view lv) {
    ss.emplace_back(lv);
    return true;
  });

  int64_t x{0};
  for (auto& s : ss) {
    // s.print();
    // printf(" %ld\n", base10(s));
    x += base10(s);
  }

  // printf("%ld %s\n", x, baseSNAFU(x).c_str());
  return {baseSNAFU(x), 0};
}
} // namespace aoc2022