#include "common/aoc.h"

namespace y2015::day19
{
  struct Input
  {
    std::multimap<std::string, std::string> conversions;
    std::string molecule;
  };

  REGISTER_DAY(2015, Day19, Input, int);

  REGISTER_TEST_EXAMPLE(2015, Day19, ExampleInput, 1, 7);
  REGISTER_TEST(2015, Day19, Input, 1, 576);
  // REGISTER_TEST_EXAMPLE(2015, Day19, ExampleInput, 2, 6);
  REGISTER_TEST(2015, Day19, Input, 2, 207);

  READ_INPUT(input)
  {
    Input i;
    std::string str;
    while (getline(input, str))
    {
      if (str.empty())
        break;

      std::string from;
      std::string to;
      std::stringstream ss{str};
      ss >> from >> "=>" >> to;
      i.conversions.emplace(from, to);
    }
    getline(input, i.molecule);
    return i;
  }

  std::vector<size_t> FindAll(const std::string& str, const std::string& find)
  {
    std::vector<size_t> positions;
    size_t pos = str.find(find);
    while (pos != std::string::npos)
    {
      positions.emplace_back(pos);
      pos = str.find(find, pos + 1);
    }
    return positions;
  }

  bool IsAtom(const std::string& sv)
  {
    std::regex regex{"^(e|[A-Z][a-z]?)$"};
    return std::regex_match(sv, regex);
  }

  std::string Prettify(const std::string& molecule)
  {
    std::string result = molecule;
    for (size_t i = 0; i < result.size(); i++)
    {
      if (Helper::StartsWith(result, "Rn", i))
      {
        result.replace(i, 2, "(");
        i--;
      }
      else if (Helper::StartsWith(result, "Ar", i))
      {
        result.replace(i, 2, ")");
        i--;
      }
      else if (Helper::StartsWith(result, "Y", i))
      {
        result.replace(i, 1, ",");
      }
    }
    return result;
  }

  std::pair<std::string, int> Solve(const std::string& molecule,
                                    const std::multimap<std::string, std::string>& conversions,
                                    std::map<std::string, std::pair<std::string, int>>& memoization);

  void Solve(std::string& molecule,
             int& steps,
             size_t& offset,
             size_t pos,
             size_t size,
             const std::multimap<std::string, std::string>& conversions,
             std::map<std::string, std::pair<std::string, int>>& memoization)
  {
    std::string toReduce = molecule.substr(pos, size);
    auto [solveStr, solveSteps] = Solve(toReduce, conversions, memoization);
    Helper::Replace(molecule, toReduce.size(), solveStr, pos);
    steps += solveSteps;
    offset -= toReduce.size() - solveStr.size();
  }

  void SolveRnAr(std::string& molecule,
                 int& steps,
                 const std::multimap<std::string, std::string>& conversions,
                 std::map<std::string, std::pair<std::string, int>>& memoization)
  {
    // Solve all parts between Rn and Ar
    int stack = 0;
    size_t position = 0;
    for (size_t i = 0; i < molecule.size(); i++)
    {
      if (Helper::StartsWith(molecule, "Rn", i))
      {
        if (stack == 0)
          position = i + 2;
        i++;
        stack++;
      }
      else if (Helper::StartsWith(molecule, "Ar", i))
      {
        stack--;
        if (stack == 0)
        {
          Solve(molecule, steps, i, position, i - position, conversions, memoization);
        }
        i++;
      }
    }
  }

  bool SolveY(std::string& molecule,
              int& steps,
              const std::multimap<std::string, std::string>& conversions,
              std::map<std::string, std::pair<std::string, int>>& memoization)
  {
    // Reduce all parts split by Y
    int stack = 0;
    size_t position = 0;
    bool hasY = false;
    for (size_t i = 0; i < molecule.size(); i++)
    {
      if (Helper::StartsWith(molecule, "Rn", i))
        stack++;
      else if (Helper::StartsWith(molecule, "Ar", i))
        stack--;

      if (Helper::StartsWith(molecule, "Y", i))
      {
        if (stack != 0)
          continue;

        hasY = true;
        Solve(molecule, steps, i, position, i - position, conversions, memoization);
        position = i + 1;
      }
    }
    if (hasY)
    {
      size_t i = 0;
      Solve(molecule, steps, i, position, molecule.size(), conversions, memoization);
      return true;
    }
    return false;
  }

  void SolveAr(std::string& molecule,
               int& steps,
               const std::multimap<std::string, std::string>& conversions,
               std::map<std::string, std::pair<std::string, int>>& memoization)
  {
    // Reduce all parts ending with Ar (except the last one)
    for (size_t i = 0; i < molecule.size(); i++)
    {
      if (Helper::StartsWith(molecule, "Ar", i) && i != molecule.size() - 2)
      {
        Solve(molecule, steps, i, 0, i + 2, conversions, memoization);
      }
    }
  }

  std::pair<std::string, int> SolveToAtom(const std::string& molecule,
                                          const std::multimap<std::string, std::string>& conversions,
                                          std::map<std::string, std::pair<std::string, int>>& memoization)
  {
    if (IsAtom(molecule))
      return {molecule, 0};

    auto it = memoization.find(molecule);
    if (it != memoization.end())
      return it->second;

    int steps = std::numeric_limits<int>::max();
    std::string result = ".";
    for (const auto& [from, to] : conversions)
    {
      std::vector<size_t> positions = FindAll(molecule, from);
      for (const auto& pos : positions)
      {
        std::string moleculeCpy = molecule;
        moleculeCpy.replace(pos, from.size(), to);
        auto [solveStr, solveSteps] = SolveToAtom(moleculeCpy, conversions, memoization);
        if (solveStr == ".")
          continue;
        if (solveSteps < steps)
        {
          steps = solveSteps;
          result = solveStr;
        }
      }
    }

    memoization[molecule] = {result, steps + 1};
    return {result, steps + 1};
  }

  void SolveToAtom(std::string& molecule,
                   int& steps,
                   const std::multimap<std::string, std::string>& conversions,
                   std::map<std::string, std::pair<std::string, int>>& memoization)
  {
    // Reduce the rest
    auto [solveStr, solveSteps] = SolveToAtom(molecule, conversions, memoization);
    if (solveStr == ".")
    {
      std::cout << "No solution found for: " << molecule << std::endl;
    }
    Helper::Replace(molecule, molecule.size(), solveStr, 0);
    steps += solveSteps;
  }

  std::pair<std::string, int> Solve(const std::string& molecule,
                                    const std::multimap<std::string, std::string>& conversions,
                                    std::map<std::string, std::pair<std::string, int>>& memoization)
  {
    if (IsAtom(molecule))
      return {molecule, 0};

    std::string result = molecule;
    int steps = 0;

    SolveRnAr(result, steps, conversions, memoization);
    if (SolveY(result, steps, conversions, memoization))
      return {result, steps};
    SolveAr(result, steps, conversions, memoization);
    SolveToAtom(result, steps, conversions, memoization);

    return {result, steps};
  }

  OUTPUT1(input)
  {
    std::set<std::string> uniques;
    for (const auto& [from, to] : input.conversions)
    {
      std::vector<size_t> positions = FindAll(input.molecule, from);
      for (const auto& pos : positions)
      {
        std::string str = input.molecule;
        Helper::Replace(str, from.size(), to, pos);
        uniques.emplace(str);
      }
    }
    return uniques.size();
  }

  OUTPUT2(input)
  {
    std::multimap<std::string, std::string> reverse;
    for (const auto& [from, to] : input.conversions)
    {
      reverse.emplace(to, from);
    }

    std::map<std::string, std::pair<std::string, int>> memoization;
    return Solve(input.molecule, reverse, memoization).second;
  }
}