#include <algorithm>
#include <charconv>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>
#include <string_view>
#include <system_error>
#include <vector>

namespace {

struct Options {
    std::size_t width = 4096;
    std::size_t height = 2048;
    std::size_t tile = 32;
    std::size_t passes = 1;
    std::size_t rounds = 5;
};

struct AddressParts {
    std::vector<std::size_t> x;
    std::vector<std::size_t> y;
};

struct TimedResult {
    double milliseconds;
    std::uint64_t checksum;
};

[[noreturn]] void fail(const std::string& message) {
    throw std::runtime_error(message);
}

std::size_t parse_positive(std::string_view text, std::string_view name) {
    std::size_t value = 0;
    const char* begin = text.data();
    const char* end = begin + text.size();
    const auto result = std::from_chars(begin, end, value);
    if (result.ec != std::errc{} || result.ptr != end || value == 0) {
        fail("--" + std::string(name) + " must be a positive integer");
    }
    return value;
}

Options parse_options(int argc, char** argv) {
    Options options;

    for (int index = 1; index < argc; ++index) {
        const std::string_view argument(argv[index]);
        if (argument == "--help" || argument == "-h") {
            std::cout
                << "Usage: tiling-rotated-stencil-benchmark [options]\n\n"
                << "  --width=N    source width in pixels (default 4096)\n"
                << "  --height=N   source height in pixels (default 2048)\n"
                << "  --tile=N     square tile edge in pixels (default 32)\n"
                << "  --passes=N   rotated 3x3 reads per sample (default 1)\n"
                << "  --rounds=N   timed samples per layout (default 5)\n";
            std::exit(0);
        }

        const auto equals = argument.find('=');
        if (equals == std::string_view::npos || argument.size() < 2 ||
            argument.substr(0, 2) != "--") {
            fail("unknown argument: " + std::string(argument));
        }

        const std::string_view name = argument.substr(2, equals - 2);
        const std::size_t value = parse_positive(argument.substr(equals + 1), name);
        if (name == "width") {
            options.width = value;
        } else if (name == "height") {
            options.height = value;
        } else if (name == "tile") {
            options.tile = value;
        } else if (name == "passes") {
            options.passes = value;
        } else if (name == "rounds") {
            options.rounds = value;
        } else {
            fail("unknown option: --" + std::string(name));
        }
    }

    if (options.width < 3 || options.height < 3) {
        fail("width and height must both be at least 3");
    }
    if (options.width % options.tile != 0 || options.height % options.tile != 0) {
        fail("width and height must both be multiples of tile");
    }
    if (options.width > std::numeric_limits<std::size_t>::max() / options.height) {
        fail("image dimensions overflow size_t");
    }
    return options;
}

std::vector<std::uint32_t> make_linear_image(std::size_t pixel_count) {
    std::vector<std::uint32_t> image(pixel_count);
    for (std::size_t index = 0; index < pixel_count; ++index) {
        const auto value = static_cast<std::uint32_t>(index);
        image[index] = (value * 2654435761u) >> 16;
    }
    return image;
}

std::vector<std::uint32_t> make_tiled_image(
    const std::vector<std::uint32_t>& linear,
    const Options& options)
{
    std::vector<std::uint32_t> tiled(linear.size());
    const std::size_t tiles_x = options.width / options.tile;
    const std::size_t tiles_y = options.height / options.tile;
    const std::size_t tile_area = options.tile * options.tile;

    for (std::size_t tile_y = 0; tile_y < tiles_y; ++tile_y) {
        for (std::size_t tile_x = 0; tile_x < tiles_x; ++tile_x) {
            const std::size_t tiled_base = (tile_y * tiles_x + tile_x) * tile_area;
            const std::size_t source_x = tile_x * options.tile;
            const std::size_t source_y = tile_y * options.tile;

            for (std::size_t local_y = 0; local_y < options.tile; ++local_y) {
                const std::size_t source_row =
                    (source_y + local_y) * options.width + source_x;
                const std::size_t tiled_row = tiled_base + local_y * options.tile;
                std::copy_n(
                    linear.data() + source_row,
                    options.tile,
                    tiled.data() + tiled_row);
            }
        }
    }
    return tiled;
}

AddressParts make_linear_parts(const Options& options) {
    AddressParts parts;
    parts.x.resize(options.width);
    parts.y.resize(options.height);
    for (std::size_t x = 0; x < options.width; ++x) {
        parts.x[x] = x;
    }
    for (std::size_t y = 0; y < options.height; ++y) {
        parts.y[y] = y * options.width;
    }
    return parts;
}

AddressParts make_tiled_parts(const Options& options) {
    AddressParts parts;
    parts.x.resize(options.width);
    parts.y.resize(options.height);
    const std::size_t tiles_x = options.width / options.tile;
    const std::size_t tile_area = options.tile * options.tile;

    for (std::size_t x = 0; x < options.width; ++x) {
        parts.x[x] = (x / options.tile) * tile_area + x % options.tile;
    }
    for (std::size_t y = 0; y < options.height; ++y) {
        parts.y[y] =
            (y / options.tile) * tiles_x * tile_area
            + (y % options.tile) * options.tile;
    }
    return parts;
}

std::uint64_t sum_rotated_stencil(
    const std::vector<std::uint32_t>& image,
    const AddressParts& address,
    const Options& options,
    std::size_t passes)
{
    std::uint64_t checksum = 0;

    // A 90-degree rotation swaps the output dimensions.  Moving along an
    // output row therefore walks vertically through the source image.
    for (std::size_t pass = 0; pass < passes; ++pass) {
        for (std::size_t output_y = 1; output_y + 1 < options.width; ++output_y) {
            const std::size_t source_x = output_y;
            const std::size_t x0 = address.x[source_x - 1];
            const std::size_t x1 = address.x[source_x];
            const std::size_t x2 = address.x[source_x + 1];

            for (std::size_t output_x = 1; output_x + 1 < options.height; ++output_x) {
                const std::size_t source_y = options.height - 1 - output_x;
                const std::size_t y0 = address.y[source_y - 1];
                const std::size_t y1 = address.y[source_y];
                const std::size_t y2 = address.y[source_y + 1];

                checksum += image[y0 + x0] + image[y0 + x1] + image[y0 + x2];
                checksum += image[y1 + x0] + image[y1 + x1] + image[y1 + x2];
                checksum += image[y2 + x0] + image[y2 + x1] + image[y2 + x2];
            }
        }
    }
    return checksum;
}

template <typename Operation>
TimedResult measure(Operation&& operation) {
    const auto start = std::chrono::steady_clock::now();
    const std::uint64_t checksum = operation();
    const auto finish = std::chrono::steady_clock::now();
    const std::chrono::duration<double, std::milli> elapsed = finish - start;
    return {elapsed.count(), checksum};
}

double median(std::vector<double> values) {
    std::sort(values.begin(), values.end());
    const std::size_t middle = values.size() / 2;
    if (values.size() % 2 == 0) {
        return (values[middle - 1] + values[middle]) / 2.0;
    }
    return values[middle];
}

void print_result(
    std::string_view name,
    const std::vector<double>& samples,
    double logical_bytes_read)
{
    const double middle = median(samples);
    const double gib_per_second =
        logical_bytes_read / (middle / 1000.0) / (1024.0 * 1024.0 * 1024.0);
    std::cout << std::left << std::setw(7) << name << std::right
              << ": median " << std::fixed << std::setprecision(3) << middle << " ms, "
              << std::setprecision(2) << gib_per_second << " logical GiB/s\n";
}

} // namespace

int main(int argc, char** argv) {
    try {
        const Options options = parse_options(argc, argv);
        const std::size_t pixel_count = options.width * options.height;
        const double bytes_per_layout =
            static_cast<double>(pixel_count * sizeof(std::uint32_t));
        const std::size_t output_pixels =
            (options.width - 2) * (options.height - 2);

        std::cout << "Source: " << options.width << 'x' << options.height
                  << ", tile: " << options.tile << 'x' << options.tile
                  << ", " << std::fixed << std::setprecision(1)
                  << bytes_per_layout / (1024.0 * 1024.0) << " MiB per layout\n"
                  << "Workload: 90-degree rotated 3x3 source reads, "
                  << options.passes << " pass(es) per sample\n";

        const std::vector<std::uint32_t> linear = make_linear_image(pixel_count);
        const std::vector<std::uint32_t> tiled = make_tiled_image(linear, options);
        const AddressParts linear_address = make_linear_parts(options);
        const AddressParts tiled_address = make_tiled_parts(options);

        const std::uint64_t warm_linear =
            sum_rotated_stencil(linear, linear_address, options, 1);
        const std::uint64_t warm_tiled =
            sum_rotated_stencil(tiled, tiled_address, options, 1);
        if (warm_linear != warm_tiled) {
            fail("layout conversion failed: warm-up checksums differ");
        }

        std::vector<double> linear_samples;
        std::vector<double> tiled_samples;
        linear_samples.reserve(options.rounds);
        tiled_samples.reserve(options.rounds);
        std::uint64_t linear_checksum = 0;
        std::uint64_t tiled_checksum = 0;

        for (std::size_t round = 0; round < options.rounds; ++round) {
            const auto run_linear = [&] {
                const TimedResult result = measure([&] {
                    return sum_rotated_stencil(
                        linear, linear_address, options, options.passes);
                });
                linear_samples.push_back(result.milliseconds);
                linear_checksum = result.checksum;
            };
            const auto run_tiled = [&] {
                const TimedResult result = measure([&] {
                    return sum_rotated_stencil(
                        tiled, tiled_address, options, options.passes);
                });
                tiled_samples.push_back(result.milliseconds);
                tiled_checksum = result.checksum;
            };

            if (round % 2 == 0) {
                run_linear();
                run_tiled();
            } else {
                run_tiled();
                run_linear();
            }
        }

        if (linear_checksum != tiled_checksum) {
            fail("measured checksums differ");
        }

        const double logical_bytes_read = static_cast<double>(output_pixels)
            * 9.0 * sizeof(std::uint32_t) * static_cast<double>(options.passes);
        print_result("linear", linear_samples, logical_bytes_read);
        print_result("tiled", tiled_samples, logical_bytes_read);
        const double speedup = median(linear_samples) / median(tiled_samples);
        std::cout << "speedup: " << std::fixed << std::setprecision(2)
                  << speedup << "x\n"
                  << "checksum: " << linear_checksum << "\n"
                  << "Note: this models image-style 2D reads on a CPU; it is not a "
                  << "direct Intel GPU hardware benchmark.\n";
        return 0;
    } catch (const std::exception& exception) {
        std::cerr << "error: " << exception.what() << '\n';
        return 1;
    }
}