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

class Image {
public:
    Image(size_t width, size_t height) {
        image.resize(height);
        for(auto &x : image) {
            x.resize(width);
        }
    }
    Image(const std::vector<std::vector<bool>> &input) : image(std::move(input)) {}

    void setPixel(size_t x, size_t y, bool pixel) {
        image[y][x] = pixel;
    }
    bool getPixel(size_t x, size_t y) const {
        return image[y][x];
    }
    void setOutsidePixelState(bool pixel) {
        outside = pixel;
    }
    bool getOutsidePixelState() const {
        return outside;
    }
    size_t getImageWidth() const {
        return image[0].size();
    }
    size_t getImageHeight() const {
        return image.size();
    }
    size_t getLitPixels() const {
        size_t pixels = 0;
        for(auto &line : image) {
            for(auto pixel : line) {
                if(pixel) {
                    pixels++;
                }
            }
        }
        return pixels;
    }
    bool getNewPixel(const std::array<bool, 512> &algorithm, int64_t x, int64_t y) const {
        uint16_t index = 0;
        for(int y_add = -1; y_add <= 1; y_add++) {
            for(int x_add = -1; x_add <= 1; x_add++) {
                index <<= 1;
                auto new_y = y + y_add;
                auto new_x = x + x_add;
                if(new_y < 0 || new_x < 0 || new_y >= getImageHeight() || new_x >= getImageWidth()) {
                    if(getOutsidePixelState()) {
                        index |= 0b1;
                    }
                } else {
                    if(getPixel(new_x, new_y)) {
                        index |= 0b1;
                    }
                }
            }
        }
        return algorithm[index];
    }
    void printImage() const {
        for(auto &line : image) {
            for(auto pixel : line) {
                std::cout << (pixel ? '#' : '.');
            }
            std::cout << std::endl;
        }
    }
private:
    bool outside = false;
    std::vector<std::vector<bool>> image;
};

std::pair<std::array<bool, 512>, Image> getInput(const std::string &file_name) {
    std::ifstream file(file_name);
    std::array<bool, 512> algorithm{};
    std::vector<std::vector<bool>> image_data{};
    int tmp = 0;
    std::string str;
    std::getline(file, str);
    for(size_t i = 0; i < 512; i++) {
        if(str[i] == '#') {
            algorithm[i] = true;
        } else {
            algorithm[i] = false;
        }
    }
    std::getline(file, str);
    while (std::getline(file, str)) {
        image_data.emplace_back();
        for(auto &c : str) {
            if(c == '#') {
                image_data.back().push_back(true);
            } else {
                image_data.back().push_back(false);
            }
        }
    }
    auto image = Image(image_data);
    return {algorithm, image};
}

Image applyAlgorithm(const std::array<bool, 512> &algorithm, const Image &startingImage) {
    Image newImage(startingImage.getImageWidth() + 2, startingImage.getImageHeight() + 2);
    for(int64_t x = 0; x < newImage.getImageWidth(); x++) {
        for(int64_t y = 0; y < newImage.getImageHeight(); y++) {
            auto startingX = x - 1;
            auto startingY = y - 1;
            auto nextPixel = startingImage.getNewPixel(algorithm, startingX, startingY);
            newImage.setPixel(x, y, nextPixel);
        }
    }
    newImage.setOutsidePixelState(algorithm[startingImage.getOutsidePixelState() ? 511 : 0]);
    return newImage;
}

size_t part1(const std::array<bool, 512> &algorithm, Image image) {
    for(int i = 0; i < 2; i++) {
        image = applyAlgorithm(algorithm, image);
    }
    return image.getLitPixels();
}

size_t part2(const std::array<bool, 512> &algorithm, Image image) {
    for(int i = 0; i < 50; i++) {
        image = applyAlgorithm(algorithm, image);
    }
    return image.getLitPixels();
}

int main(int argc, char **argv) {
    if (argc < 2) {
        std::cerr << "You must provide input file!" << std::endl;
        return 1;
    }
    auto input = getInput(argv[1]);
    std::cout << "There are \033[91;1m" << part1(input.first, input.second)
              << "\033[0m lit pixels after 2 applies." << std::endl;
    std::cout << "There are \033[91;1m" << part2(input.first, input.second)
              << "\033[0m window depth increases." << std::endl;
    return 0;
}