cmake_minimum_required(VERSION 3.28 FATAL_ERROR)

project(HelloLibtorch
    VERSION 1.0.0
    LANGUAGES CXX CUDA)

set(CMAKE_CXX_STANDARD 20 )
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
set(CMAKE_COLOR_DIAGNOSTICS ON)

# Додавання шляху до бібліотеки cusparselt
find_library(CUSPARSELT_LIB cusparseLt
  PATHS 
  /шлях до cusparseLt/lib/
)


# Додавання Python, matplotlib, numpy, matplotlibcpp
find_package(Python3 COMPONENTS Interpreter Development NumPy REQUIRED)
set(MATPLOTLIB_CPP /шлях до matplotlib_cpp)
include_directories(${MATPLOTLIB_CPP})

if(NOT CUSPARSELT_LIB)
    message(FATAL_ERROR "libcusparseLt.so not found. Please install CUDA Toolkit with cuSPARSELt.")
endif()

# Додавання бібліотеки Libtorch, що знаходиться в окремій теці
set(CMAKE_PREFIX_PATH /шлях до libtorch)
set(Torch_DIR "${CMAKE_SOURCE_DIR}/libtorch/share/cmake/Torch")
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

add_executable(${PROJECT_NAME} main.cpp)

target_include_directories(${PROJECT_NAME} PRIVATE
    ${Python3_INCLUDE_DIRS}
    ${Python3_NumPy_INCLUDE_DIRS}
)
  
target_link_libraries(${PROJECT_NAME}
  "${TORCH_LIBRARIES}"
  ${CUSPARSELT_LIB}
  ${Python3_LIBRARIES}
)

set_property(TARGET ${PROJECT_NAME} PROPERTY  CXX_STANDARD 20)


# Потрібне для автозапуску програми після компіляції
add_custom_command(
    TARGET ${PROJECT_NAME}
    POST_BUILD
    COMMAND ${PROJECT_NAME}
)

#include <cmath>
#include <cstdio>
#include <iostream>
#include <iomanip>
#include <vector>
#include <map>
#include <random>
#include <string>

#include <torch/torch.h>
#include <matplotlibcpp.h>

namespace plt = matplotlibcpp;

// The size of the noise vector fed to the generator.
const int64_t kNoiseSize = 100;

// The batch size for training.
const int64_t kBatchSize = 64;

// Where to find the MNIST dataset.
const char* kDataFolder = "../../mnist/MNIST/raw/";

// After how many batches to create a new checkpoint periodically.
const int64_t kCheckpointEvery = 200;

// How many images to sample at every checkpoint.
const int64_t kNumberOfSamplesPerCheckpoint = 10;

// Set to `true` to restore models and optimizers from previously saved
// checkpoints.
const bool kRestoreFromCheckpoint = false;

// After how many batches to log a new update with the loss value.
const int64_t kLogInterval = 10;

using namespace torch;

struct DCGANGeneratorImpl : nn::Module {
    DCGANGeneratorImpl(int kNoiseSize)
        : conv1(nn::ConvTranspose2dOptions(kNoiseSize, 256, 4).bias(false))
        , batch_norm1(256)
        , conv2(nn::ConvTranspose2dOptions(256, 128, 3)
                    .stride(2)
                    .padding(1)
                    .bias(false))
        , batch_norm2(128)
        , conv3(nn::ConvTranspose2dOptions(128, 64, 4)
                    .stride(2)
                    .padding(1)
                    .bias(false))
        , batch_norm3(64)
        , conv4(nn::ConvTranspose2dOptions(64, 1, 4).stride(2).padding(1).bias(
              false))
    {
        // register_module() is needed if we want to use the parameters()
        // method later on
        register_module("conv1", conv1);
        register_module("conv2", conv2);
        register_module("conv3", conv3);
        register_module("conv4", conv4);
        register_module("batch_norm1", batch_norm1);
        register_module("batch_norm2", batch_norm2);
        register_module("batch_norm3", batch_norm3);
    }

    torch::Tensor forward(torch::Tensor x)
    {
        x = torch::relu(batch_norm1(conv1(x)));
        x = torch::relu(batch_norm2(conv2(x)));
        x = torch::relu(batch_norm3(conv3(x)));
        x = torch::tanh(conv4(x));
        return x;
    }

    nn::ConvTranspose2d conv1, conv2, conv3, conv4;
    nn::BatchNorm2d batch_norm1, batch_norm2, batch_norm3;
};

TORCH_MODULE(DCGANGenerator);

nn::Sequential create_discriminator()
{
    return nn::Sequential(
        // Layer 1
        nn::Conv2d(
            nn::Conv2dOptions(1, 64, 4).stride(2).padding(1).bias(false)),
        nn::LeakyReLU(nn::LeakyReLUOptions().negative_slope(0.2)),
        // Layer 2
        nn::Conv2d(
            nn::Conv2dOptions(64, 128, 4).stride(2).padding(1).bias(false)),
        nn::BatchNorm2d(128),
        nn::LeakyReLU(nn::LeakyReLUOptions().negative_slope(0.2)),
        // Layer 3
        nn::Conv2d(
            nn::Conv2dOptions(128, 256, 4).stride(2).padding(1).bias(false)),
        nn::BatchNorm2d(256),
        nn::LeakyReLU(nn::LeakyReLUOptions().negative_slope(0.2)),
        // Layer 4
        nn::Conv2d(
            nn::Conv2dOptions(256, 1, 3).stride(1).padding(0).bias(false)),
        nn::Sigmoid());
}

template <typename T>
auto tensor_to_vector_copy(const torch::Tensor& tensor) -> std::vector<T>
{
    // Функція для перетворення тензора на std::vector обраного типу даних.
    // Для matplotlibcpp це повинен бути std::vector<float>.

    // Тензор потрібно перенести до ЦПУ та зробити неперервним
    torch::Tensor cur_tensor = tensor.detach().cpu().contiguous();

    if (cur_tensor.scalar_type() != torch::CppTypeToScalarType<T>()) {
        throw std::runtime_error("Скалярний тип тензора не придатний для "
                                 "копіювання до стандартного типу даних С++.");
    }

    std::vector<T> vec(cur_tensor.numel());

    std::copy(cur_tensor.data_ptr<T>(),
              cur_tensor.data_ptr<T>() + cur_tensor.numel(), vec.begin());
    return vec;
}

auto get_random_number(const int batch_size) -> int
{
    // Функція для отримання випадкового числа
    auto rand_dev { std::random_device() };
    std::mt19937_64 gen(rand_dev());
    std::uniform_int_distribution<> distr(0, batch_size - 1);
    auto random_number { distr(gen) };

    return random_number;
}

void show_images(const torch::data::Example<>& batch, const int image_amount,
                 const long cols = 1)
{
    // Функція для показу випадкових зображень із партії(batch) тензора
    // за допомогою бібліотеки matplotlibcpp.

    // Кількість елементів в партії(batch) тензора
    const auto batch_item_amount { batch.data.size(0) };
    // Кількість кольорових каналів (зазвичай 1 або 3)
    const auto color_channels { batch.data.size(1) };
    // Висота зображення в пікселях
    const auto image_height { batch.data.size(2) };
    // Ширина зображення в пікселях
    const auto image_width { batch.data.size(3) };

    if (image_amount >= batch_item_amount || (image_amount % cols != 0)) {
        throw std::runtime_error(
            "image_amount має бути додатнім числом меншим за batch_size та "
            "ділитися націло на columns.");
    }

    const long row { image_amount
                     / cols }; // Кількість рядків із зображеннями на графіку

    long pos { 1 }; // Положення зображення на графіку
    /* Кількість рядків і стовпців із зображеннями має бути сталою.
       Максимальне положення зображення не має перевищувати добуток рядків та
       стовпців.
     */

    // Параметри для imshow, які відповідають за відображення кольорової гами
    // Для чорнобілого зображення треба вказувати cmap = gray
    // інакше будуть два будь-які кольори (зазвичай жовтий і зелений)
    std::map<std::string, std::string> imshow_keywords;
    if (color_channels == 1) {
        imshow_keywords["cmap"] = "gray";
    }

    plt::suptitle("Випадкові зображення з бази даних");
    for (std::size_t i = 0; i < static_cast<std::size_t>(image_amount); ++i) {
        // Випадковий тензор із зображенням і його мітка-значення
        auto random_number { get_random_number(batch_item_amount) };
        torch::Tensor img_tensor = batch.data[random_number]; 
        const int label = batch.target[random_number].item<int>();

        // Переводимо тензор до ЦПУ і робимо одновимірним
        // Якщо тип не float, виправляємо це
        img_tensor = img_tensor.squeeze(0).detach().cpu();
        if (img_tensor.dtype() != torch::kFloat32) {
            img_tensor = img_tensor.to(torch::kFloat32);
        }

        const auto vector_float { tensor_to_vector_copy<float>(img_tensor) };

        plt::subplot(row, cols, pos);
        ++pos;
        plt::imshow(vector_float.data(), image_height, image_width,
                    color_channels, imshow_keywords);
        plt::title("Мітка: " + std::to_string(label));
    }
    // Налаштування відступу між зображеннями
    // wspace — горизонтальний відступ
    // hspace — вертикальний відступ
    std::map<std::string, double> adjust_keywords;
    adjust_keywords["hspace"] = 0.75;
    adjust_keywords["wspace"] = 0.2;
    plt::subplots_adjust(adjust_keywords);
    plt::show();
}

void show_test_images(std::vector<torch::Tensor>& vector_with_tensors,
                      const long columns = 1)
{
    const int color_channels { static_cast<int>(
        vector_with_tensors[0].size(1)) };
    const int image_height { static_cast<int>(
        vector_with_tensors[0].size(2)) };
    const int image_width { static_cast<int>(vector_with_tensors[0].size(3)) };
    const long images_amount { static_cast<long>(vector_with_tensors.size()) };
    if (images_amount % columns != 0) {
        throw std::runtime_error(
            "Кількість зображень має націло ділитися на кількість стовпців.");
    }
    const long row { images_amount / columns };
    long pos { 1 };

    std::map<std::string, std::string> imshow_keywords;
    if (color_channels == 1) {
        imshow_keywords["cmap"] = "gray";
    }
    for (torch::Tensor& img_tensor : vector_with_tensors) {
        img_tensor = img_tensor.squeeze(0).detach().cpu();
        if (img_tensor.dtype() != torch::kFloat32) {
            img_tensor = img_tensor.to(torch::kFloat32);
        }

        const auto vector_float { tensor_to_vector_copy<float>(img_tensor) };

        plt::subplot(row, columns, pos);
        pos++;
        plt::imshow(vector_float.data(), image_height, image_width,
                    color_channels, imshow_keywords);
    }

    plt::suptitle("Тестове зображення");
    std::map<std::string, double> adjust_keywords;
    adjust_keywords["hspace"] = 0.75;
    adjust_keywords["wspace"] = 0.2;
    plt::subplots_adjust(adjust_keywords);
    plt::show();
}

auto main() -> int
{

    // The number of epochs to train, default value is 30.
    const int64_t kNumberOfEpochs = 1;
    std::cout << "Traning with number of epochs: " << kNumberOfEpochs
              << std::endl;

    torch::manual_seed(1);

    // Create the device we pass around based on whether CUDA is available.
    torch::Device device(torch::kCPU);
    if (torch::cuda::is_available()) {
        std::cout << "CUDA is available! Training on GPU." << std::endl;
        device = torch::Device(torch::kCUDA);
    }

    DCGANGenerator generator(kNoiseSize);
    generator->to(device);

    nn::Sequential discriminator = create_discriminator();
    discriminator->to(device);

    // Assume the MNIST dataset is available under `kDataFolder`;
    auto dataset = torch::data::datasets::MNIST(kDataFolder)
                       .map(torch::data::transforms::Normalize<>(0.5, 0.5))
                       .map(torch::data::transforms::Stack<>());
    const int64_t batches_per_epoch = static_cast<int64_t>(
        std::ceil(dataset.size().value() / static_cast<double>(kBatchSize)));

    auto data_loader = torch::data::make_data_loader(
        std::move(dataset),
        torch::data::DataLoaderOptions().batch_size(kBatchSize).workers(2));

    try {
        std::vector<torch::data::Example<torch::Tensor, torch::Tensor>>
            all_batches;
        all_batches.reserve(batches_per_epoch);

        for (const auto& batch : *data_loader) {
            // Clone to keep independent copy
            all_batches.push_back(
                { batch.data.clone(), batch.target.clone() });
        }

        constexpr int image_amount { 15 };
        constexpr long columns { 5 };
        auto rand_number { get_random_number(all_batches.size()) };
        show_images(all_batches[rand_number], image_amount, columns);

    } catch (const c10::Error& error) {
        std::cout << "LibTorch error: " << error.msg() << "\n";
    } catch (const std::exception& error) {
        std::cout << "Error: " << error.what() << "\n";
    }

    torch::optim::Adam generator_optimizer(
        generator->parameters(),
        torch::optim::AdamOptions(2e-4).betas(std::make_tuple(0.5, 0.5)));
    torch::optim::Adam discriminator_optimizer(
        discriminator->parameters(),
        torch::optim::AdamOptions(2e-4).betas(std::make_tuple(0.5, 0.5)));

    if (kRestoreFromCheckpoint) {
        torch::load(generator, "generator-checkpoint.pt");
        torch::load(generator_optimizer, "generator-optimizer-checkpoint.pt");
        torch::load(discriminator, "discriminator-checkpoint.pt");
        torch::load(discriminator_optimizer,
                    "discriminator-optimizer-checkpoint.pt");
    }

    int64_t checkpoint_counter = 1;
    for (int64_t epoch = 1; epoch <= kNumberOfEpochs; ++epoch) {
        int64_t batch_index = 0;
        for (const torch::data::Example<>& batch : *data_loader) {
            // Train discriminator with real images.
            discriminator->zero_grad();
            torch::Tensor real_images = batch.data.to(device);
            torch::Tensor real_labels
                = torch::empty(batch.data.size(0), device).uniform_(0.8, 1.0);
            torch::Tensor real_output = discriminator->forward(real_images)
                                            .reshape(real_labels.sizes());
            torch::Tensor d_loss_real
                = torch::binary_cross_entropy(real_output, real_labels);
            d_loss_real.backward();

            // Train discriminator with fake images.
            torch::Tensor noise = torch::randn(
                { batch.data.size(0), kNoiseSize, 1, 1 }, device);
            torch::Tensor fake_images = generator->forward(noise);
            torch::Tensor fake_labels
                = torch::zeros(batch.data.size(0), device);
            torch::Tensor fake_output
                = discriminator->forward(fake_images.detach())
                      .reshape(fake_labels.sizes());
            torch::Tensor d_loss_fake
                = torch::binary_cross_entropy(fake_output, fake_labels);
            d_loss_fake.backward();

            torch::Tensor d_loss = d_loss_real + d_loss_fake;
            discriminator_optimizer.step();

            // Train generator.
            generator->zero_grad();
            fake_labels.fill_(1);
            fake_output = discriminator->forward(fake_images)
                              .reshape(fake_labels.sizes());
            torch::Tensor g_loss
                = torch::binary_cross_entropy(fake_output, fake_labels);
            g_loss.backward();
            generator_optimizer.step();
            batch_index++;
            if (batch_index % kLogInterval == 0) {
                std::cout << "\r[" << epoch << "/" << kNumberOfEpochs << "]"
                          << "[" << batch_index << "/" << batches_per_epoch
                          << "]" << "D_loss: " << std::setprecision(4)
                          << d_loss.item<float>()
                          << " | G_loss: " << g_loss.item<float>() << "\n";
            }

            if (batch_index % kCheckpointEvery == 0) {
                // Checkpoint the model and optimizer state.
                torch::save(generator, "generator-checkpoint.pt");
                torch::save(generator_optimizer,
                            "generator-optimizer-checkpoint.pt");
                torch::save(discriminator, "discriminator-checkpoint.pt");
                torch::save(discriminator_optimizer,
                            "discriminator-optimizer-checkpoint.pt");
                // Sample the generator and save the images.
                torch::Tensor samples = generator->forward(torch::randn(
                    { kNumberOfSamplesPerCheckpoint, kNoiseSize, 1, 1 },
                    device));

                torch::save(
                    (samples + 1.0) / 2.0,
                    torch::str("dcgan-sample-", checkpoint_counter, ".pt"));
                std::cout << "\n-> checkpoint " << ++checkpoint_counter
                          << '\n';
            }
        }
    }

    try {
        // Кількість зображень має націло ділитися на кількість стовпців
        constexpr std::size_t image_amount { 15 };
        constexpr long columns { 5 };
        std::vector<torch::Tensor> test_images(image_amount);

        for (std::size_t i = 0; i < image_amount; ++i) {
            // Створюємо шум, другий розмір має співпадати з шумом при навчанні
            // kNoisesize
            torch::Tensor noise
                = torch::randn({ 1, kNoiseSize, 1, 1 }, device);
            // Генеруємо фейкове зображення для тестування генератора
            torch::Tensor test_image = generator->forward(noise);
            test_images[i] = std::move(test_image);
        }

        plt::clf(); // Очищає figure від попереднього малюнка і закриває його
        show_test_images(test_images, columns);
    } catch (const c10::Error& error) {
        std::cout << "LibTorch error: " << error.msg() << "\n";
    } catch (const std::exception& error) {
        std::cout << "Error: " << error.what() << "\n";
    }

    std::cout << "Training complete!" << std::endl;
    // Необхідно чітко вказувати закриття show цією командою інтерпретатора
    // Інакше буде подвійне видалення і segfault,
    // бо show закривається приховано без цього при закритті графіка
    // користувачем
    plt::detail::_interpreter::kill();
    return EXIT_SUCCESS;
}