#include "common/aoc.h"

namespace y2016::day11
{
  struct Item
  {
    int floorPos;
    bool generator;
    int id;

    bool operator==(const Item& item) const
    {
      return floorPos == item.floorPos && id == item.id && generator == item.generator;
    }
  };

  static int FLOORS = 4;

  REGISTER_DAY(2016, Day11, std::vector<Item>, int);

  REGISTER_TEST_EXAMPLE(2016, Day11, ExampleInput, 1, 11);
  REGISTER_TEST(2016, Day11, Input, 1, 37);
  REGISTER_TEST(2016, Day11, Input, 2, 61);

  Item ParseItem(const std::string_view& str, int floorPos, std::map<std::string, int>& nameToId)
  {
    size_t posDash = str.find('-');
    std::string name;
    bool generator;
    if (posDash == std::string::npos)
    {
      name = str.substr(0, str.find(' '));
      generator = true;
    }
    else
    {
      name = str.substr(0, posDash);
      generator = false;
    }
    auto it = nameToId.find(name);
    if (it != nameToId.end())
      return Item{floorPos, generator, it->second};
    int size = nameToId.size();
    nameToId[name] = size;

    return Item{floorPos, generator, size};
  }

  READ_INPUT(input)
  {
    std::vector<Item> vec;
    std::string str;
    std::map<std::string, int> nameToId;
    for (int i = 0; i < FLOORS; i++)
    {
      getline(input, str);
      std::vector<std::string_view> split = Helper::Split(str, " a ");
      for (size_t j = 1; j < split.size(); j++)
      {
        vec.emplace_back(ParseItem(split[j], i, nameToId));
      }
    }
    return vec;
  }

  struct State
  {
    std::vector<Item> items;
    int elevatorPos;

    int amountOfPairs;

    std::vector<std::pair<int, int>> GetPairs() const
    {
      // Convert the data to a format where the pairs of generator/microchip don't matter. All that matters are their
      // positions
      std::vector<std::pair<int, int>> pairs(amountOfPairs);
      for (const auto& item : items)
      {
        if (item.generator)
          pairs[item.id].first = item.floorPos;
        else
          pairs[item.id].second = item.floorPos;
      }
      std::sort(pairs.begin(), pairs.end());
      return pairs;
    }

    bool operator<(const State& state) const
    {
      if (elevatorPos != state.elevatorPos)
        return elevatorPos < state.elevatorPos;
      return GetPairs() < state.GetPairs();
    }
  };

  bool HasItem(const State& state, int floor, bool generator, int id)
  {
    return std::find(state.items.begin(), state.items.end(), Item{floor, generator, id}) != state.items.end();
  }

  bool IsValidState(const State& state)
  {
    for (int i = 0; i < FLOORS; i++)
    {
      bool hasChipNotConnected = false;
      int generators = 0;
      for (const auto& item : state.items)
      {
        if (item.floorPos != i)
          continue;
        if (item.generator)
        {
          generators++;
          continue;
        }

        if (HasItem(state, i, true, item.id))
          continue;
        hasChipNotConnected = true;
      }
      if (hasChipNotConnected && generators > 0)
      {
        return false;
      }
    }
    return true;
  }

  std::vector<std::pair<int, State>> Branch(const std::vector<Item>& input, const State& state)
  {
    std::vector<std::pair<int, State>> branches;
    for (int i = -1; i <= 1; i += 2)
    {
      if (state.elevatorPos + i < 0 || state.elevatorPos + i > 3)
        continue;

      State newState = state;
      newState.elevatorPos += i;
      for (auto& item : newState.items)
      {
        if (item.floorPos != state.elevatorPos)
          continue;

        item.floorPos = newState.elevatorPos;
        if (IsValidState(newState))
          branches.emplace_back(1, newState);
        item.floorPos = state.elevatorPos;
      }

      for (size_t i = 0; i < newState.items.size(); i++)
      {
        if (newState.items[i].floorPos != state.elevatorPos)
          continue;
        newState.items[i].floorPos = newState.elevatorPos;

        for (size_t j = i + 1; j < newState.items.size(); j++)
        {
          if (newState.items[j].floorPos != state.elevatorPos)
            continue;
          newState.items[j].floorPos = newState.elevatorPos;
          if (IsValidState(newState))
            branches.emplace_back(1, newState);
          newState.items[j].floorPos = state.elevatorPos;
        }
        newState.items[i].floorPos = state.elevatorPos;
      }
    }
    return branches;
  }

  bool Goal(const std::vector<Item>& input, const State& state)
  {
    for (size_t i = 0; i < state.items.size(); i++)
    {
      if (state.items[i].floorPos != 3)
        return false;
    }
    return true;
  }

  int Heuristic(const std::vector<Item>& input, const State& state)
  {
    int heuristic = 0;
    int moves = 0;
    for (size_t i = 0; i < state.items.size(); i++)
    {
      moves += (3 - state.items[i].floorPos);
    }
    heuristic += moves / 2;
    return heuristic;
  }

  OUTPUT1(input)
  {
    int max = 0;
    for (const auto& item : input)
    {
      max = std::max(max, item.id);
    }
    State state{input, 0, max + 1};
    return Helper::AStar(input, state, Heuristic, Branch, Goal);
  }

  OUTPUT2(input)
  {
    std::vector<Item> items = input;
    int max = 0;
    for (const auto& item : items)
    {
      max = std::max(max, item.id);
    }
    items.emplace_back(Item{0, true, max + 1});
    items.emplace_back(Item{0, false, max + 1});
    items.emplace_back(Item{0, true, max + 2});
    items.emplace_back(Item{0, false, max + 2});
    State state{items, 0, max + 3};
    return Helper::AStar(input, state, Heuristic, Branch, Goal);
  }
}