#include <algorithm>
#include <cstdint>
#include <iostream>
#include <iterator>
#include <iterator>
#include <sstream>
#include <string>
#include <vector>

#include <boost/phoenix.hpp>
#include <boost/range/algorithm.hpp>
#include <boost/range/irange.hpp>
#include <doctest.h>

#include <aocpp/Startup.hpp>

namespace
{
  namespace phx = boost::phoenix;
  using phx::placeholders::arg1;

  auto Parse(std::istream &in) -> std::vector<std::string>
  {
    using It = std::istream_iterator<std::string>;
    std::vector<std::string> result;
    std::move(It{in}, It{}, std::back_inserter(result));

    if (result.empty())
    {
      throw std::runtime_error{"Empty input"};
    }

    if (std::any_of(std::next(result.begin()), result.end(),
                    [W = result[0].size()](auto const &s)
                    { return s.size() != W; }))
    {
      throw std::runtime_error{"input line length mismatch"};
    }

    return result;
  }

  auto Part1(std::vector<std::string> const &entries) -> std::uint64_t
  {
    auto const N{entries.size()};
    auto const W{entries[0].size()};

    std::uint64_t lo{};
    std::uint64_t hi{};

    for (auto const i : boost::irange(W))
    {
      auto const bit_test = '1' == arg1[i];

      auto const n = boost::count_if(entries, bit_test);
      lo <<= 1;
      hi <<= 1;
      (2 * n >= N ? hi : lo)++;
    }

    return lo * hi;
  }

  auto Part2_(std::vector<std::string> entries, bool const big) -> std::uint64_t
  {
    auto const N{entries.size()};
    auto const W{entries[0].size()};

    std::uint64_t result{};
    auto begin{entries.begin()};
    auto end{entries.end()};

    for (auto const i : boost::irange(W))
    {
      auto const bit_test = '1' == arg1[i];

      result <<= 1;

      // If there's only one number left we should select it
      if (std::distance(begin, end) == 1)
      {
        result |= bit_test(*begin);
      }
      else
      {
        auto const mid = std::partition(begin, end, bit_test);
        auto const ones = std::distance(begin, mid);
        auto const zeros = std::distance(mid, end);
        auto const bit = big == (ones >= zeros);
        result |= bit;
        (bit ? end : begin) = mid;
      }
    }

    return result;
  }

  auto Part2(std::vector<std::string> entries) -> std::uint64_t
  {
    return Part2_(entries, false) * Part2_(entries, true);
  }
}

TEST_SUITE("2021-03")
{
  TEST_CASE("example")
  {
    std::istringstream in{
        "00100\n"
        "11110\n"
        "10110\n"
        "10111\n"
        "10101\n"
        "01111\n"
        "00111\n"
        "11100\n"
        "10000\n"
        "11001\n"
        "00010\n"
        "01010\n"};
    auto entries = Parse(in);
    CHECK(Part1(entries) == 198);
    CHECK(Part2_(entries, false) == 10);
    CHECK(Part2_(entries, true) == 23);
    CHECK(Part2(std::move(entries)) == 230);
  }
}

auto Main(std::istream &in, std::ostream &out) -> void
{
  auto entries{Parse(in)};
  out << "Part 1: " << Part1(entries) << std::endl;
  out << "Part 2: " << Part2(std::move(entries)) << std::endl;
}