#include <algorithm>
#include <array>
#include <charconv>
#include <cstddef>
#include <expected>
#include <format>
#include <fstream>
#include <functional>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <mdspan>
#include <numeric>
#include <print>
#include <ranges>
#include <set>
#include <sstream>
#include <string>
#include <string_view>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

struct Diagram {
    std::vector<char> data;
    size_t width;
    size_t height;
    size_t start_location;

    using GridView = std::mdspan<char, std::dextents<size_t, 2>>;

    auto grid() { return GridView(data.data(), height, width); }
};
struct CoordinateHash {
    std::size_t operator()(const std::pair<long long, long long> &p) const {

        auto h1 = std::hash<long long>{}(p.first);
        auto h2 = std::hash<long long>{}(p.second);

        return h1 ^ (h2 + 0x9e3779b9 + (h1 << 6) + (h1 >> 2));
    }
};

auto printDiagram(Diagram &diagram) {
    auto grid = diagram.grid();
    for (auto row = 0UZ; row != grid.extent(0); row++) {
        for (auto col = 0UZ; col != grid.extent(1); col++) {
            std::print("{}", grid[row, col]);
        }
        std::println();
    }
    // std::println("Map: {}", diagram.data);
    std::println("Width: {}", diagram.width);
    std::println("Height: {}", diagram.height);
}

static auto parseInput(const std::string &filename) -> std::expected<Diagram, std::string> {
    std::ifstream inputF{filename};

    if (!inputF) {
        return std::unexpected{"Some file open error."};
    }

    Diagram diagram{};

    std::string puzzleLine;
    while (std::getline(inputF, puzzleLine)) {
        if (!puzzleLine.empty()) {
            diagram.data.append_range(puzzleLine |
                                      std::views::transform([](char c) -> char { return c == '.' ? ' ' : c; }));
            diagram.height++;
            diagram.width = puzzleLine.length();
        }
    }
    // printDiagram(diagram);
    return diagram;
}

static auto decorateTree(Diagram &diagram) -> long long {
    auto grid = diagram.grid();
    long long count{};
    for (auto row = 0UZ; row != grid.extent(0); row++) {
        for (auto col = 0UZ; col != grid.extent(1); col++) {
            auto &curChar = grid[row, col];

            if (curChar == 'S') {
                diagram.start_location = col;
                grid[row + 1, col] = '|';
            } else if (curChar == '^' && grid[row - 1, col] == '|') {
                grid[row, col - 1] = '|';
                grid[row, col + 1] = '|';
                count++;

            } else if (grid[row - 1, col] == '|') {
                grid[row, col] = '|';
            }
        }
    }
    // printDiagram(diagram);

    return count;
}

using Coordinate = std::pair<int, int>;
using SparseGrid = std::unordered_map<std::pair<long long, long long>, long long, CoordinateHash>;

auto countTree(Diagram &diagram, Coordinate coord, SparseGrid &sparseGrid) -> long long {
    if (coord.first + 1 == diagram.height) {
        return 1;
    }
    auto grid = diagram.grid();
    // std::println("sparseGrid: {}", sparseGrid);
    // std::println("sparseGrid size: {}", sparseGrid.size());

    if (sparseGrid.contains(coord)) {
        // std::println("Found value: {}", sparseGrid[coord]);
        return sparseGrid[coord];
    }
    auto cur = grid[coord.first + 1, coord.second];
    // std::println("{}: Under me: {}", coord, cur);
    if (cur == '^') {
        auto res = countTree(diagram, {coord.first + 1, coord.second - 1}, sparseGrid) +
                   countTree(diagram, {coord.first + 1, coord.second + 1}, sparseGrid);
        sparseGrid[coord] = res;
        return res;
    }
    if (cur == '|') {
        auto res = countTree(diagram, {coord.first + 1, coord.second}, sparseGrid);
        sparseGrid[coord] = res;
        return res;
    }
}

static auto treePartTwo(Diagram &diagram) -> long long {
    std::println("Counting big tree...");
    SparseGrid sparseGrid;
    return countTree(diagram, {1, diagram.start_location}, sparseGrid);
}

auto main() -> int {
    auto testCase = parseInput("test_input");
    if (testCase) {

        auto testResult = decorateTree(*testCase);
        std::println("P1 Testcase result: {}", testResult);

        auto testResultP2 = treePartTwo(*testCase);
        std::println("P2 Testcase result: {}", testResultP2);
    } else {
        std::print("{}\n", testCase.error());
    }

    auto realPuzzle = parseInput("puzzle_input");
    if (realPuzzle) {
        auto realResult = decorateTree(*realPuzzle);
        std::println("P1 Real result: {}", realResult);

        auto realResultP2 = treePartTwo(*realPuzzle);
        std::println("P1 Real result: {}", realResultP2);
    } else {
        std::print("{}\n", realPuzzle.error());
    }

    return 0;
}