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

int getTypePrice(char type) {
    switch(type) {
        case 'A':
            return 1;
        case 'B':
            return 10;
        case 'C':
            return 100;
        case 'D':
            return 1000;
    }
    return 0;
}

std::vector<std::vector<char>> getDungeon(const std::string &file_name) {
    std::ifstream file(file_name);
    std::vector<std::vector<char>> result{};
    result.resize(3);
    for(int i = 0; i < 3; i++) {
        result[i].resize(11);
    }
    for(int j = 0; j < 11; j++) {
        result[0][j] = '.';
    }
    for(int i = 1; i < 3; i++) {
        for(int j = 0; j < 11; j++) {
            result[i][j] = '#';
        }
    }
    int tmp = 0;
    std::string str;
    std::getline(file, str);
    std::getline(file, str);
    auto index = 1;
    for(int i = 0; i < 2; i++) {
        std::getline(file, str);
        std::stringstream ss(str);
        result[index][2] = str[3];
        result[index][4] = str[5];
        result[index][6] = str[7];
        result[index][8] = str[9];
        index++;
    }
    return result;
}

bool canGoToSpace(const std::pair<int,int> &choice, const std::pair<int,int> &space, const std::vector<std::vector<char>> &dungeon) {
    for(int i = choice.first - 1; i >= space.first; i--) {
        if(dungeon[i][choice.second] != '.') {
            return false;
        }
    }
    auto addition = 1;
    if(choice.second > space.second) {
        addition = -1;
    }
    for(int i = choice.second + addition; i != space.second; i += addition) {
        if(dungeon[0][i] != '.') {
            return false;
        }
    }
    return dungeon[0][space.second] == '.';
}

bool isHome(const std::pair<int,int> &choice, const std::vector<std::vector<char>> &dungeon, int height) {
    int index = -1;
    switch(dungeon[choice.first][choice.second]) {
        case 'A':
            index = 2;
            break;
        case 'B':
            index = 4;
            break;
        case 'C':
            index = 6;
            break;
        case 'D':
            index = 8;
            break;
        default:
            return false;
    }
    if(choice.second != index) {
        return false;
    }
    for(int i = 1; i < height; i++) {
        if(dungeon[i][index] != dungeon[choice.first][choice.second] && dungeon[i][index] != '.') {
            return false;
        }
    }
    return dungeon[height][index] == dungeon[choice.first][choice.second];
}

int getHomeIndex(char type) {
    switch(type) {
        case 'A':
            return 2;
        case 'B':
            return 4;
        case 'C':
            return 6;
        case 'D':
            return 8;
        default:
            return -1;
    }
}

bool canGoHome(int index, const std::vector<std::vector<char>> &dungeon, int height) {
    auto target = getHomeIndex(dungeon[0][index]);
    if(dungeon[1][target] != '.') {
        return false;
    } // require at least 1 vacant space
    for(int i = 2; i <= height; i++) {
        if(dungeon[i][target] != '.' && dungeon[i][target] != dungeon[0][index]) {
            return false;
        }
    }
    auto addition = 1;
    if(target < index) {
        addition = -1;
    }
    for(int i = index + addition; i != target; i += addition) {
        if(dungeon[0][i] != '.') {
            return false;
        }
    }
    return dungeon[0][target] == '.';
}

int goHome(int index, std::vector<std::vector<char>> &dungeon, int height) {
    auto target = getHomeIndex(dungeon[0][index]);
    auto target_y = height;
    while(dungeon[target_y][target] != '.') {
        target_y -= 1;
    }
    dungeon[target_y][target] = dungeon[0][index];
    dungeon[0][index] = '.';
    return (abs(target - index) + target_y) * getTypePrice(dungeon[target_y][target]);
}

bool allHome(const std::vector<std::vector<char>> &dungeon, int height) {
    for(int i = 1; i <= height; i++) {
        if(dungeon[i][2] != 'A') {
            return false;
        }
    }
    for(int i = 1; i <= height; i++) {
        if(dungeon[i][4] != 'B') {
            return false;
        }
    }
    for(int i = 1; i <= height; i++) {
        if(dungeon[i][6] != 'C') {
            return false;
        }
    }
    for(int i = 1; i <= height; i++) {
        if(dungeon[i][8] != 'D') {
            return false;
        }
    }
    return true;
}

std::string getCacheIndex(const std::vector<std::vector<char>> &dungeon, std::pair<int,int> choice, int height) {
    std::string index = "";
    for(auto &c : dungeon[0]) {
        index += c;
    }
    for(int i = 1; i <= height; i++) {
        for(int j = 2; j <= 8; j+=2) {
            index += dungeon[i][j];
        }
    }
    index += std::to_string(choice.first);
    index += std::to_string(choice.second);
    return index;
}

int performMovement(const std::vector<std::vector<char>> &dungeon, std::pair<int,int> choice, std::unordered_map<std::string, int64_t> &cache, int height) {
    auto cache_index = getCacheIndex(dungeon, choice, height);
    if(cache.find(cache_index) != cache.end() && cache[cache_index] != -1) {
        return cache[cache_index];
    }
    std::vector<std::pair<int,int>> possibleSpaces = {{0,0}, {0,1}, {0,3}, {0,5}, {0,7}, {0,9}, {0,10}};
    int64_t minPrice = -1;
    for(auto &space : possibleSpaces) {
        auto dungeoncopy = dungeon;
        if(isHome(choice, dungeoncopy, height) || !canGoToSpace(choice, space, dungeoncopy)) {
            continue;
        }
        dungeoncopy[space.first][space.second] = dungeoncopy[choice.first][choice.second];
        dungeoncopy[choice.first][choice.second] = '.';
        auto price = (abs(space.second - choice.second) + choice.first) * getTypePrice(dungeoncopy[space.first][space.second]);
        for(int i = 0; i < 7; i++) {
            for(int j = 0; j < 11; j++) {
                if(canGoHome(j, dungeoncopy, height)) {
                    price += goHome(j, dungeoncopy, height);
                }
            }
        }
        if(allHome(dungeoncopy, height)) {
            if(price < minPrice || minPrice == -1) {
                minPrice = price;
            }
            continue;
        }
        for(int j = 1; j <= height; j++) {
            for(int i = 2; i <= 8; i+=2) {
                if(dungeoncopy[j][i] != '.') {
                    auto nextPrice = price;
                    auto add = performMovement(dungeoncopy, {j,i}, cache, height);
                    if(add == -1) {
                        continue;
                    }
                    nextPrice += add;
                    if(nextPrice < minPrice || minPrice == -1) {
                        minPrice = nextPrice;
                    }
                }
            }
        }
    }
    cache[cache_index] = minPrice;
    return minPrice;
}

int part1(const std::vector<std::vector<char>> &input) {
    std::unordered_map<std::string, int64_t> cache{};
    auto min = -1;
    for(int i = 2; i <= 8; i+=2) {
        auto price = performMovement(input, {1,i}, cache, 2);
        if((price < min || min == -1) && price != -1) {
            min = price;
        }
    }
    return min;
}

int part2(const std::vector<std::vector<char>> &input) {
    std::unordered_map<std::string, int64_t> cache{};
    auto inputCopy = input;
    inputCopy.resize(5);
    inputCopy[3].resize(11);
    inputCopy[4].resize(11);
    for(int i = 0; i < inputCopy[0].size(); i++) {
        inputCopy[4][i] = inputCopy[2][i];
    }
    inputCopy[2] = {'#', '#', 'D', '#', 'C', '#', 'B', '#', 'A', '#', '#'};
    inputCopy[3] = {'#', '#', 'D', '#', 'B', '#', 'A', '#', 'C', '#', '#'};
    auto min = -1;
    for(int i = 2; i <= 8; i+=2) {
        auto price = performMovement(inputCopy, {1,i}, cache, 4);
        if((price < min || min == -1) && price != -1) {
            min = price;
        }
    }
    return min;
}

int main(int argc, char **argv) {
    if (argc < 2) {
        std::cerr << "You must provide input file!" << std::endl;
        return 1;
    }
    auto dungeon = getDungeon(argv[1]);
    std::cout << "The lowest energy price to sort amphipods is \033[91;1m" << part1(dungeon)
              << "\033[0m." << std::endl;
    std::cout << "The lowest energy price to sort amphipods (with the full map this time) is \033[91;1m" << part2(dungeon)
              << "\033[0m." << std::endl;
    return 0;
}