main.cpp

#include "Utils.h"
#include "benchmark.hpp"
#include "image_io.hpp"
#include "Mask.h"
#include "CPUConvolution.h"
#include "CUDABasicConvolution.h"
#include "CUDATiledInputConvolution.h"
#include "CUDATiledOutputConvolution.h"
#include "Globals.h"

#include <algorithm>
#include <cctype>
#include <chrono>
#include <ctime>
#include <exception>
#include <filesystem>
#include <iomanip>
#include <iostream>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

// ─── Variant descriptor ─────────────────────────────────────────────

struct Variant {
    std::string name;
    IConvolution* conv;
    bool enabled;
};

struct InputImageCase {
    std::string name;
    std::string path;
    GrayImage image;
};

static std::string make_run_timestamp() {
    const auto now = std::chrono::system_clock::now();
    const std::time_t tt = std::chrono::system_clock::to_time_t(now);
    std::tm tm{};
    localtime_r(&tt, &tm);

    std::ostringstream oss;
    oss << std::put_time(&tm, "%Y%m%d_%H%M%S");
    return oss.str();
}

// ─── Main ────────────────────────────────────────────────────────────

int main() {
    std::cout << "=== CUDA Convolution Benchmark Setup ===\n";

    // ── Image source ──
    const std::string images_dir = Utils::prompt_string("Images directory", "images");
    const std::filesystem::path images_path(images_dir);
    if (!std::filesystem::exists(images_path) || !std::filesystem::is_directory(images_path)) {
        throw std::runtime_error("Images directory not found: " + images_dir);
    }

    std::vector<std::filesystem::path> image_paths;
    for (const auto& entry : std::filesystem::directory_iterator(images_path)) {
        if (!entry.is_regular_file()) continue;
        image_paths.push_back(entry.path());
    }
    std::sort(image_paths.begin(), image_paths.end());

    std::vector<InputImageCase> image_cases;
    std::unordered_map<std::string, int> image_name_counts;
    for (const auto& path : image_paths) {
        try {
            GrayImage image = load_image(path.string());
            const std::string base_name = "image" + std::to_string(image.width) + "x" + std::to_string(image.height);
            int& count = image_name_counts[base_name];
            ++count;
            const std::string image_name = (count == 1) ? base_name : (base_name + "_" + std::to_string(count));
            image_cases.push_back({image_name, path.string(), std::move(image)});
        } catch (const std::exception& ex) {
            std::cerr << "[SKIP IMAGE] " << path.string() << " -> " << ex.what() << "\n";
        }
    }
    if (image_cases.empty()) {
        throw std::runtime_error("No supported images found in: " + images_dir);
    }

    // ── Mask selection: specific name or "all" ──
    const std::string mask_input = Utils::prompt_string(
        "Mask name (blur3|blur5|blur9|blur17|gaussian3|gaussian5|gaussian9|gaussian17|identity|sharpen|edge|emboss) or 'all'",
        "all");

    std::vector<MaskDef> masks_to_run;
    if (mask_input == "all") {
        masks_to_run = all_masks();
    } else {
        masks_to_run.push_back(find_mask(mask_input));
    }

    // ── Runs ──
    const int runs = Utils::prompt_int("Measured runs", 5, 1);
    const int warmup = Utils::prompt_int("Warm-up runs", 1, 0);

    // ── Implementations ──
    const bool run_cpu = Utils::prompt_yes_no("Run CPU sequential? (needed for correctness check & speedup)", false);
    const bool run_basic = Utils::prompt_yes_no("Run CUDA basic global?", true);
    const bool run_tiled_input = Utils::prompt_yes_no("Run CUDA tiled input (A1)?", true);
    const bool run_tiled_output = Utils::prompt_yes_no("Run CUDA tiled output (A2)?", true);

    // ── Execution plan ──
    std::cout << "\n=== Execution plan ===\n";
    std::cout << "Images directory: " << images_dir << "\n";
    std::cout << "Images to run (" << image_cases.size() << "):\n";
    for (const auto& image_case : image_cases) {
        std::cout << " - " << image_case.name
                  << " <- " << image_case.path
                  << " (" << image_case.image.width << "x" << image_case.image.height << ")\n";
    }
    std::cout << "Measured runs: " << runs << ", warm-up runs: " << warmup << "\n";
    std::cout << "Implementations:";
    if (run_cpu) std::cout << " CPU";
    if (run_basic) std::cout << " Basic";
    if (run_tiled_input) std::cout << " TiledInput";
    if (run_tiled_output) std::cout << " TiledOutput";
    
    std::cout << "\nMasks to run: ";
    for (size_t i = 0; i < masks_to_run.size(); ++i) {
        std::cout << masks_to_run[i].name << "(" << masks_to_run[i].width << "x" << masks_to_run[i].width << ")";
        if (i + 1 < masks_to_run.size()) std::cout << ", ";
    }
    std::cout << "\n\n";

    std::filesystem::create_directories("results");
    std::filesystem::path run_dir = std::filesystem::path("results") / ("run_" + make_run_timestamp());
    int suffix = 2;
    while (std::filesystem::exists(run_dir)) {
        run_dir = std::filesystem::path("results") / ("run_" + make_run_timestamp() + "_" + std::to_string(suffix));
        ++suffix;
    }
    std::filesystem::create_directories(run_dir);
    std::cout << "Run output dir: " << run_dir.string() << "\n";
    std::vector<BenchmarkRow> all_rows;

    // ── Instantiations ──
    CPUConvolution cpu_impl;
    CudaBasicConvolution basic_impl;

    std::vector<std::unique_ptr<CudaTiledInputConvolution>> tiled_input_impls;
    std::vector<std::unique_ptr<CudaTiledOutputConvolution>> tiled_output_impls;

    for (int tw : TILE_WIDTHS) {
        tiled_input_impls.push_back(std::make_unique<CudaTiledInputConvolution>(tw));
        tiled_output_impls.push_back(std::make_unique<CudaTiledOutputConvolution>(tw));
    }

    std::vector<Variant> gpu_variants;
    gpu_variants.push_back({"cuda_basic_global", &basic_impl, run_basic});

    for (auto& impl : tiled_input_impls) {
        const int tw = impl->getTileWidth();
        gpu_variants.push_back({"cuda_tiled_input_tw" + std::to_string(tw),
                                impl.get(), run_tiled_input});
    }

    for (auto& impl : tiled_output_impls) {
        const int tw = impl->getTileWidth();
        gpu_variants.push_back({"cuda_tiled_output_tw" + std::to_string(tw),
                                impl.get(), run_tiled_output});
    }

    for (const auto& image_case : image_cases) {
        const int W = image_case.image.width;
        const int H = image_case.image.height;
        const std::vector<float>& input_pixels = image_case.image.pixels;
        const std::filesystem::path image_dir = run_dir / image_case.name;
        std::filesystem::create_directories(image_dir);

        std::cout << "\n============================================================\n";
        std::cout << "[IMAGE START] " << image_case.name << " -> " << image_case.path
                  << " (" << W << "x" << H << ")\n";

        for (const MaskDef& mask_def : masks_to_run) {
            const std::string& mask_name = mask_def.name;
            const std::vector<float>& mask = mask_def.data;
            const int mask_width = mask_def.width;
            std::vector<BenchmarkRow> test_rows;

            std::cout << "------------------------------------------------------------\n";
            std::cout << "[MASK START] " << mask_name << "\n";
            std::cout << "[MASK INFO] width=" << mask_width << "x" << mask_width << "\n";

            std::vector<float> cpu_out;
            std::vector<double> cpu_samples;
            TimingStats cpu_stats{};

            if (run_cpu) {
                cpu_samples = Utils::run_variant_measure(cpu_impl, "cpu_sequential",
                                                         input_pixels, W, H, mask, mask_width, warmup, runs);
                cpu_impl.apply(input_pixels, cpu_out, W, H, mask, mask_width);
                cpu_stats = compute_stats(cpu_samples);
                test_rows.push_back({"cpu_sequential", image_case.name, mask_name, W, H, mask_width, runs, cpu_stats, 1.0});
                all_rows.push_back(test_rows.back());
            } else {
                std::cout << "[CPU] Skipped\n";
            }

            for (auto& v : gpu_variants) {
                if (!v.enabled) continue;

                try {
                    std::vector<double> samples = Utils::run_variant_measure(
                        *v.conv, v.name, input_pixels, W, H, mask, mask_width, warmup, runs);

                    std::vector<float> gpu_out;
                    v.conv->apply(input_pixels, gpu_out, W, H, mask, mask_width);

                    if (run_cpu) {
                        int mismatch = 0;
                        float max_err = 0.0f;
                        Utils::compare_outputs(cpu_out, gpu_out, 1e-3f, mismatch, max_err);
                        std::cout << "[" << image_case.name << "][" << mask_name << "] " << v.name
                                  << " mismatches=" << mismatch
                                  << " max_abs_error=" << max_err
                                  << (mismatch == 0 ? " PASS" : " FAIL") << "\n";
                    }

                    const TimingStats stats = compute_stats(samples);
                    const double speedup = (run_cpu && stats.mean_ms > 0)
                                               ? (cpu_stats.mean_ms / stats.mean_ms) : 0.0;

                    test_rows.push_back({v.name, image_case.name, mask_name, W, H, mask_width, runs, stats, speedup});
                    all_rows.push_back(test_rows.back());

                    try {
                        save_pgm((image_dir / ("output_" + v.name + "_" + mask_name + ".pgm")).string(),
                                 GrayImage{W, H, gpu_out});
                    } catch (const std::exception& ex) {
                        std::cerr << "Failed to write PGM for " << v.name << ": " << ex.what() << "\n";
                    }
                } catch (const std::invalid_argument& ex) {
                    std::cout << "[SKIP] " << v.name << " mask=" << mask_name << ": " << ex.what() << "\n";
                }
            }

            if (run_cpu) {
                try {
                    save_pgm((image_dir / ("output_cpu_" + mask_name + ".pgm")).string(),
                             GrayImage{W, H, cpu_out});
                } catch (const std::exception& ex) {
                    std::cerr << "Failed to write CPU PGM: " << ex.what() << "\n";
                }
            }

            std::cout << "\nTiming summary (" << image_case.name << ", " << mask_name << ", "
                      << W << "x" << H << ", ms):\n";
            for (const auto& r : test_rows) {
                std::cout << " - " << std::left << std::setw(35) << r.variant
                          << " mean=" << std::fixed << std::setprecision(3) << r.stats.mean_ms
                          << " min=" << r.stats.min_ms
                          << " max=" << r.stats.max_ms
                          << " stddev=" << r.stats.stddev_ms;
                if (r.speedup_vs_cpu > 0.0 && r.variant != "cpu_sequential") {
                    std::cout << " speedup_vs_cpu=" << std::setprecision(2) << r.speedup_vs_cpu << "x";
                }
                std::cout << "\n";
            }

            std::cout << "Pairwise comparisons (" << image_case.name << ", " << mask_name << "):\n";
            std::vector<const BenchmarkRow*> gpu_ptrs;
            for (const auto& r : test_rows) {
                if (r.variant != "cpu_sequential") gpu_ptrs.push_back(&r);
            }
            for (size_t i = 0; i + 1 < gpu_ptrs.size(); ++i) {
                for (size_t j = i + 1; j < gpu_ptrs.size(); ++j) {
                    std::cout << " - " << gpu_ptrs[i]->variant << " vs " << gpu_ptrs[j]->variant
                              << " speedup = " << std::fixed << std::setprecision(2)
                              << (gpu_ptrs[i]->stats.mean_ms / gpu_ptrs[j]->stats.mean_ms) << "x\n";
                }
            }

            std::cout << "[MASK END] " << mask_name << "\n";
        }

        std::cout << "[IMAGE OUTPUT] " << image_dir.string() << "\n";
        std::cout << "[IMAGE END] " << image_case.name << "\n\n";
    }

    write_timings_csv((run_dir / "timings.csv").string(), all_rows);
    std::cout << "[TIMINGS OUTPUT] " << (run_dir / "timings.csv").string() << "\n";
    std::cout << "[RUN OUTPUT END] " << run_dir.string() << "\n";
    std::cout << "=== Benchmark execution completed ===\n";
    return 0;
}