#include <fstream>
#include <iostream>
#include <set>
#include <sstream>
#include <vector>

std::vector<std::vector<std::set<char>>>
getSegments(const std::string &file_name) {
    std::vector<std::vector<std::set<char>>> segments;

    std::ifstream file(file_name);
    std::string str;
    while (std::getline(file, str)) {
        std::string tmp{};
        char tmp_char{};
        std::stringstream ss(str);
        std::vector<std::set<char>> segment{};
        std::vector<std::set<char>> unknown{};
        segment.resize(14);

        for (int i = 0; i < 10; i++) {
            ss >> tmp;
            int index = -1;
            if (tmp.length() == 2) {
                index = 1;
            } else if (tmp.length() == 3) {
                index = 7;
            } else if (tmp.length() == 4) {
                index = 4;
            } else if (tmp.length() == 7) {
                index = 8;
            }
            if (index != -1) {
                for (auto &c : tmp) {
                    segment[index].insert(c);
                }
            } else {
                unknown.emplace_back();
                for (auto &c : tmp) {
                    unknown.back().insert(c);
                }
            }
        }

        ss >> tmp_char;
        for (int i = 10; i < 14; i++) {
            ss >> tmp;
            for (auto &c : tmp) {
                segment[i].insert(c);
            }
        }

        while (unknown.size() > 2) {
            for (int i = 0; i < unknown.size(); i++) {
                if (unknown[i].size() == 5) {
                    // 5-seg number which contains 1 must be 3
                    if (std::includes(unknown[i].begin(), unknown[i].end(),
                                      segment[1].begin(), segment[1].end())) {
                        segment[3].swap(unknown[i]);
                        unknown.erase(unknown.begin() + i);
                        i--;
                        continue;
                    }
                } else if (unknown[i].size() == 6) {
                    // 6-seg number which does not contain 1 must be 6
                    if (!std::includes(unknown[i].begin(), unknown[i].end(),
                                       segment[1].begin(), segment[1].end())) {
                        segment[6].swap(unknown[i]);
                        unknown.erase(unknown.begin() + i);
                        i--;
                        continue;
                    }
                    // 6-seg number which contains 3 must be 9
                    if (segment[3].size() != 0 &&
                        std::includes(unknown[i].begin(), unknown[i].end(),
                                      segment[3].begin(), segment[3].end())) {
                        segment[9].swap(unknown[i]);
                        unknown.erase(unknown.begin() + i);
                        i--;
                        continue;
                    }
                    // 6-seg number which does not contain 3 and contains
                    // 1 must be 0
                    if (segment[3].size() != 0 &&
                        !std::includes(unknown[i].begin(), unknown[i].end(),
                                       segment[3].begin(), segment[3].end()) &&
                        std::includes(unknown[i].begin(), unknown[i].end(),
                                      segment[1].begin(), segment[1].end())) {
                        segment[0].swap(unknown[i]);
                        unknown.erase(unknown.begin() + i);
                        i--;
                        continue;
                    }
                }
            }
        }
        for (int i = 0; i < unknown.size(); i++) {
            // only 5 and 2 are remaining, 5 is 9 with a segment missing
            // so 9 must contain 5, but not 2
            if (std::includes(segment[9].begin(), segment[9].end(),
                              unknown[i].begin(), unknown[i].end())) {
                segment[5].swap(unknown[i]);
                unknown.erase(unknown.begin() + i);
                break;
            }
        }
        // 2 is the only segment missing
        segment[2].swap(unknown[0]);
        segments.push_back(std::move(segment));
    }
    return segments;
}

uint64_t decodeDigits(const std::vector<std::set<char>> &segment) {
    uint64_t result = 0;
    for (int i = 10; i < 14; i++) {
        for (int j = 0; j < 10; j++) {
            if (segment[i] == segment[j]) {
                result *= 10;
                result += j;
                break;
            }
        }
    }
    return result;
}

uint64_t part1(const std::vector<std::vector<std::set<char>>> &segments) {
    uint64_t count = 0;
    for (auto &segment : segments) {
        for (int i = 10; i < 14; i++) {
            if (segment[i] == segment[1] || segment[i] == segment[4] ||
                segment[i] == segment[7] || segment[i] == segment[8]) {
                count++;
            }
        }
    }
    return count;
}

uint64_t part2(const std::vector<std::vector<std::set<char>>> &segments) {
    uint64_t sum = 0;
    for (auto &segment : segments) {
        sum += decodeDigits(segment);
    }
    return sum;
}

int main(int argc, char **argv) {
    if (argc < 2) {
        std::cerr << "You must provide input file!" << std::endl;
        return 1;
    }
    auto segments = getSegments(argv[1]);
    std::cout << "There are \033[91;1m" << part1(segments)
              << "\033[0m instances of 1, 4, 7, or 8." << std::endl;
    std::cout << "There sum of all decoded entries is \033[91;1m"
              << part2(segments) << "\033[0m." << std::endl;
    return 0;
}