#include <algorithm>
#include <concepts>
#include <cstdint>
#include <iostream>
#include <iterator>
#include <map>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>

#include <doctest.h>

#include <aocpp/Startup.hpp>
#include <aocpp/Coord.hpp>
#include <aocpp/Grid.hpp>

using namespace aocpp;

namespace {

auto FindBugs(Grid const& grid) -> std::vector<Coord> {
  std::vector<Coord> result;

  std::int64_t w = grid.Cols();
  std::int64_t h = grid.Rows();

  grid.each([&](Coord c, char v) {
    if (v == '#') result.push_back(c);
  });
  return result;
}

struct Neighbor1 {
  auto operator()(Coord c, auto f) const -> void {
    if (c.x > 0) f(Left(c));
    if (c.x < 4) f(Right(c));
    if (c.y > 0) f(Up(c));
    if (c.y < 4) f(Down(c));
  }
};

struct Neighbor2 {
  using C3 = std::pair<Coord, std::int64_t>;
  auto operator()(C3 cd, auto f) const -> void {

    auto left_neighbors = [cd,f](auto k_, auto k) {
      auto const [c,d] = cd;
      auto c_ = k_(c);
      if (c_.x == 1 && c_.y == 0) {
        for (std::int64_t yi = -2; yi <= 2; yi++) {
          f({k({2,yi}),d+1});
        }
      } else if (c_.x > -2) {
        f({k(Left(c_)),d});
      } else {
        f({k({-1,0}),d-1});
      }
    };
    auto id = [](Coord i) { return i; };
    left_neighbors(id, id);
    left_neighbors(Turn180, Turn180);
    left_neighbors(CW, CCW);
    left_neighbors(CCW, CW);
  }
};

template<std::totally_ordered C, std::invocable<C, void(C const&)> F>
auto Step(std::vector<C> const& bugs, F with_neighbors = F{})
{
  std::map<C, int> adjacent;
  for (auto const& x : bugs) {
    with_neighbors(x, [&adjacent](C const& n) -> void { adjacent[n]++; });
  }

  std::vector<C> result;
  for (auto const& [k,v] : adjacent) {
    switch (v) {
    case 1:
      result.push_back(k);
      break;
    case 2:
      if (!std::binary_search(bugs.begin(), bugs.end(), k)) {
        result.push_back(k);
      }
      break;
    }
  }

  return result;
}

auto Bio(std::vector<Coord> const& bugs) -> std::uint32_t {
  std::uint32_t result = 0;
  for (auto const& c : bugs) {
    result |= 1U << (5 * c.y + c.x);
  }
  return result;
}

auto Part1(std::vector<Coord> const& bugs) -> std::uint32_t {
  auto cursor = bugs;
  std::set<std::uint32_t> seen;
  std::uint32_t bio;
  while (seen.insert(bio = Bio(cursor)).second) {
    cursor = Step<Coord, Neighbor1>(cursor);
  }
  return bio;
}

auto Part2(std::vector<Coord> const& bugs, int minutes) {
  std::vector<Neighbor2::C3> bugs2;
  for (auto const& [x,y] : bugs) {
    bugs2.push_back({{x-2,y-2},0});
  }
  for (int i = 0; i < minutes; i++) {
    bugs2 = Step<Neighbor2::C3, Neighbor2>(bugs2);
  }
  return bugs2.size();
}

} // namespace

TEST_SUITE("documented examples") {
  TEST_CASE("part 1") {
    std::istringstream in {"....#\n#..#.\n#..##\n..#..\n#....\n"};
    REQUIRE(Part1(FindBugs(Grid::Parse(in))) == 2129920);
  }
  TEST_CASE("part 2") {
    std::istringstream in {"....#\n#..#.\n#.?##\n..#..\n#....\n"};
    REQUIRE(Part2(FindBugs(Grid::Parse(in)), 10) == 99);

  }
}

auto Main(std::istream & in) -> void
{
  auto const bugs {FindBugs(Grid::Parse(in))};
  std::cout << "Part 1: " << Part1(bugs) << std::endl;
  std::cout << "Part 2: " << Part2(std::move(bugs), 200) << std::endl;
}