#include "aoc.h"
#include <algorithm>
#include <deque>
#include <iostream>
#include <set>
#include <vector>

namespace aoc2022 {
struct valve_m {
  valve* v;
  int m;

  friend bool operator<(valve_m v1, valve_m v2) { return v1.m < v2.m ? true : v1.m > v2.m ? false : v1.v < v2.v; }
  friend bool operator>(valve_m v1, valve_m v2) { return v1.m > v2.m ? true : v1.m < v2.m ? false : v1.v > v2.v; }
};

std::map<line_view, valve*> valves = {};
std::map<valve*, std::vector<valve_m>> diagram;

int maxopened = 0;

static valve* get(line_view lv) {
  auto p = valves.insert({lv, nullptr});
  return p.first->second;
}

// bfs Dijkstra’s algorithm
// You estimate it will take you one minute to open a single valve
// and one minute to follow any tunnel from one valve to another.
void rank_neighbours(valve* v, std::map<valve*, int>& visited) {
  std::deque<valve*> q;
  q.push_back(v);
  visited.insert({v, 0});

  while (!q.empty()) {
    auto s = q.size();
    while (s-- > 0) {
      valve* vx = q.front();
      q.pop_front();
      for (auto& lv : vx->others) {
        valve* vn = get(lv);
        auto it = visited.find(vn);
        int m = visited[vx] + 1;
        if (it == visited.end()) {
          visited.insert({vn, m});
          q.push_back(vn);
        } else {
          if (m < it->second) {
            it->second = m;
            q.push_back(vn);
          }
        }
      }
    }
  }
}

void update_total(int* total, std::map<valve*, int>& opened, int m) {
  int t{0};
  for (auto& kv : opened) {
    t += kv.first->rate * (kv.second - m);
  }
  *total = t;
}

bool is_open(valve* v, std::map<valve*, int>& opened) { return opened.find(v) != opened.end(); }

// m minutes left at v
void flow(int m, valve* v, std::map<valve*, int> opened, int os, int total, int* max) {
  update_total(&total, opened, m);
  if (m == 0) {
    if (*max < total) {
      *max = total;
    }
  } else {
    if (os >= maxopened) {
      flow(m - 1, v, opened, os, total, max);
    } else {
      if (!is_open(v, opened)) {
        if (v->rate > 0) {
          os += 1;
          m -= 1;
        }
        opened.insert({v, m});
      }
      if (os >= maxopened) {
        flow(m - 1, v, opened, os, total, max);
      } else {
        // choose next
        auto f{false};
        for (auto& n : diagram[v]) {
          if (m >= n.m && !is_open(n.v, opened)) {
            f = true;
            flow(m - n.m, n.v, opened, os, total, max);
          }
        }
        if (!f && m > 0) {
          flow(m - 1, v, opened, os, total, max);
        }
      }
    }
  }
}

bool update(valve_m vs[2], valve_m n0, bool g0, bool m0, valve_m n1, bool g1, bool m1) {
  if (g0 && m0 && g1 && m1) {
    vs[0].m -= n0.m;
    vs[0].v = n0.v;
    vs[1].m -= n1.m;
    vs[1].v = n1.v;
    return true;
  }
  if (g0 && m0 && (!g1 || !m1)) {
    vs[0].m -= n0.m;
    vs[0].v = n0.v;
    vs[1].m -= vs[1].m > 0 ? 1 : 0;
    return true;
  }
  if ((!g0 || !m0) && g1 && m1) {
    vs[0].m -= vs[0].m > 0 ? 1 : 0;
    vs[1].m -= n1.m;
    vs[1].v = n1.v;
    return true;
  }
  return false;
}

struct valve_mm {
  valve_m vm1;
  valve_m vm2;

  friend bool operator<(valve_mm m1, valve_mm m2) {
    return m1.vm1 < m2.vm1 ? true : m1.vm1 > m2.vm1 ? false : m1.vm2 < m2.vm2;
  }
};

void flow(valve_m vs[2], std::map<valve_mm, int>& visited, std::map<valve*, int> opened, int os, int total, int* max) {
  for (auto i = 0; i < 2; i++) {
    update_total(&total, opened, vs[i].m);
  }
  if (vs[0].m == 0 && vs[1].m == 0) {
    if (*max < total) {
      *max = total;
    }
  } else {
    auto p = visited.insert({{vs[0], vs[1]}, total});
    if (!p.second) {
      return;
      // if (total <= p.first->second) {
      //   return;
      // }
    }
    // for (auto i = 0; i < 2; i++) {
    //   std::cout << i << " at " << vs[i].v->name << " when " << vs[i].m << " minutes left, total is " << total
    //             << std::endl;
    // }
    if (os >= maxopened) {
      vs[0].m -= vs[0].m > 0 ? 1 : 0;
      vs[1].m -= vs[1].m > 0 ? 1 : 0;
      flow(vs, visited, opened, os, total, max);
    } else {
      for (auto i = 0; i < 2; i++) {
        if (!is_open(vs[i].v, opened)) {
          if (vs[i].v->rate > 0) {
            os += 1;
            vs[i].m -= 1;
          }
          opened.insert({vs[i].v, vs[i].m});
        }
      }
      if (os >= maxopened) {
        vs[0].m -= vs[0].m > 0 ? 1 : 0;
        vs[1].m -= vs[1].m > 0 ? 1 : 0;
        flow(vs, visited, opened, os, total, max);
      } else {
        auto& ns0 = diagram[vs[0].v];
        auto& ns1 = diagram[vs[1].v];
        for (auto& n0 : ns0) {
          for (auto& n1 : ns1) {
            if (n0.v != n1.v) {
              valve_m vx[2];
              vx[0] = vs[0];
              vx[1] = vs[1];

              auto g0 = !is_open(n0.v, opened);
              auto m0 = vx[0].m > n0.m;
              auto g1 = !is_open(n1.v, opened);
              auto m1 = vx[1].m > n1.m;
              if (update(vx, n0, g0, m0, n1, g1, m1)) {
                auto v = visited;
                flow(vx, v, opened, os, total, max);
              }
            }
          }
        }
      }
    }
  }
}

std::pair<int, int> day16(line_view file) {
  per_line(file, [](line_view lv) {
    valve* v = new valve{lv};
    valves[v->name] = v;

    if (v->rate > 0) {
      maxopened += 1;
    }
    return true;
  });

  for (auto& kv : valves) {
    valve* v = kv.second;
    std::map<valve*, int> visited;
    rank_neighbours(v, visited);
    std::vector<valve_m> vs;
    for (auto& kv : visited) {
      if (kv.first != v && kv.first->rate > 0) {
        vs.emplace_back(valve_m{kv.first, kv.second});
      }
    }
    std::sort(vs.begin(), vs.end(), [](valve_m v1, valve_m v2) {
      return v1.m < v2.m ? true : v1.m > v2.m ? false : v1.v->rate > v2.v->rate;
    });
    diagram.insert({v, vs});
  }

  // for (auto& kv : diagram) {
  //   std::cout << kv.first->name << ": ";
  //   for (auto& m : kv.second) {
  //     std::cout << m.v->name << "(" << m.m << "," << m.v->rate << ") ";
  //   }
  //   std::cout << std::endl;
  // }

  int m1{INT32_MIN};
  std::map<valve*, int> opened;
  flow(30, get("AA"), opened, 0, 0, &m1);

  opened.clear();
  int m2{INT32_MIN};
  valve_m vs[2] = {
      {get("AA"), 26},
      {get("AA"), 26},
  };
  std::map<valve_mm, int> visited;
  flow(vs, visited, opened, 0, 0, &m2);
  printf("%d %d\n", m1, m2);

  return {m1, 0};
}

} // namespace aoc2022