#include <algorithm>
#include <cstdint>
#include <iostream>
#include <sstream>
#include <tuple>
#include <cstdio>
#include <cstdlib>
#include <numeric>
#include <map>
#include <set>
#include <functional>
#include <stdexcept>

#include <doctest.h>

#include <aocpp/Startup.hpp>
#include <aocpp/Parsing.hpp>
#include <dlx.hpp>

using aocpp::Startup;
using aocpp::SplitOn;
using dlx::Dlx;

namespace {

using Input = std::vector<std::pair<std::vector<std::string>, std::vector<std::string>>>;

auto Parse(std::istream & in) -> Input
{
  Input results;
  std::string line;
  while (std::getline(in, line)) {
    char const sep[] = " (contains ";
    auto i = line.find(sep);
    if (i == std::string::npos || line.back() != ')') { throw std::runtime_error{"bad input"}; }
    line.pop_back();
    results.emplace_back(
      SplitOn(line.substr(0, i), " "),
      SplitOn(line.substr(i + std::size(sep) - 1), ", ")
    );
  }
  return results;
}

auto Arrange(Input const& input) -> std::map<std::string, std::vector<std::string>>
{
  std::map<std::string, std::vector<std::string>> possible;

  for (auto const& [ingredients, allergens] : input) {
    for (auto const& allergen : allergens) {
      auto ingredients_ = ingredients;
      std::sort(ingredients_.begin(), ingredients_.end());
      
      if (auto it = possible.find(allergen); it != possible.end()) {
        auto & previous = it->second;
        std::vector<std::string> intersection;
        std::set_intersection(
          previous.begin(), previous.end(),
          ingredients_.begin(), ingredients_.end(),
          std::back_inserter(intersection)
        );
        previous = std::move(intersection);
      } else {
        possible.emplace(allergen, std::move(ingredients_));
      }
    }
  }
  return possible;
}

auto Part1(Input const& input) -> std::size_t
{
  std::map<std::string, std::vector<std::string>> possible = Arrange(input);
  std::map<std::string, std::size_t> occurences;

  for (auto const& [ingredients, allergens] : input) {
    for (auto const& ingredient : ingredients) {
      occurences[ingredient]++;
    }
  }

  // Remove all candidates
  for (auto const& [k,vs] : possible) {
    for (auto const& v : vs) {
      occurences.erase(v);
    }
  }

  std::size_t result = 0;
  for (auto const& [_,v] : occurences) {
    result += v;
  }

  return result;
}

auto Part2(Input const& input) {
  std::map<std::string, std::size_t> a_ids;
  std::map<std::string, std::size_t> i_ids;
  std::vector<std::string> id_to_a;
  std::vector<std::string> id_to_i;

  for (auto const& [is, as] : input) {
    for (auto const& i : is) {
      if (i_ids.emplace(i, i_ids.size()).second) { id_to_i.push_back(i); }
    }
    for (auto const& a : as) {
      if (a_ids.emplace(a, a_ids.size()).second) { id_to_a.push_back(a); }
    }
  }

  // rows: allergen + allergens * ingredient
  // columns: all allergens, then all ingredients
  Dlx dlx;

  for (auto const& [allergen, ingredients] : Arrange(input)) {
    auto const a_id = a_ids[allergen];
    for (auto const& ingredient : ingredients) {
      auto const i_id = i_ids[ingredient];
      auto const row_id = a_id + a_ids.size() * i_id;
      dlx.Set(row_id, a_id);
      dlx.Set(row_id, a_ids.size() + i_id);
    }
  }

  // ingredients don't have to be assigned to anything
  for (auto const& [i,n] : i_ids) {
    dlx.MarkOptional(a_ids.size() + n);
  }

  // extract assignment map from dlx row ids
  std::map<std::string, std::string> entries;
  dlx::ForallCover(dlx, [&](auto & sln){
    for (auto const row : sln) {
      entries.emplace(id_to_a[row % a_ids.size()], id_to_i[row / a_ids.size()]);
    }
    return true; // exit early
  });

  // build comma-separated list of ingredients
  std::string result;
  for (auto const & [_,i] : entries) {
    if (!result.empty()) result += ',';
    result += i;
  }
  return result;
}

} // namespace

TEST_SUITE("documented examples") {

std::istringstream in {
  "mxmxvkd kfcds sqjhc nhms (contains dairy, fish)\n"
  "trh fvjkl sbzzf mxmxvkd (contains dairy)\n"
  "sqjhc fvjkl (contains soy)\n"
  "sqjhc mxmxvkd sbzzf (contains fish)\n"
};
auto input = Parse(in);

TEST_CASE("parser") {
  Input expected {
    {{"mxmxvkd", "kfcds", "sqjhc", "nhms"}, {"dairy", "fish"}},
    {{"trh", "fvjkl", "sbzzf", "mxmxvkd"}, {"dairy"}},
    {{"sqjhc", "fvjkl"}, {"soy"}},
    {{"sqjhc", "mxmxvkd", "sbzzf"}, {"fish"}}
  };
  REQUIRE(input == expected);
}

TEST_CASE("part 1") {
  REQUIRE(Part1(input) == 5);
}

TEST_CASE("part 2") {
  REQUIRE(Part2(input) == "mxmxvkd,sqjhc,fvjkl");
}

}

auto main(int argc, char** argv) -> int {
  auto const input = Parse(*Startup(argc, argv));
  std::cout << "Part 1: " << Part1(input) << std::endl;
  std::cout << "Part 2: " << Part2(input) << std::endl;
}