#include <cstddef>
#include <cstdint>
#include <iostream>
#include <stdexcept>
#include <sstream>
#include <unordered_map>
#include <unordered_set>

#include <doctest.h>

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

auto Part1(aocpp::Grid const &input, std::int64_t reps = 10'000) -> std::int64_t
{
  std::unordered_set<aocpp::Coord> infected;

  // Initialize set of infected locations
  input.each([&infected](auto const pos, auto const cell)
  {
    if (cell == '#')
    {
      infected.insert(pos);
    }
  });

  // Start in the middle
  aocpp::Coord pos;
  pos.x = input.Cols() / 2;
  pos.y = input.Rows() / 2;

  // Start going "up"
  aocpp::Coord vel{.x = 0, .y = -1};

  // Count the number of iterations that cause an infection
  std::uint64_t infections = 0;

  while (reps --> 0)
  {
    auto const [it, added] = infected.insert(pos);
    if (added)
    {
      // was clean
      vel = CCW(vel); // turn "left"
      infections++;
    }
    else
    {
      // was infected
      vel = CW(vel); // turn "right"
      infected.erase(it); // clean
    }
    pos += vel; // advance
  }

  return infections;
}

auto Part2(aocpp::Grid const &input, std::int64_t reps = 10'000'000) -> std::int64_t
{
  std::unordered_map<aocpp::Coord, char> cells;

  // Initialize set of infected locations
  input.each([&cells](auto const pos, auto const cell)
  {
    if (cell == '#')
    {
      cells.try_emplace(pos, '#');
    }
  });

  // Start in the middle
  aocpp::Coord pos;
  pos.x = input.Cols() / 2;
  pos.y = input.Rows() / 2;

  // Start going "up"
  aocpp::Coord vel{.x = 0, .y = -1};

  // Count the number of iterations that cause an infection
  std::uint64_t infections = 0;

  while (reps --> 0)
  {
    auto const [it, added] = cells.try_emplace(pos, 'W'); // clean becomes weakened
    if (added)
    {
      // already inserted 'W'
      vel = CCW(vel); // turn "left"
    }
    else
    {
      switch (it->second)
      {
      case 'W':
        // no turn
        it->second = '#';
        infections++;
        break;
      case '#':
        vel = CW(vel); // turn "right"
        it->second = 'F';
        break;
      case 'F':
        vel = Turn180(vel);
        cells.erase(it); // clean
        break;
      default:
        throw std::logic_error{"unexpected cell in steps"};
      }
    }
    pos += vel; // advance
  }

  return infections;
}

TEST_SUITE("2017-22 examples")
{
  TEST_CASE("part 1")
  {
    std::istringstream in{
        "..#\n"
        "#..\n"
        "...\n"};
    auto const input = aocpp::Grid::Parse(in);
    CHECK(5 == Part1(input, 7));
    CHECK(41 == Part1(input, 70));
    CHECK(5587 == Part1(input, 10'000));
  }
  TEST_CASE("part 2")
  {
    std::istringstream in{
        "..#\n"
        "#..\n"
        "...\n"};
    auto const input = aocpp::Grid::Parse(in);
    CHECK(26 == Part2(input, 100));
    CHECK(2'511'944 == Part2(input, 10'000'000));
  }
}

auto Main(std::istream &in, std::ostream &out) -> void
{
  auto const input = aocpp::Grid::Parse(in);
  out << "Part 1: " << Part1(input) << std::endl;
  out << "Part 2: " << Part2(input) << std::endl;
}