#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, const 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);
        }
      }
    }
  }
}

// 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;
//   }
// };

std::vector<valve_m> cango(valve* v, int m0, int m, const std::map<valve*, int>& opened) {
  std::vector<valve_m> t;
  for (auto& vm : diagram[v]) {
    if (!is_open(vm.v, opened) && m > vm.m && m == m0) {
      t.emplace_back(vm);
    }
  }
  return t;
}

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

            vx[1].m -= n.m;
            vx[1].v = n.v;
            auto v = visited;
            flow(m - 1, vx, v, opened, os, total, max);
          }
        }
        if (ns0.size() > 0 && ns1.size() == 0) {
          for (auto& n : ns0) {
            valve_m vx[2];
            vx[0] = vs[0];
            vx[1] = vs[1];

            vx[0].m -= n.m;
            vx[0].v = n.v;
            auto v = visited;
            flow(m - 1, vx, v, opened, os, total, max);
          }
        }
        if (ns0.size() > 0 && ns1.size() > 0) {
          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];

                vx[0].m -= n0.m;
                vx[0].v = n0.v;
                vx[1].m -= n1.m;
                vx[1].v = n1.v;
                auto v = visited;
                flow(m - 1, 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<int, int> visited;
  // flow(26, vs, visited, opened, 0, 0, &m2);

  return {m1, 1999};
}

} // namespace aoc2022