#include "common/aoc.h"

namespace y2016::day24
{
  REGISTER_DAY(2016, Day24, Array2D<char>, int);

  REGISTER_TEST_EXAMPLE(2016, Day24, ExampleInput, 1, 14);
  REGISTER_TEST(2016, Day24, Input, 1, 448);
  REGISTER_TEST_EXAMPLE(2016, Day24, ExampleInput, 2, 20);
  REGISTER_TEST(2016, Day24, Input, 2, 672);

  READ_INPUT(input)
  {
    return Input::ReadArray2D(input);
  }

  struct State
  {
    Index2D pos;
    char goalChar;

    bool operator<(const State& other) const
    {
      return pos < other.pos;
    }
  };

  std::vector<std::pair<int, State>> Branch(const Array2D<char>& map, const State& state)
  {
    std::vector<std::pair<int, State>> branches;
    for (const auto& pos : map.GetNeighbors(state.pos, false))
    {
      if (map[pos] != '#')
        branches.emplace_back(1, State{pos, state.goalChar});
    }
    return branches;
  }

  bool Goal(const Array2D<char>& map, const State& state)
  {
    return map[state.pos] == state.goalChar;
  }

  Graph<char, int> GetGraph(const Array2D<char>& map)
  {
    std::vector<Index2D> positions(10);
    for (auto it = map.begin(); it != map.end(); it++)
    {
      if (*it >= '0' && *it <= '9')
      {
        positions[*it - '0'] = it.index;
      }
    }
    while (!positions.back().IsValid())
      positions.pop_back();

    Graph<char, int> graph;
    for (size_t i = 0; i < positions.size(); i++)
    {
      for (size_t j = i + 1; j < positions.size(); j++)
      {
        int distance = Helper::Dijkstras(map, State{positions[i], (char)(j + '0')}, Branch, Goal);
        graph.AddTwoWayEdge(i + '0', j + '0', distance);
      }
    }
    return graph;
  }

  OUTPUT1(input)
  {
    return Helper::TSP(GetGraph(input), '0', false);
  }

  OUTPUT2(input)
  {
    return Helper::TSP(GetGraph(input), '0', true);
  }
}