#include <cstdint>
#include <deque>
#include <iostream>
#include <map>
#include <numeric>
#include <optional>
#include <set>
#include <sstream>
#include <string>
#include <tuple>
#include <vector>

#include <doctest.h>

#include <aocpp/Startup.hpp>
#include <aocpp/Grid.hpp>
using namespace aocpp;

namespace {

auto IsUnit(char c) -> bool {
  return c == 'G' || c == 'E';
}

auto IsOpposed(char x, char y) -> bool {
  return x == 'G' && y == 'E' || x == 'E' && y == 'G';
}

struct IsBefore {
  auto operator()(Coord const& a, Coord const& b) const -> bool {
    return a.y < b.y || a.y == b.y && a.x < b.x;
  }
};

int const initial_hp = 200;
Coord const reading_order[] {{0,-1}, {-1,0}, {1,0}, {0,1}};

struct Game {
  Grid grid_;
  std::map<Coord, int> hp_;
  int e_damage_;
  int g_damage_;
  bool no_elves_can_die;
  std::ostream * debug_out_;

  Game(Grid grid)
  : grid_(std::move(grid))
  , hp_{}
  , e_damage_(3)
  , g_damage_(3)
  , no_elves_can_die{false}
  , debug_out_{}
  {
    // set each unit to its initial hit points
    grid_.each([&](auto k, auto v) {
      if (IsUnit(v)) {
        hp_[k] = initial_hp;
      }
    });
  }

  auto Simulate() -> std::optional<int>;
  auto PickTarget(Coord my_coord, char my_team)
    -> std::optional<std::map<Coord, int>::iterator>;
  auto Move(Coord my_coord, char my_team) -> std::optional<Coord>;
};

auto operator<<(std::ostream & out, Game const& game) -> std::ostream & {
  for (std::size_t y = 0; y < game.grid_.rows.size(); ++y) {
    out << game.grid_.rows[y];
    bool first = true;
    for (std::size_t x = 0; x < game.grid_.rows[y].size(); ++x) {
      auto c = game.grid_.rows[y][x];
      if (IsUnit(c)) {
        if (first) {
          first = false;
          out << "   ";
        } else {
          out << ", ";
        }
        Coord coord {std::int64_t(x),std::int64_t(y)};
        out << c << "(" << game.hp_.at(coord) << ")";
      }
    }
    out << std::endl;
  }
  return out;
}

// Determine the best target to attack from the current position, if there is one.
auto Game::PickTarget(Coord my_coord, char my_team)
-> std::optional<std::map<Coord, int>::iterator>
{
  std::optional<std::map<Coord, int>::iterator> best;
  int best_hp;
  for (auto const dir : reading_order) {
    auto here = dir + my_coord;
    if (IsOpposed(my_team, grid_[here])) {
      auto it = hp_.find(here);
      if (!best || best_hp > it->second) {
          best = it;
          best_hp = it->second;
      }
    }
  }
  return best;
}

// Determine the best location to move to, if there is one.
auto Game::Move(Coord my_coord, char my_team) -> std::optional<Coord>
{
  int best_distance;
  std::optional<Coord> best_start;
  Coord best_end;

  std::set<Coord> seen;
  std::deque<std::tuple<Coord, Coord, int>> todo;

  for (auto const dir : reading_order) {
    auto start = my_coord + dir;
    if (grid_[start] == '.') {
      todo.push_back({start,start,1});
    }
  }

  while (!todo.empty()) {
    auto [start, cursor, steps] = todo.front();
    todo.pop_front();

    if (best_start) {
      if (best_distance < steps) break;
      if (best_start && !IsBefore()(cursor, best_end)) continue;
    } else {
      if (!seen.insert(cursor).second) { continue; }
      for (auto const dir : reading_order) {
        auto const next = cursor + dir;
        if (grid_[next] == '.') {
          todo.push_back({start, cursor + dir, steps+1});
        }
      }
    }
    if (PickTarget(cursor, my_team)) {
      best_start = start;
      best_end = cursor;
      best_distance = steps;
    }
  }
  return best_start;
}

auto Game::Simulate() -> std::optional<int> {

  for (int rounds = 0;; rounds++) {
    if (debug_out_) {
      *debug_out_ << "After " << rounds << " rounds:" << std::endl
                  << *this << std::endl;
    }

    // determine order of unit updates
    std::set<Coord, IsBefore> turn_order;
    for (auto const& [k,_] : hp_) {
      turn_order.insert(k);
    }

    for (auto active_coord : turn_order) {
      char active_team = grid_[active_coord];

      // Detect when no targets remain for this unit and end the game
      bool game_over = true;
      grid_.each([&](auto k, auto v) {
        if (IsOpposed(active_team, v)) game_over = false;
      });
      if (game_over) {
        int total_hp = 0;
        for (auto const& [_,v] : hp_) { total_hp += v; }
        return {rounds * total_hp};
      }

      // Try to pick a target without moving
      auto target = PickTarget(active_coord, active_team);

      // only move when necessary
      if (!target) {
        // move if we can
        if (auto const destination = Move(active_coord, active_team)) {
          std::swap(grid_[*destination], grid_[active_coord]);

          auto const it = hp_.find(active_coord);
          hp_[*destination] = it->second;
          hp_.erase(it);

          active_coord = *destination;
        }

        target = PickTarget(active_coord, active_team);
      }

      // target in range, do the attack
      if (target) {
        auto & [target_coord, target_hp] = **target;
        auto damage = active_team == 'G' ? g_damage_ : e_damage_;

        // lethal
        if (target_hp <= damage) {
          if (no_elves_can_die && grid_[target_coord] == 'E') return {};

          grid_[target_coord] = '.';
          turn_order.erase(target_coord);
          hp_.erase(*target); // invalidates target, target_coord, target_hp
        } else {
          target_hp -= damage;
        }
      }
    }
  }
}

auto Part2(Game game) {
  game.no_elves_can_die = true;
  int best;

  auto const attempt = [&](int const power) {
    auto g = game;
    g.e_damage_ = power;
    return g.Simulate();
  };

  int too_low = 3;
  int high = 4;
  for(;;) {
    if (auto outcome = attempt(high)) {
      best = *outcome; break;
    }
    too_low = high;
    high *= 2;
  }

  while (too_low+1 < high) {
    auto x = std::midpoint(too_low, high);
    if (auto outcome = attempt(x)) {
      best = *outcome; high = x;
    } else {
      too_low = x;
    }
  }

  return std::make_pair(best, high);
}

} // namespace

TEST_SUITE("2018-15 examples") {

  auto case1(int expect, std::string start, std::string end) -> void {
    std::istringstream in { std::move(start) };
    std::ostringstream out;
    Game game {Grid::Parse(in)};
    auto outcome = game.Simulate();
    out << game;
    CHECK(out.str() == end);
    CHECK(outcome == std::optional<int>{expect});
  }

  TEST_CASE("example 1") {
    case1(27730,
      "#######\n"
      "#.G...#\n"
      "#...EG#\n"
      "#.#.#G#\n"
      "#..G#E#\n"
      "#.....#\n"
      "#######\n",
      "#######\n"
      "#G....#   G(200)\n"
      "#.G...#   G(131)\n"
      "#.#.#G#   G(59)\n"
      "#...#.#\n"
      "#....G#   G(200)\n"
      "#######\n");
  }
  TEST_CASE("example 2") {
    case1(36334,
      "#######\n"
      "#G..#E#\n"
      "#E#E.E#\n"
      "#G.##.#\n"
      "#...#E#\n"
      "#...E.#\n"
      "#######\n",
      "#######\n"
      "#...#E#   E(200)\n"
      "#E#...#   E(197)\n"
      "#.E##.#   E(185)\n"
      "#E..#E#   E(200), E(200)\n"
      "#.....#\n"
      "#######\n");
  }
  TEST_CASE("example 3") {
    case1(39514,
      "#######\n"
      "#E..EG#\n"
      "#.#G.E#\n"
      "#E.##E#\n"
      "#G..#.#\n"
      "#..E#.#\n"
      "#######\n",
      "#######\n"
      "#.E.E.#   E(164), E(197)\n"
      "#.#E..#   E(200)\n"
      "#E.##.#   E(98)\n"
      "#.E.#.#   E(200)\n"
      "#...#.#\n"
      "#######\n");
  }
  TEST_CASE("example 4") {
    case1(27755,
      "#######\n"
      "#E.G#.#\n"
      "#.#G..#\n"
      "#G.#.G#\n"
      "#G..#.#\n"
      "#...E.#\n"
      "#######\n",
      "#######\n"
      "#G.G#.#   G(200), G(98)\n"
      "#.#G..#   G(200)\n"
      "#..#..#\n"
      "#...#G#   G(95)\n"
      "#...G.#   G(200)\n"
      "#######\n");
  }
  TEST_CASE("example 5") {
    case1(28944,
      "#######\n"
      "#.E...#\n"
      "#.#..G#\n"
      "#.###.#\n"
      "#E#G#G#\n"
      "#...#G#\n"
      "#######\n",
      "#######\n"
      "#.....#\n"
      "#.#G..#   G(200)\n"
      "#.###.#\n"
      "#.#.#.#\n"
      "#G.G#G#   G(98), G(38), G(200)\n"
      "#######\n");
  }
  TEST_CASE("example 6") {
    case1(18740,
      "#########\n"
      "#G......#\n"
      "#.E.#...#\n"
      "#..##..G#\n"
      "#...##..#\n"
      "#...#...#\n"
      "#.G...G.#\n"
      "#.....G.#\n"
      "#########\n",
      "#########\n"
      "#.G.....#   G(137)\n"
      "#G.G#...#   G(200), G(200)\n"
      "#.G##...#   G(200)\n"
      "#...##..#\n"
      "#.G.#...#   G(200)\n"
      "#.......#\n"
      "#.......#\n"
      "#########\n"
    );
  }

  auto case2(int expect, std::string start, std::string end) -> void {
    std::istringstream in { std::move(start) };
    std::ostringstream out;
    Game game {Grid::Parse(in)};
    auto [outcome, e_damage] = Part2(game);
    game.e_damage_ = e_damage;
    game.Simulate();
    out << game;
    CHECK(out.str() == end);
    CHECK(outcome == expect);
  }

  TEST_CASE("example 7") {
    case2(4988,
      "#######\n"
      "#.G...#\n"
      "#...EG#\n"
      "#.#.#G#\n"
      "#..G#E#\n"
      "#.....#\n"
      "#######\n",
      "#######\n"
      "#..E..#   E(158)\n"
      "#...E.#   E(14)\n"
      "#.#.#.#\n"
      "#...#.#\n"
      "#.....#\n"
      "#######\n"
    );
  }

  TEST_CASE("example 8") {
    case2(31284,
      "#######\n"
      "#E..EG#\n"
      "#.#G.E#\n"
      "#E.##E#\n"
      "#G..#.#\n"
      "#..E#.#\n"
      "#######\n",
      "#######\n"
      "#.E.E.#   E(200), E(23)\n"
      "#.#E..#   E(200)\n"
      "#E.##E#   E(125), E(200)\n"
      "#.E.#.#   E(200)\n"
      "#...#.#\n"
      "#######\n");
  }

  TEST_CASE("example 9") {
    case2(3478,
      "#######\n"
      "#E.G#.#\n"
      "#.#G..#\n"
      "#G.#.G#\n"
      "#G..#.#\n"
      "#...E.#\n"
      "#######\n",
      "#######\n"
      "#.E.#.#   E(8)\n"
      "#.#E..#   E(86)\n"
      "#..#..#\n"
      "#...#.#\n"
      "#.....#\n"
      "#######\n");
  }

  TEST_CASE("example 10") {
    case2(6474,
      "#######\n"
      "#.E...#\n"
      "#.#..G#\n"
      "#.###.#\n"
      "#E#G#G#\n"
      "#...#G#\n"
      "#######\n",
      "#######\n"
      "#...E.#   E(14)\n"
      "#.#..E#   E(152)\n"
      "#.###.#\n"
      "#.#.#.#\n"
      "#...#.#\n"
      "#######\n");
  }

  TEST_CASE("example 11") {
    case2(1140,
      "#########\n"
      "#G......#\n"
      "#.E.#...#\n"
      "#..##..G#\n"
      "#...##..#\n"
      "#...#...#\n"
      "#.G...G.#\n"
      "#.....G.#\n"
      "#########\n",
      "#########\n"
      "#.......#\n"
      "#.E.#...#   E(38)\n"
      "#..##...#\n"
      "#...##..#\n"
      "#...#...#\n"
      "#.......#\n"
      "#.......#\n"
      "#########\n");
  }
}

auto main(int argc, char** argv) -> int {
  auto grid = Grid::Parse(*aocpp::Startup(argc, argv));
  Game game {std::move(grid)};
  //game.debug_out_ = &std::cout;
  std::cout << "Part 1: " << *Game{game}.Simulate() << std::endl;
  std::cout << "Part 2: " << Part2(std::move(game)).first << std::endl;
}