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

std::vector<int> getFish(const std::string &file_name) {
    std::vector<int> fish;

    std::ifstream file(file_name);
    std::string str;
    std::getline(file, str);
    int tmp{};
    char tmp_char{};
    std::stringstream ss(str);
    while (ss >> tmp) {
        fish.push_back(tmp);
        ss >> tmp_char;
    }
    return fish;
}

uint64_t computeFish(std::vector<int> fish, int lifecycle, int young_addition,
                     int days) {
    // index of bucket = days until reproduction
    std::vector<uint64_t> fish_buckets{};
    for (int i = 0; i < lifecycle + young_addition; i++) {
        fish_buckets.push_back(0);
    }
    for (auto &f : fish) {
        fish_buckets[f] += 1;
    }
    for (int i = 0; i < days; i++) {
        auto new_fish = fish_buckets[0];
        for (int j = 1; j < fish_buckets.size(); j++) {
            fish_buckets[j - 1] = fish_buckets[j];
        }
        fish_buckets[lifecycle - 1] += new_fish;
        fish_buckets.back() = new_fish;
    }
    uint64_t population = 0;
    for (auto &f : fish_buckets) {
        population += f;
    }
    return population;
}

uint64_t part1(const std::vector<int> &fish) {
    return computeFish(fish, 7, 2, 80);
}

uint64_t part2(const std::vector<int> &fish) {
    return computeFish(fish, 7, 2, 256);
}

int main(int argc, char **argv) {
    if (argc < 2) {
        std::cerr << "You must provide input file!" << std::endl;
        return 1;
    }
    auto fish = getFish(argv[1]);
    std::cout << "The number of fish after 80 days is \033[91;1m" << part1(fish)
              << "\033[0m." << std::endl;
    std::cout << "The number of fish after 256 days is \033[91;1m"
              << part2(fish) << "\033[0m." << std::endl;
    return 0;
}