#include <algorithm>
#include <deque>
#include <fstream>
#include <iostream>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using Rules =
    std::unordered_map<std::string, std::pair<std::string, std::string>>;
using Growth =
    std::unordered_map<std::string, std::unordered_map<char, uint64_t>>;

std::pair<std::string, Rules> getRules(const std::string &file_name) {
    Rules rules{};
    std::string input{};

    std::ifstream file(file_name);
    std::string str;
    bool isFold = false;
    std::getline(file, input);
    std::getline(file, str);
    while (std::getline(file, str)) {
        auto in = str.substr(0, 2);
        auto mid = str.substr(6, 1);
        rules[in] = { in[0] + mid, mid + in[1] };
    }
    return { input, rules };
}

uint64_t calculateDifferenceAfterRounds(const Rules &rules,
                                        const std::string &text, int rounds) {
    std::unordered_map<std::string, uint64_t> pair_count{};
    for (auto &rule : rules) {
        pair_count[rule.first] = 0;
    }
    for (int i = 0; i < text.size() - 1; i++) {
        pair_count[text.substr(i, 2)] += 1;
    }
    for (int i = 0; i < rounds; i++) {
        std::unordered_map<std::string, uint64_t> new_pair_count{};
        for (auto &rule : rules) {
            new_pair_count[rule.first] = 0;
        }
        for (auto &pair : pair_count) {
            auto new_pairs = rules.at(pair.first);
            new_pair_count[new_pairs.first] += pair.second;
            new_pair_count[new_pairs.second] += pair.second;
        }
        pair_count = new_pair_count;
    }

    std::unordered_map<char, uint64_t> char_count{};
    for (auto &pair : pair_count) {
        char c = pair.first[0];
        if (char_count.find(c) == char_count.end()) {
            char_count[c] = pair.second;
        } else {
            char_count[c] += pair.second;
        }
    }
    char_count[text.back()] += 1;

    uint64_t min = -1, max = 0;
    for (auto &c : char_count) {
        if (c.second < min) {
            min = c.second;
        }
        if (c.second > max) {
            max = c.second;
        }
    }
    return max - min;
}

uint64_t part1(Rules &rules, const std::string &text) {
    return calculateDifferenceAfterRounds(rules, text, 10);
}

uint64_t part2(Rules &rules, const std::string &text) {
    return calculateDifferenceAfterRounds(rules, text, 40);
}

int main(int argc, char **argv) {
    if (argc < 2) {
        std::cerr << "You must provide input file!" << std::endl;
        return 1;
    }
    auto text_and_rules = getRules(argv[1]);
    auto &text = text_and_rules.first;
    auto &rules = text_and_rules.second;

    std::cout << "The difference between the most common and least common "
                 "letters after 10 rounds is \033[91;1m"
              << part1(rules, text) << "\033[0m." << std::endl;

    std::cout << "The difference between the most common and least common "
                 "letters after 40 rounds is \033[91;1m"
              << part2(rules, text) << "\033[0m." << std::endl;
    return 0;
}