#define NOMINMAX
#include <windows.h>
#include <bcrypt.h>
#include <urlmon.h>

#include <algorithm>
#include <array>
#include <chrono>
#include <cctype>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <locale>
#include <numeric>
#include <random>
#include <sstream>
#include <span>
#include <stdexcept>
#include <string>
#include <string_view>
#include <system_error>
#include <vector>

#pragma comment(lib, "bcrypt.lib")
#pragma comment(lib, "urlmon.lib")

namespace {

namespace fs = std::filesystem;

constexpr std::size_t image_side = 28;
constexpr std::size_t pixel_count = image_side * image_side;
constexpr std::size_t official_train_count = 60'000;
constexpr std::size_t official_test_count = 10'000;
constexpr std::size_t dataset_row_count =
    official_train_count + official_test_count;
constexpr std::uintmax_t expected_file_size = 155'225'241;
constexpr std::string_view expected_md5 =
    "cdfc9c58cb9fe86ffaa76af247ae2ef2";
constexpr wchar_t dataset_url[] =
    L"https://www.openml.org/data/download/18238735/Fashion-MNIST.arff";

static_assert(sizeof(float) == 4, "o relatório físico pressupõe FP32 de 4 bytes");

struct Dataset {
    std::size_t dimensions{};
    std::vector<std::uint8_t> pixels;
    std::vector<std::uint8_t> labels;

    [[nodiscard]] std::size_t size() const noexcept {
        return labels.size();
    }

    [[nodiscard]] std::span<const std::uint8_t> image(
        const std::size_t index) const {
        return {pixels.data() + index * dimensions, dimensions};
    }
};

struct Model {
    std::size_t dimensions{};
    std::size_t latent_dimensions{};
    // Each contiguous row contains one encoder direction. The decoder reuses
    // the transpose, so there is only one trainable matrix.
    std::vector<float> weights;
};

struct RunConfig {
    std::string_view name;
    std::size_t train_count;
    std::size_t validation_count;
    std::size_t test_count;
    std::size_t latent_dimensions;
    std::size_t epochs;
    std::size_t batch_size;
    float learning_rate;
    float momentum;
    std::uint32_t seed;
};

struct Options {
    fs::path dataset_path = fs::path{L"dados"} / L"Fashion-MNIST.arff";
    fs::path output_directory = L"resultados_autoencoder";
    RunConfig run{
        "didatico", 10'000, 2'000, 2'000, 16, 12, 128, 0.005F, 0.9F,
        20260728U};
    bool allow_download = true;
    bool self_test_only = false;
};

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

[[nodiscard]] std::string lowercase(std::string text) {
    std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) {
        return static_cast<char>(std::tolower(c));
    });
    return text;
}

[[nodiscard]] std::string trim_copy(std::string_view text) {
    const auto first = text.find_first_not_of(" \t\r\n");
    if (first == std::string_view::npos) {
        return {};
    }
    const auto last = text.find_last_not_of(" \t\r\n");
    return std::string{text.substr(first, last - first + 1)};
}

void check_ntstatus(const NTSTATUS status, const char* operation) {
    if (status < 0) {
        std::ostringstream message;
        message << operation << " falhou com NTSTATUS 0x" << std::hex
                << static_cast<unsigned long>(status);
        fail(message.str());
    }
}

[[nodiscard]] std::string md5_file(const fs::path& path) {
    BCRYPT_ALG_HANDLE algorithm = nullptr;
    BCRYPT_HASH_HANDLE hash = nullptr;
    std::vector<std::uint8_t> hash_object;

    try {
        check_ntstatus(
            BCryptOpenAlgorithmProvider(
                &algorithm, BCRYPT_MD5_ALGORITHM, nullptr, 0),
            "BCryptOpenAlgorithmProvider");

        DWORD object_size = 0;
        DWORD result_size = 0;
        check_ntstatus(
            BCryptGetProperty(
                algorithm,
                BCRYPT_OBJECT_LENGTH,
                reinterpret_cast<PUCHAR>(&object_size),
                sizeof(object_size),
                &result_size,
                0),
            "BCryptGetProperty(BCRYPT_OBJECT_LENGTH)");

        DWORD digest_size = 0;
        check_ntstatus(
            BCryptGetProperty(
                algorithm,
                BCRYPT_HASH_LENGTH,
                reinterpret_cast<PUCHAR>(&digest_size),
                sizeof(digest_size),
                &result_size,
                0),
            "BCryptGetProperty(BCRYPT_HASH_LENGTH)");

        hash_object.resize(object_size);
        std::vector<std::uint8_t> digest(digest_size);
        check_ntstatus(
            BCryptCreateHash(
                algorithm,
                &hash,
                hash_object.data(),
                static_cast<ULONG>(hash_object.size()),
                nullptr,
                0,
                0),
            "BCryptCreateHash");

        std::ifstream input(path, std::ios::binary);
        if (!input) {
            fail("não foi possível abrir o arquivo para calcular o MD5");
        }

        // Keep the 1 MiB I/O block off the default Windows thread stack.
        std::vector<std::uint8_t> buffer(1U << 20U);
        while (input) {
            input.read(
                reinterpret_cast<char*>(buffer.data()),
                static_cast<std::streamsize>(buffer.size()));
            const auto bytes_read = input.gcount();
            if (bytes_read > 0) {
                check_ntstatus(
                    BCryptHashData(
                        hash,
                        buffer.data(),
                        static_cast<ULONG>(bytes_read),
                        0),
                    "BCryptHashData");
            }
        }
        if (!input.eof()) {
            fail("erro de leitura durante o cálculo do MD5");
        }

        check_ntstatus(
            BCryptFinishHash(
                hash,
                digest.data(),
                static_cast<ULONG>(digest.size()),
                0),
            "BCryptFinishHash");

        BCryptDestroyHash(hash);
        hash = nullptr;
        BCryptCloseAlgorithmProvider(algorithm, 0);
        algorithm = nullptr;

        std::ostringstream hex;
        hex << std::hex << std::setfill('0');
        for (const auto byte : digest) {
            hex << std::setw(2) << static_cast<unsigned>(byte);
        }
        return hex.str();
    } catch (...) {
        if (hash != nullptr) {
            BCryptDestroyHash(hash);
        }
        if (algorithm != nullptr) {
            BCryptCloseAlgorithmProvider(algorithm, 0);
        }
        throw;
    }
}

[[nodiscard]] bool has_expected_fingerprint(const fs::path& path) {
    std::error_code error;
    if (!fs::is_regular_file(path, error) || error) {
        return false;
    }
    if (fs::file_size(path, error) != expected_file_size || error) {
        return false;
    }
    return md5_file(path) == expected_md5;
}

void ensure_dataset(const fs::path& path, const bool allow_download) {
    if (has_expected_fingerprint(path)) {
        std::cout << "Dados íntegros encontrados em " << path.string() << "\n";
        return;
    }
    if (!allow_download) {
        fail("o conjunto não existe ou falhou na verificação de integridade");
    }

    const auto parent = path.parent_path();
    if (!parent.empty()) {
        fs::create_directories(parent);
    }
    auto temporary = path;
    temporary += L".part";
    std::error_code error;
    fs::remove(temporary, error);

    std::wcout << L"Baixando Fashion-MNIST do OpenML para "
               << temporary.wstring() << L"\n";
    const HRESULT result = URLDownloadToFileW(
        nullptr, dataset_url, temporary.c_str(), 0, nullptr);
    if (FAILED(result)) {
        std::ostringstream message;
        message << "URLDownloadToFileW falhou com HRESULT 0x" << std::hex
                << static_cast<unsigned long>(result);
        fail(message.str());
    }
    if (!has_expected_fingerprint(temporary)) {
        fail("o download terminou, mas tamanho ou MD5 não corresponde ao OpenML");
    }

    // The previous file is replaced only after the new download has passed both
    // checks, so an interrupted transfer cannot destroy a valid local copy.
    fs::remove(path, error);
    error.clear();
    fs::rename(temporary, path, error);
    if (error) {
        fail("não foi possível instalar o arquivo verificado: " + error.message());
    }
}

[[nodiscard]] bool starts_with_case_insensitive(
    const std::string& text,
    const std::string_view prefix) {
    if (text.size() < prefix.size()) {
        return false;
    }
    for (std::size_t i = 0; i < prefix.size(); ++i) {
        const auto left = static_cast<unsigned char>(text[i]);
        const auto right = static_cast<unsigned char>(prefix[i]);
        if (std::tolower(left) != std::tolower(right)) {
            return false;
        }
    }
    return true;
}

void parse_data_row(
    const std::string_view row,
    const std::size_t dimensions,
    std::span<std::uint8_t> destination,
    std::uint8_t& label) {
    std::size_t position = 0;
    for (std::size_t field = 0; field <= dimensions; ++field) {
        while (position < row.size() &&
               (row[position] == ' ' || row[position] == '\t')) {
            ++position;
        }
        if (position == row.size() || row[position] < '0' || row[position] > '9') {
            fail("campo ARFF ausente ou não inteiro");
        }

        unsigned value = 0;
        while (position < row.size() && row[position] >= '0' &&
               row[position] <= '9') {
            value = value * 10U + static_cast<unsigned>(row[position] - '0');
            ++position;
        }
        while (position < row.size() &&
               (row[position] == ' ' || row[position] == '\t')) {
            ++position;
        }

        const unsigned maximum = field == dimensions ? 9U : 255U;
        if (value > maximum) {
            fail("valor fora do intervalo permitido no ARFF");
        }
        if (field < dimensions) {
            destination[field] = static_cast<std::uint8_t>(value);
        } else {
            label = static_cast<std::uint8_t>(value);
        }

        if (field < dimensions) {
            if (position == row.size() || row[position] != ',') {
                fail("linha ARFF com número incorreto de campos");
            }
            ++position;
        } else if (position != row.size() && row[position] != '\r') {
            fail("conteúdo excedente ao fim da linha ARFF");
        }
    }
}

[[nodiscard]] Dataset read_arff_stream(
    std::istream& input,
    const std::size_t expected_rows,
    const std::size_t dimensions) {
    Dataset dataset;
    dataset.dimensions = dimensions;
    dataset.pixels.resize(expected_rows * dimensions);
    dataset.labels.resize(expected_rows);

    bool found_relation = false;
    bool found_data = false;
    std::size_t attribute_count = 0;
    std::size_t row_count = 0;
    std::string line;
    line.reserve(4'096);

    while (std::getline(input, line)) {
        const auto trimmed = trim_copy(line);
        if (trimmed.empty() || trimmed.front() == '%') {
            continue;
        }
        if (!found_data) {
            if (starts_with_case_insensitive(trimmed, "@relation")) {
                found_relation = true;
            } else if (starts_with_case_insensitive(trimmed, "@attribute")) {
                ++attribute_count;
            } else if (lowercase(trimmed) == "@data") {
                found_data = true;
            } else {
                fail("diretiva desconhecida no cabeçalho ARFF");
            }
            continue;
        }

        if (row_count == expected_rows) {
            fail("o ARFF contém mais exemplos do que o contrato declara");
        }
        auto destination = std::span{
            dataset.pixels.data() + row_count * dimensions, dimensions};
        parse_data_row(trimmed, dimensions, destination, dataset.labels[row_count]);
        ++row_count;
    }

    if (!input.eof()) {
        fail("erro de leitura no arquivo ARFF");
    }
    if (!found_relation || !found_data) {
        fail("cabeçalho ARFF incompleto");
    }
    if (attribute_count != dimensions + 1) {
        fail("dimensão declarada no ARFF não corresponde ao experimento");
    }
    if (row_count != expected_rows) {
        fail("arquivo ARFF truncado ou com número inesperado de exemplos");
    }
    return dataset;
}

[[nodiscard]] Dataset read_fashion_mnist(const fs::path& path) {
    std::ifstream input(path, std::ios::binary);
    if (!input) {
        fail("não foi possível abrir o Fashion-MNIST");
    }
    return read_arff_stream(input, dataset_row_count, pixel_count);
}

[[nodiscard]] std::vector<std::size_t> shuffled_indices(
    const std::size_t first,
    const std::size_t available,
    const std::size_t selected,
    const std::uint32_t seed) {
    if (selected > available) {
        fail("a configuração solicita mais exemplos do que a partição contém");
    }
    std::vector<std::size_t> indices(available);
    std::iota(indices.begin(), indices.end(), first);
    std::mt19937 generator(seed);
    std::shuffle(indices.begin(), indices.end(), generator);
    indices.resize(selected);
    return indices;
}

[[nodiscard]] std::vector<float> training_mean(
    const Dataset& dataset,
    const std::span<const std::size_t> train_indices) {
    std::vector<double> sums(dataset.dimensions, 0.0);
    for (const auto index : train_indices) {
        const auto image = dataset.image(index);
        for (std::size_t j = 0; j < dataset.dimensions; ++j) {
            sums[j] += static_cast<double>(image[j]);
        }
    }

    std::vector<float> mean(dataset.dimensions);
    const double scale = 1.0 /
        (255.0 * static_cast<double>(train_indices.size()));
    for (std::size_t j = 0; j < dataset.dimensions; ++j) {
        mean[j] = static_cast<float>(sums[j] * scale);
    }
    return mean;
}

void orthonormalize_rows(Model& model, std::mt19937& generator) {
    std::normal_distribution<float> normal(0.0F, 1.0F);
    const auto D = model.dimensions;
    for (std::size_t k = 0; k < model.latent_dimensions; ++k) {
        auto row = std::span{model.weights.data() + k * D, D};
        for (auto& value : row) {
            value = normal(generator);
        }
        for (std::size_t previous = 0; previous < k; ++previous) {
            const auto basis = std::span{
                model.weights.data() + previous * D, D};
            double projection = 0.0;
            for (std::size_t j = 0; j < D; ++j) {
                projection += static_cast<double>(row[j]) * basis[j];
            }
            for (std::size_t j = 0; j < D; ++j) {
                row[j] -= static_cast<float>(projection) * basis[j];
            }
        }
        double squared_norm = 0.0;
        for (const auto value : row) {
            squared_norm += static_cast<double>(value) * value;
        }
        if (squared_norm <= std::numeric_limits<double>::min()) {
            fail("a inicialização produziu uma direção degenerada");
        }
        const auto inverse_norm = static_cast<float>(1.0 / std::sqrt(squared_norm));
        for (auto& value : row) {
            value *= inverse_norm;
        }
    }
}

[[nodiscard]] Model make_model(
    const std::size_t dimensions,
    const std::size_t latent_dimensions,
    const std::uint32_t seed) {
    if (latent_dimensions == 0 || latent_dimensions > dimensions) {
        fail("dimensão latente inválida");
    }
    Model model{dimensions, latent_dimensions,
                std::vector<float>(dimensions * latent_dimensions)};
    std::mt19937 generator(seed);
    orthonormalize_rows(model, generator);
    return model;
}

void center_image(
    const std::span<const std::uint8_t> image,
    const std::span<const float> mean,
    const std::span<float> centered) {
    for (std::size_t j = 0; j < image.size(); ++j) {
        centered[j] = static_cast<float>(image[j]) / 255.0F - mean[j];
    }
}

void encode_decode_error(
    const Model& model,
    const std::span<const float> input,
    const std::span<float> latent,
    const std::span<float> error) {
    const auto D = model.dimensions;
    const auto K = model.latent_dimensions;
    for (std::size_t k = 0; k < K; ++k) {
        const auto row = std::span{
            model.weights.data() + k * D, D};
        float sum = 0.0F;
        for (std::size_t j = 0; j < D; ++j) {
            sum += row[j] * input[j];
        }
        latent[k] = sum;
    }
    std::copy(input.begin(), input.end(), error.begin());
    for (auto& value : error) {
        value = -value;
    }
    for (std::size_t k = 0; k < K; ++k) {
        const auto row = std::span{
            model.weights.data() + k * D, D};
        const float activation = latent[k];
        for (std::size_t j = 0; j < D; ++j) {
            error[j] += activation * row[j];
        }
    }
}

void accumulate_gradient(
    const Model& model,
    const std::span<const float> input,
    const std::span<const float> latent,
    const std::span<const float> error,
    const std::span<float> projected_error,
    const std::span<float> gradient) {
    const auto D = model.dimensions;
    const auto K = model.latent_dimensions;
    for (std::size_t k = 0; k < K; ++k) {
        const auto row = std::span{
            model.weights.data() + k * D, D};
        float sum = 0.0F;
        for (std::size_t j = 0; j < D; ++j) {
            sum += row[j] * error[j];
        }
        projected_error[k] = sum;
    }

    for (std::size_t k = 0; k < K; ++k) {
        auto gradient_row = std::span{gradient.data() + k * D, D};
        const float z = latent[k];
        const float q = projected_error[k];
        for (std::size_t j = 0; j < D; ++j) {
            gradient_row[j] += z * error[j] + q * input[j];
        }
    }
}

[[nodiscard]] double evaluate_mse(
    const Dataset& dataset,
    const std::span<const std::size_t> indices,
    const std::span<const float> mean,
    const Model& model) {
    std::vector<float> input(model.dimensions);
    std::vector<float> error(model.dimensions);
    std::vector<float> latent(model.latent_dimensions);
    double squared_error = 0.0;
    for (const auto index : indices) {
        center_image(dataset.image(index), mean, input);
        encode_decode_error(model, input, latent, error);
        for (const auto value : error) {
            squared_error += static_cast<double>(value) * value;
        }
    }
    const auto denominator = static_cast<double>(indices.size()) *
        static_cast<double>(model.dimensions);
    return squared_error / denominator;
}

struct WorkBuffers {
    explicit WorkBuffers(const Model& model)
        : input(model.dimensions),
          error(model.dimensions),
          latent(model.latent_dimensions),
          projected_error(model.latent_dimensions),
          gradient(model.weights.size()),
          velocity(model.weights.size()) {}

    std::vector<float> input;
    std::vector<float> error;
    std::vector<float> latent;
    std::vector<float> projected_error;
    std::vector<float> gradient;
    std::vector<float> velocity;
};

struct PhysicalCost {
    std::uint64_t dataset_pixel_bytes{};
    std::uint64_t dataset_label_bytes{};
    std::uint64_t weight_parameters{};
    std::uint64_t weight_bytes{};
    std::uint64_t optimizer_bytes{};
    std::uint64_t per_image_buffer_bytes{};
    std::uint64_t hot_training_bytes{};
    std::uint64_t input_fp32_bytes{};
    std::uint64_t latent_fp32_bytes{};
    std::uint64_t forward_flops_per_image{};
    std::uint64_t forward_flops_per_full_batch{};
    std::uint64_t training_flops_per_image{};
    std::uint64_t training_flops_per_epoch{};
};

[[nodiscard]] PhysicalCost estimate_physical_cost(
    const std::size_t dataset_examples,
    const std::size_t dimensions,
    const std::size_t latent_dimensions,
    const std::size_t train_examples,
    const std::size_t batch_size) {
    if (batch_size == 0) {
        fail("o tamanho do minilote não pode ser zero");
    }

    const auto D = static_cast<std::uint64_t>(dimensions);
    const auto K = static_cast<std::uint64_t>(latent_dimensions);
    const auto N = static_cast<std::uint64_t>(train_examples);
    const auto B = static_cast<std::uint64_t>(batch_size);
    const auto examples = static_cast<std::uint64_t>(dataset_examples);
    const auto float_bytes = static_cast<std::uint64_t>(sizeof(float));
    const auto weights = D * K;
    const auto batches = (N + B - 1U) / B;

    PhysicalCost cost;
    cost.dataset_pixel_bytes = examples * D * sizeof(std::uint8_t);
    cost.dataset_label_bytes = examples * sizeof(std::uint8_t);
    cost.weight_parameters = weights;
    cost.weight_bytes = weights * float_bytes;
    // Momentum needs one gradient accumulator and one velocity per weight.
    cost.optimizer_bytes = 2U * cost.weight_bytes;
    // The implementation converts and processes one image at a time, so the
    // batch size does not multiply these activation buffers.
    cost.per_image_buffer_bytes = (2U * D + 2U * K) * float_bytes;
    cost.hot_training_bytes = cost.weight_bytes + D * float_bytes +
        cost.optimizer_bytes + cost.per_image_buffer_bytes;
    cost.input_fp32_bytes = D * float_bytes;
    cost.latent_fp32_bytes = K * float_bytes;
    // A multiply followed by an accumulation counts as two floating-point
    // operations, independently of whether the compiler emits an FMA.
    cost.forward_flops_per_image = 4U * D * K;
    cost.forward_flops_per_full_batch = B * cost.forward_flops_per_image;
    cost.training_flops_per_image = 10U * D * K;
    cost.training_flops_per_epoch =
        N * cost.training_flops_per_image + batches * 5U * D * K;
    return cost;
}

void print_physical_cost(
    const Dataset& dataset,
    const Model& model,
    const RunConfig& config) {
    const auto cost = estimate_physical_cost(
        dataset.size(),
        model.dimensions,
        model.latent_dimensions,
        config.train_count,
        config.batch_size);
    const auto flags = std::cout.flags();
    const auto precision = std::cout.precision();
    const auto kib = [](const std::uint64_t bytes) {
        return static_cast<double>(bytes) / 1024.0;
    };

    std::cout << "Custo físico estimado para esta configuração:\n"
              << "  conjunto em RAM: " << cost.dataset_pixel_bytes
              << " bytes de pixels + " << cost.dataset_label_bytes
              << " bytes de rótulos\n"
              << "  pesos ligados: " << cost.weight_parameters
              << " FP32 = " << cost.weight_bytes << " bytes ("
              << std::fixed << std::setprecision(2) << kib(cost.weight_bytes)
              << " KiB)\n"
              << "  gradiente + momento: " << cost.optimizer_bytes
              << " bytes (" << kib(cost.optimizer_bytes) << " KiB)\n"
              << "  buffers por imagem: " << cost.per_image_buffer_bytes
              << " bytes (" << kib(cost.per_image_buffer_bytes) << " KiB)\n"
              << "  pesos + média + buffers de treino: "
              << cost.hot_training_bytes << " bytes ("
              << kib(cost.hot_training_bytes) << " KiB)\n"
              << "  entrada FP32 -> latente FP32: " << cost.input_fp32_bytes
              << " -> " << cost.latent_fp32_bytes << " bytes ("
              << static_cast<double>(cost.input_fp32_bytes) /
                     static_cast<double>(cost.latent_fp32_bytes)
              << "x)\n"
              << "  imagem uint8 -> latente FP32: " << model.dimensions
              << " -> " << cost.latent_fp32_bytes << " bytes ("
              << static_cast<double>(model.dimensions) /
                     static_cast<double>(cost.latent_fp32_bytes)
              << "x)\n"
              << "  passagem direta: " << cost.forward_flops_per_image
              << " FLOPs/imagem; " << cost.forward_flops_per_full_batch
              << " FLOPs/lote cheio\n"
              << "  núcleo de treino: " << cost.training_flops_per_image
              << " FLOPs/imagem; " << cost.training_flops_per_epoch
              << " FLOPs/época\n";
    std::cout.flags(flags);
    std::cout.precision(precision);
}

void train_one_epoch(
    const Dataset& dataset,
    std::vector<std::size_t>& train_indices,
    const std::span<const float> mean,
    Model& model,
    WorkBuffers& work,
    const RunConfig& config,
    std::mt19937& generator) {
    std::shuffle(train_indices.begin(), train_indices.end(), generator);
    for (std::size_t begin = 0; begin < train_indices.size();
         begin += config.batch_size) {
        const auto end = std::min(begin + config.batch_size, train_indices.size());
        std::fill(work.gradient.begin(), work.gradient.end(), 0.0F);

        for (std::size_t position = begin; position < end; ++position) {
            center_image(dataset.image(train_indices[position]), mean, work.input);
            encode_decode_error(model, work.input, work.latent, work.error);
            accumulate_gradient(
                model,
                work.input,
                work.latent,
                work.error,
                work.projected_error,
                work.gradient);
        }

        const float inverse_batch =
            1.0F / static_cast<float>(end - begin);
        for (std::size_t i = 0; i < model.weights.size(); ++i) {
            work.velocity[i] = config.momentum * work.velocity[i] +
                inverse_batch * work.gradient[i];
            model.weights[i] -= config.learning_rate * work.velocity[i];
        }
    }
}

[[nodiscard]] std::vector<float> reconstruct(
    const Dataset& dataset,
    const std::size_t index,
    const std::span<const float> mean,
    const Model& model) {
    std::vector<float> input(model.dimensions);
    std::vector<float> error(model.dimensions);
    std::vector<float> latent(model.latent_dimensions);
    center_image(dataset.image(index), mean, input);
    encode_decode_error(model, input, latent, error);

    std::vector<float> reconstruction(model.dimensions);
    for (std::size_t j = 0; j < model.dimensions; ++j) {
        reconstruction[j] = std::clamp(error[j] + input[j] + mean[j], 0.0F, 1.0F);
    }
    return reconstruction;
}

void write_contact_sheet(
    const fs::path& path,
    const Dataset& dataset,
    const std::span<const std::size_t> test_indices,
    const std::span<const float> mean,
    const Model& model) {
    constexpr std::size_t columns = 8;
    constexpr std::size_t gap = 2;
    const auto examples = std::min(columns, test_indices.size());
    const auto width = examples * image_side + (examples - 1) * gap;
    const auto height = 2 * image_side + gap;
    std::vector<std::uint8_t> canvas(width * height, 255);

    for (std::size_t column = 0; column < examples; ++column) {
        const auto index = test_indices[column];
        const auto original = dataset.image(index);
        const auto reconstruction = reconstruct(dataset, index, mean, model);
        const auto x_offset = column * (image_side + gap);
        for (std::size_t y = 0; y < image_side; ++y) {
            for (std::size_t x = 0; x < image_side; ++x) {
                const auto source = y * image_side + x;
                canvas[y * width + x_offset + x] = original[source];
                canvas[(y + image_side + gap) * width + x_offset + x] =
                    static_cast<std::uint8_t>(std::lround(
                        255.0F * reconstruction[source]));
            }
        }
    }

    std::ofstream output(path, std::ios::binary);
    if (!output) {
        fail("não foi possível criar a folha de reconstruções");
    }
    output << "P5\n" << width << ' ' << height << "\n255\n";
    output.write(
        reinterpret_cast<const char*>(canvas.data()),
        static_cast<std::streamsize>(canvas.size()));
    if (!output) {
        fail("erro ao escrever a folha de reconstruções");
    }
}

[[nodiscard]] double half_sse(
    const Model& model,
    const std::span<const float> examples,
    const std::size_t example_count) {
    std::vector<float> latent(model.latent_dimensions);
    std::vector<float> error(model.dimensions);
    double total = 0.0;
    for (std::size_t i = 0; i < example_count; ++i) {
        const auto input = examples.subspan(i * model.dimensions, model.dimensions);
        encode_decode_error(model, input, latent, error);
        for (const auto value : error) {
            total += 0.5 * static_cast<double>(value) * value;
        }
    }
    return total / static_cast<double>(example_count);
}

void test_arff_reader() {
    constexpr std::string_view header =
        "@RELATION tiny\n"
        "@ATTRIBUTE pixel1 real\n"
        "@ATTRIBUTE pixel2 real\n"
        "@ATTRIBUTE class {0,1,2,3,4,5,6,7,8,9}\n"
        "@DATA\n";
    {
        std::istringstream input(std::string{header} + "0,255,1\n128,64,9\n");
        const auto data = read_arff_stream(input, 2, 2);
        if (data.size() != 2 || data.pixels[1] != 255 || data.labels[1] != 9) {
            fail("o autoteste do leitor ARFF não reproduziu os dados válidos");
        }
    }

    const auto expect_failure = [](const std::string& text,
                                   const std::size_t rows,
                                   const std::size_t dimensions) {
        try {
            std::istringstream input(text);
            static_cast<void>(read_arff_stream(input, rows, dimensions));
        } catch (const std::exception&) {
            return;
        }
        fail("o leitor ARFF aceitou deliberadamente uma entrada inválida");
    };
    expect_failure(std::string{header} + "0,255,1\n", 2, 2);
    expect_failure(std::string{header} + "0,256,1\n128,64,9\n", 2, 2);
    expect_failure(
        "@RELATION tiny\n@ATTRIBUTE pixel1 real\n@ATTRIBUTE class {0,1}\n"
        "@DATA\n0,1\n",
        1,
        2);
}

void test_tied_gradient() {
    Model model{3, 2, {0.30F, -0.20F, 0.10F, -0.40F, 0.25F, 0.15F}};
    const std::array examples{
        0.20F, -0.10F, 0.30F,
        -0.40F, 0.25F, 0.10F};
    std::vector<float> analytical(model.weights.size(), 0.0F);
    std::vector<float> input(model.dimensions);
    std::vector<float> latent(model.latent_dimensions);
    std::vector<float> error(model.dimensions);
    std::vector<float> projected(model.latent_dimensions);

    for (std::size_t i = 0; i < 2; ++i) {
        std::copy_n(examples.data() + i * model.dimensions,
                    model.dimensions,
                    input.data());
        encode_decode_error(model, input, latent, error);
        accumulate_gradient(
            model, input, latent, error, projected, analytical);
    }
    for (auto& value : analytical) {
        value /= 2.0F;
    }

    constexpr float epsilon = 1.0e-3F;
    for (std::size_t i = 0; i < model.weights.size(); ++i) {
        const float original = model.weights[i];
        model.weights[i] = original + epsilon;
        const double plus = half_sse(model, examples, 2);
        model.weights[i] = original - epsilon;
        const double minus = half_sse(model, examples, 2);
        model.weights[i] = original;
        const double numerical = (plus - minus) / (2.0 * epsilon);
        if (std::abs(numerical - analytical[i]) > 2.0e-4) {
            fail("a diferença finita não confirmou o gradiente dos pesos ligados");
        }
    }
}

void test_physical_cost() {
    const auto cost = estimate_physical_cost(
        dataset_row_count, pixel_count, 16, 10'000, 128);
    if (cost.dataset_pixel_bytes != 54'880'000 ||
        cost.dataset_label_bytes != 70'000 ||
        cost.weight_parameters != 12'544 ||
        cost.weight_bytes != 50'176 ||
        cost.optimizer_bytes != 100'352 ||
        cost.per_image_buffer_bytes != 6'400 ||
        cost.hot_training_bytes != 160'064 ||
        cost.input_fp32_bytes != 3'136 ||
        cost.latent_fp32_bytes != 64 ||
        cost.forward_flops_per_image != 50'176 ||
        cost.forward_flops_per_full_batch != 6'422'528 ||
        cost.training_flops_per_image != 125'440 ||
        cost.training_flops_per_epoch != 1'259'354'880) {
        fail("o autoteste do relatório físico encontrou uma contagem divergente");
    }
}

void run_self_tests() {
    test_arff_reader();
    test_tied_gradient();
    test_physical_cost();
    std::cout
        << "Autotestes: leitor ARFF, gradiente e custo físico confirmados.\n";
}

[[nodiscard]] Options parse_options(const int argc, wchar_t* argv[]) {
    Options options;
    for (int i = 1; i < argc; ++i) {
        const std::wstring_view argument = argv[i];
        if (argument == L"--smoke") {
            options.run = {
                "diagnostico", 512, 128, 128, 4, 2, 64, 0.01F, 0.9F,
                20260728U};
        } else if (argument == L"--full") {
            options.run = {
                "completo", 55'000, 5'000, 10'000, 32, 20, 128, 0.003F,
                0.9F, 20260728U};
        } else if (argument == L"--data" && i + 1 < argc) {
            options.dataset_path = argv[++i];
        } else if (argument == L"--output" && i + 1 < argc) {
            options.output_directory = argv[++i];
        } else if (argument == L"--no-download") {
            options.allow_download = false;
        } else if (argument == L"--self-test") {
            options.self_test_only = true;
        } else if (argument == L"--help") {
            std::cout
                << "Uso: autoencoder_fashion_mnist [--smoke|--full] "
                   "[--data ARQUIVO] [--output DIRETORIO] [--no-download] "
                   "[--self-test]\n";
            std::exit(EXIT_SUCCESS);
        } else {
            fail("argumento desconhecido; use --help");
        }
    }
    return options;
}

void run_experiment(const Options& options) {
    ensure_dataset(options.dataset_path, options.allow_download);
    std::cout << "Lendo e validando 70.000 linhas ARFF...\n";
    const auto dataset = read_fashion_mnist(options.dataset_path);
    const auto& config = options.run;

    if (config.train_count + config.validation_count > official_train_count) {
        fail("treino e validação excedem a partição oficial de treino");
    }
    auto development_indices = shuffled_indices(
        0,
        official_train_count,
        official_train_count,
        config.seed);
    auto train_indices = std::vector<std::size_t>{
        development_indices.begin(),
        development_indices.begin() +
            static_cast<std::ptrdiff_t>(config.train_count)};
    const auto validation_indices = std::vector<std::size_t>{
        development_indices.begin() +
            static_cast<std::ptrdiff_t>(config.train_count),
        development_indices.begin() +
            static_cast<std::ptrdiff_t>(
                config.train_count + config.validation_count)};
    const auto test_indices = shuffled_indices(
        official_train_count,
        official_test_count,
        config.test_count,
        config.seed + 1U);
    const auto mean = training_mean(dataset, train_indices);
    auto model = make_model(pixel_count, config.latent_dimensions, config.seed + 2U);
    WorkBuffers work(model);
    std::mt19937 training_generator(config.seed + 3U);

    fs::create_directories(options.output_directory);
    const auto metrics_path = options.output_directory / L"metricas.csv";
    std::ofstream metrics(metrics_path);
    if (!metrics) {
        fail("não foi possível criar o CSV de métricas");
    }
    metrics.imbue(std::locale::classic());
    metrics << "epoca,mse_treino,mse_validacao,mse_teste,segundos_acumulados\n";

    std::cout << "Execução " << config.name << ": "
              << config.train_count << " treino, " << config.validation_count
              << " validação, " << config.test_count << " teste, D="
              << pixel_count << ", d="
              << config.latent_dimensions << ", " << model.weights.size()
              << " parâmetros FP32 ligados.\n";
    print_physical_cost(dataset, model, config);

    const auto start = std::chrono::steady_clock::now();
    auto report = [&](const std::size_t epoch, const bool measure_test) {
        const double train_mse = evaluate_mse(dataset, train_indices, mean, model);
        const double validation_mse =
            evaluate_mse(dataset, validation_indices, mean, model);
        const double test_mse = measure_test
            ? evaluate_mse(dataset, test_indices, mean, model)
            : std::numeric_limits<double>::quiet_NaN();
        const double seconds = std::chrono::duration<double>(
            std::chrono::steady_clock::now() - start).count();
        if (!std::isfinite(train_mse) || !std::isfinite(validation_mse) ||
            (measure_test && !std::isfinite(test_mse))) {
            fail("o treinamento produziu uma perda não finita");
        }
        metrics << epoch << ',' << std::setprecision(9) << train_mse << ','
                << validation_mse << ',';
        if (measure_test) {
            metrics << test_mse;
        }
        metrics << ',' << std::setprecision(6) << seconds << '\n';
        metrics.flush();
        std::cout << "época=" << std::setw(2) << epoch
                  << " mse_treino=" << std::fixed << std::setprecision(6)
                  << train_mse << " mse_validacao=" << validation_mse;
        if (measure_test) {
            std::cout << " mse_teste=" << test_mse;
        }
        std::cout << " tempo=" << std::setprecision(2) << seconds << " s\n";
    };

    report(0, true);
    for (std::size_t epoch = 1; epoch <= config.epochs; ++epoch) {
        train_one_epoch(
            dataset,
            train_indices,
            mean,
            model,
            work,
            config,
            training_generator);
        report(epoch, epoch == config.epochs);
    }

    const auto sheet_path = options.output_directory / L"reconstrucoes.pgm";
    write_contact_sheet(sheet_path, dataset, test_indices, mean, model);
    std::cout << "Rótulos das oito imagens (somente inspeção pós-treino): ";
    for (std::size_t i = 0; i < std::min<std::size_t>(8, test_indices.size()); ++i) {
        std::cout << static_cast<unsigned>(dataset.labels[test_indices[i]])
                  << (i + 1 == std::min<std::size_t>(8, test_indices.size()) ? '\n' : ' ');
    }
    std::cout << "Arquivos: " << metrics_path.string() << " e "
              << sheet_path.string() << "\n";
}

}  // namespace

int wmain(const int argc, wchar_t* argv[]) {
    try {
        const auto options = parse_options(argc, argv);
        run_self_tests();
        if (!options.self_test_only) {
            run_experiment(options);
        }
        return EXIT_SUCCESS;
    } catch (const std::exception& error) {
        std::cerr << "erro: " << error.what() << '\n';
        return EXIT_FAILURE;
    }
}