#include <array>
#include <cstdint>
#include <iostream>
#include <optional>
#include <set>
#include <sstream>
#include <vector>

#include <boost/phoenix.hpp>
#include <boost/spirit/include/qi.hpp>

#include <doctest.h>

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

namespace {

namespace phx = boost::phoenix;
namespace qi = boost::spirit::qi;

struct Entry {
  aocpp::Coord sensor;
  aocpp::Coord beacon;
};

using Input = std::vector<Entry>;

// Input file grammar
class Grammar : public qi::grammar<std::string::const_iterator, Input()> {
  qi::rule<iterator_type, aocpp::Coord> coord;
  qi::rule<iterator_type, Entry()> entry;
  qi::rule<iterator_type, Input()> input;

public:
  Grammar() : Grammar::base_type{input} {
    using namespace qi::labels;

    coord = "x=" >>
            qi::long_long [ phx::bind(&aocpp::Coord::x, _val) = _1 ] >>
            ", y=" >>
            qi::long_long [ phx::bind(&aocpp::Coord::y, _val) = _1 ];
    entry = "Sensor at " >>
             coord [ phx::bind(&Entry::sensor, _val) = _1 ] >>
            ": closest beacon is at " >>
            coord [ phx::bind(&Entry::beacon, _val) = _1 ] >>
            "\n";
    input = *entry;
  }
};

struct Range {
  std::int64_t lo, hi;
  auto operator<=>(Range const& rhs) const = default;
  auto size() const -> std::size_t {
    return hi - lo;
  }
};

template <std::size_t N>
auto subtract_to(
  std::array<Range, N> x,
  std::array<Range, N> const& y,
  std::vector<std::array<Range, N>> & out
) -> void
{
  for (std::size_t i = 0; i < N; ++i) {
    if (x[i].lo < y[i].lo) {
      auto z = x;
      z[i].hi = std::min(z[i].hi, y[i].lo);
      out.push_back(z);
    }
    if (x[i].hi > y[i].hi) {
      auto z = x;
      z[i].lo = std::max(z[i].lo, y[i].hi);
      out.push_back(z);
    }
    x[i].lo = std::max(x[i].lo, y[i].lo);
    x[i].hi = std::min(x[i].hi, y[i].hi);
    if (x[i].lo >= x[i].hi) {
      return;
    }
  }
}

template <std::size_t N>
auto subtract(
  std::array<Range, N> const& x,
  std::array<Range, N> const& y 
) -> std::vector<std::array<Range, N>>
{
  std::vector<std::array<Range, N>> result;
  subtract_to(x, y, result);
  return result;
}

auto Segment(std::int64_t center, std::int64_t radius) -> Range {
  return { center - radius, center + radius + 1};
}

auto Square(aocpp::Coord center, std::int64_t radius) -> std::array<Range, 2> {
  return { Segment(center.x + center.y, radius), Segment(center.x - center.y, radius)};
}

auto Corner(std::array<Range, 2> const& s) -> aocpp::Coord
{
  return { (s[0].lo + s[1].lo) / 2, (s[0].lo - s[1].lo) / 2};
}

auto RowSlice(aocpp::Coord sensor, std::int64_t radius, std::int64_t row)
-> std::optional<std::array<Range, 1>>
{
  auto dy = std::abs(row - sensor.y);
  auto dx = radius - dy;
  if (dx >= 0) {
    return {{Segment(sensor.x, dx)}};
  } else {
    return {};
  }
} 

auto Part1(Input const& input, std::int64_t search) -> std::size_t {
  auto region = std::vector<std::array<Range, 1>>{};
  auto region_next = std::vector<std::array<Range, 1>>{};

  for (auto const& entry : input) {
    auto radius = aocpp::Norm1(entry.beacon - entry.sensor);
    if (auto s = RowSlice(entry.sensor, radius, search)) {
      region_next.push_back(*s);
      for (auto const& x : region) {
        subtract_to(x, *s, region_next);
      }
      std::swap(region, region_next);
      region_next.clear();
    }
  }

  std::size_t result {};
  for (auto const& x : region) {
    result += x[0].hi - x[0].lo;
  }

  std::set<std::int64_t> seen;
  for (auto const& entry : input) {
    if (entry.beacon.y == search && seen.insert(entry.beacon.x).second) {
      result--;
    }
  }

  return result;
}

auto Part2(Input const& input, std::int64_t search) -> std::int64_t {

  auto region = std::vector<std::array<Range, 2>>{Square({search/2, search/2}, search)};
  auto region_next = std::vector<std::array<Range, 2>>{};

  for (auto const& entry : input) {
    auto radius = aocpp::Norm1(entry.beacon - entry.sensor);
    auto s = Square(entry.sensor, radius);
    for (auto const& x : region) {
      subtract_to(x, s, region_next);
    }
    std::swap(region, region_next);
    region_next.clear();
  }

  for (auto const& entry : region) {
    auto corner = Corner(entry);
    if (0 <= corner.x && corner.x <= search &&
        0 <= corner.y && corner.y <= search)
    {
      return 4'000'000 * corner.x + corner.y;
    }
  }

  throw std::runtime_error{"no solution to part 2"};
}

} // namespace

TEST_SUITE("2022-15 examples") {
  TEST_CASE("subtraction") {
    CHECK(subtract<1>({1,3}, {4,5}) == std::vector<std::array<Range, 1>>{{1,3}});
    CHECK(subtract<1>({4,5}, {1,3}) == std::vector<std::array<Range, 1>>{{4,5}});
    CHECK(subtract<1>({1,3}, {2,4}) == std::vector<std::array<Range, 1>>{{1,2}});
    CHECK(subtract<1>({1,3}, {0,2}) == std::vector<std::array<Range, 1>>{{2,3}});
    CHECK(subtract<1>({1,3}, {1,3}) == std::vector<std::array<Range, 1>>{});
    CHECK(subtract<1>({1,4}, {2,3}) == std::vector<std::array<Range, 1>>{{1,2}, {3,4}});
    CHECK(subtract<1>({1,4}, {2,4}) == std::vector<std::array<Range, 1>>{{1,2}});
    CHECK(subtract<1>({1,4}, {1,3}) == std::vector<std::array<Range, 1>>{{3,4}});

    CHECK(subtract<2>(
      std::array<Range,2>{Range{0,3}, Range{0,4}}, std::array<Range,2>{Range{1,2},Range{2,3}})
      ==
      std::vector<std::array<Range, 2>>{
        std::array<Range, 2>{Range{0,1}, Range{0,4}},
        std::array<Range, 2>{Range{2,3}, Range{0,4}},
        std::array<Range, 2>{Range{1,2}, Range{0,2}},
        std::array<Range, 2>{Range{1,2}, Range{3,4}}
      });
  }

  TEST_CASE("documented example") {
    auto in = std::istringstream{
      "Sensor at x=2, y=18: closest beacon is at x=-2, y=15\n"
      "Sensor at x=9, y=16: closest beacon is at x=10, y=16\n"
      "Sensor at x=13, y=2: closest beacon is at x=15, y=3\n"
      "Sensor at x=12, y=14: closest beacon is at x=10, y=16\n"
      "Sensor at x=10, y=20: closest beacon is at x=10, y=16\n"
      "Sensor at x=14, y=17: closest beacon is at x=10, y=16\n"
      "Sensor at x=8, y=7: closest beacon is at x=2, y=10\n"
      "Sensor at x=2, y=0: closest beacon is at x=2, y=10\n"
      "Sensor at x=0, y=11: closest beacon is at x=2, y=10\n"
      "Sensor at x=20, y=14: closest beacon is at x=25, y=17\n"
      "Sensor at x=17, y=20: closest beacon is at x=21, y=22\n"
      "Sensor at x=16, y=7: closest beacon is at x=15, y=3\n"
      "Sensor at x=14, y=3: closest beacon is at x=15, y=3\n"
      "Sensor at x=20, y=1: closest beacon is at x=15, y=3\n"
    };
    auto const input = aocpp::ParseGrammar(Grammar{}, in);
    CHECK(26 == Part1(input, 10));
    CHECK(56000011 == Part2(input, 20));
  }
}

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