#include "sudoku.h"
#include <nd-wfc/wfc.hpp>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <chrono>
#include <algorithm>

// Helper function to load multiple puzzles from a file (one puzzle per line)
std::vector<Sudoku> loadPuzzlesFromFile(const std::string& filename) {
    std::vector<Sudoku> puzzles;
    std::ifstream file(filename);

    if (!file.is_open()) {
        std::cerr << "Failed to open file: " << filename << std::endl;
        return puzzles;
    }

    std::string line;
    while (std::getline(file, line)) {
        // Remove whitespace
        line.erase(std::remove_if(line.begin(), line.end(),
                   [](char c) { return std::isspace(c); }), line.end());

        if (line.empty()) continue;

        Sudoku sudoku;
        if (sudoku.loadFromString(line)) {
            puzzles.push_back(std::move(sudoku));
        }
    }

    return puzzles;
}

// Helper function to save failing puzzles to a file
void saveFailingPuzzles(const std::vector<std::string>& failingPuzzles, const std::string& outputFile) {
    std::ofstream file(outputFile);
    if (!file.is_open()) {
        std::cerr << "Failed to open output file: " << outputFile << std::endl;
        return;
    }

    for (const auto& puzzle : failingPuzzles) {
        file << puzzle << std::endl;
    }

    std::cout << "Saved " << failingPuzzles.size() << " failing puzzles to " << outputFile << std::endl;
}

int main() {
    // Create WFC solver
    auto sudokuSolver = WFC::Builder<Sudoku, uint8_t>()
        .DefineIDs<1, 2, 3, 4, 5, 6, 7, 8, 9>()
        .Variable<1, 2, 3, 4, 5, 6, 7, 8, 9>([](Sudoku&, size_t index, WFC::WorldValue<uint8_t> val, auto& constrainer) {
            size_t x = index % 9;
            size_t y = index / 9;

            // Add row constraints (same row, different columns)
            for (size_t i = 0; i < 9; ++i) {
                if (i != x) constrainer.Exclude(val, i + y * 9);
            }

            // Add column constraints (same column, different rows)
            for (size_t i = 0; i < 9; ++i) {
                if (i != y) constrainer.Exclude(val,x + i * 9);
            }

            // Add box constraints (3x3 box)
            size_t box_x = (x / 3) * 3;
            size_t box_y = (y / 3) * 3;
            for (size_t j = 0; j < 3; ++j) {
                for (size_t k = 0; k < 3; ++k) {
                    size_t cell_x = box_x + j;
                    size_t cell_y = box_y + k;
                    size_t cell_index = cell_x + cell_y * 9;
                    if (cell_index != index) {
                        constrainer.Exclude(val, cell_index);
                    }
                }
            }
        })
        .build();

    // File paths
    const std::string dataPath = "/home/connor/repos/nd-wfc/demos/sudoku/data";
    const std::vector<std::string> inputFiles = {
        dataPath + "/Sudoku_easy.txt",
        dataPath + "/Sudoku_medium.txt",
        dataPath + "/Sudoku_hard.txt",
        dataPath + "/Sudoku_diabolical.txt"
    };
    const std::string outputFile = dataPath + "/Sudoku_failing.txt";

    // Collect all failing puzzles
    std::vector<std::string> allFailingPuzzles;
    std::vector<std::pair<std::string, int>> failureStats; // filename -> count

    std::cout << "Analyzing Sudoku puzzles for solver failures..." << std::endl;
    std::cout << "=================================================" << std::endl;

    for (const auto& inputFile : inputFiles) {
        std::cout << "\nProcessing " << inputFile << "..." << std::endl;

        // Load puzzles from file
        auto puzzles = loadPuzzlesFromFile(inputFile);
        if (puzzles.empty()) {
            std::cout << "No puzzles loaded from " << inputFile << std::endl;
            continue;
        }

        std::cout << "Loaded " << puzzles.size() << " puzzles" << std::endl;

        // Test each puzzle
        int solvedCount = 0;
        int failedCount = 0;
        std::vector<std::string> failingPuzzles;

        auto start = std::chrono::high_resolution_clock::now();

        for (size_t i = 0; i < puzzles.size(); ++i) {
            auto& sudoku = puzzles[i];

            // Validate puzzle before solving
            if (!sudoku.isValid()) {
                std::cout << "Puzzle " << i << " is invalid, skipping" << std::endl;
                failedCount++;
                continue;
            }

            // Try to solve
            auto puzzleStart = std::chrono::high_resolution_clock::now();
            sudokuSolver.Run(sudoku, false);  // false = disable verbose output for speed
            auto puzzleEnd = std::chrono::high_resolution_clock::now();

            bool solved = sudoku.isSolved();

            if (solved) {
                solvedCount++;
            } else {
                failedCount++;
                // Get the original puzzle string for saving
                std::string puzzleStr = sudoku.toString();
                failingPuzzles.push_back(puzzleStr);

                auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(puzzleEnd - puzzleStart);
                std::cout << "Puzzle " << i << " failed to solve in " << duration.count() << "ms" << std::endl;
            }

            // Progress indicator for large files
            if ((i + 1) % 1000 == 0) {
                std::cout << "Progress: " << (i + 1) << "/" << puzzles.size()
                          << " (" << solvedCount << " solved, " << failedCount << " failed)" << std::endl;
            }
        }

        auto end = std::chrono::high_resolution_clock::now();
        auto totalDuration = std::chrono::duration_cast<std::chrono::seconds>(end - start);

        std::cout << "\nResults for " << inputFile << ":" << std::endl;
        std::cout << "  Total puzzles in file: " << puzzles.size() << std::endl;
        std::cout << "  Solved: " << solvedCount << std::endl;
        std::cout << "  Failed: " << failedCount << std::endl;
        std::cout << "  Success rate: " << (puzzles.size() > 0 ? (solvedCount * 100.0 / puzzles.size()) : 0) << "%" << std::endl;
        std::cout << "  Total time: " << totalDuration.count() << " seconds" << std::endl;

        // Add failing puzzles to the global list
        allFailingPuzzles.insert(allFailingPuzzles.end(), failingPuzzles.begin(), failingPuzzles.end());
        failureStats.push_back({inputFile, failedCount});
    }

    // Save all failing puzzles to output file
    std::cout << "\n=================================================" << std::endl;
    std::cout << "Analysis complete!" << std::endl;
    std::cout << "Total failing puzzles across all files: " << allFailingPuzzles.size() << std::endl;

    if (!allFailingPuzzles.empty()) {
        saveFailingPuzzles(allFailingPuzzles, outputFile);

        std::cout << "\nFailure statistics by file:" << std::endl;
        for (const auto& stat : failureStats) {
            std::cout << "  " << stat.first << ": " << stat.second << " failures" << std::endl;
        }
    } else {
        std::cout << "No failing puzzles found!" << std::endl;
    }

    return 0;
}