main.cpp

#include "autoencoder.h" // CPU Model
#include "cifar10_dataset.h"
#include "gpu_autoencoder.h" // GPU Model

#include <algorithm>
#include <chrono>
#include <cstring> // Cho strcmp
#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>

// =================================================================================
// CÁC HÀM TIỆN ÍCH CHUNG (UTILS)
// =================================================================================

// Logger class (Dùng chung cho cả 2 mode)
class Logger {
public:
  Logger(const std::string &filename) {
    log_file_.open(filename);
    log_file_ << "Epoch,Loss,Time_Seconds\n";
  }
  ~Logger() {
    if (log_file_.is_open())
      log_file_.close();
  }

  void log_epoch(int epoch, float avg_loss, float duration) {
    log_file_ << epoch << "," << avg_loss << "," << duration << "\n";
    log_file_.flush();
  }

private:
  std::ofstream log_file_;
};

// Hàm lưu ảnh PPM (Dùng cho Visualization)
void save_image_ppm(const std::vector<float> &img_data,
                    const std::string &filename, int W = 32, int H = 32) {
  std::ofstream ofs(filename, std::ios::binary);
  if (!ofs)
    return;
  ofs << "P6\n" << W << " " << H << "\n255\n";
  for (int i = 0; i < W * H; ++i) {
    unsigned char r = static_cast<unsigned char>(
        std::min(255.0f, std::max(0.0f, img_data[3 * i] * 255.0f)));
    unsigned char g = static_cast<unsigned char>(
        std::min(255.0f, std::max(0.0f, img_data[3 * i + 1] * 255.0f)));
    unsigned char b = static_cast<unsigned char>(
        std::min(255.0f, std::max(0.0f, img_data[3 * i + 2] * 255.0f)));
    ofs << r << g << b;
  }
}

// Hàm gọi SVM script python
void run_svm_script() {
  std::cout << "\n>>> RUNNING SVM (Python) <<<" << std::endl;
  int ret = system("python svm_benchmark.py");
  if (ret != 0)
    std::cerr << "Warning: Could not run svm_benchmark.py" << std::endl;
}

// =================================================================================
// PHẦN 1: CPU PIPELINE
// =================================================================================

void extract_cpu(Autoencoder &model,
                 const std::vector<std::vector<float>> &images,
                 const std::vector<uint8_t> &labels, const std::string &prefix,
                 size_t limit) {
  std::string feat_file = prefix + "_features.bin";
  std::string lbl_file = prefix + "_labels.bin";
  std::ofstream out_feat(feat_file, std::ios::binary);
  std::ofstream out_lbl(lbl_file, std::ios::binary);

  size_t N = std::min(images.size(), limit);
  std::cout << "Extracting " << prefix << " (" << N << ") on CPU..."
            << std::endl;

  for (size_t i = 0; i < N; ++i) {
    std::vector<float> latent;
    model.extract_features(images[i], latent);
    out_feat.write(reinterpret_cast<char *>(latent.data()),
                   latent.size() * sizeof(float));
    out_lbl.write(reinterpret_cast<const char *>(&labels[i]), sizeof(uint8_t));
  }
  std::cout << "Done." << std::endl;
}

void run_cpu_pipeline(CIFAR10Dataset &dataset, int epochs = 20,
                      int n_images = 2000) {
  std::cout << "\n========== STARTING CPU MODE ==========" << std::endl;

  // Config: Chạy Subset nhỏ để test (vì CPU chậm)
  const int EPOCHS = epochs;
  const int NUM_IMAGES = n_images; // Subset / Full
  const int BATCH_SIZE = 32;
  const int NUM_BATCHES = NUM_IMAGES / BATCH_SIZE;

  Autoencoder model;
  std::ifstream weight_in("autoencoder_weights_cpu.bin", std::ios::binary);
  if (weight_in.good()) {
    std::cout << "Loading weights from autoencoder_weights_cpu.bin..."
              << std::endl;
    model.load_weights("autoencoder_weights_cpu.bin");
  } else {
    model.initialize_weights();
  }
  Logger logger("training_log_cpu.csv");

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

  for (int epoch = 1; epoch <= EPOCHS; ++epoch) {
    float epoch_loss = 0.0f;
    auto start_epoch = std::chrono::high_resolution_clock::now();
    size_t processed = 0;

    dataset.shuffle(42);

    for (int b = 0; b < NUM_BATCHES; ++b) {
      std::vector<std::vector<float>> batch_imgs;
      std::vector<uint8_t> batch_lbls;
      dataset.get_batch(BATCH_SIZE, batch_imgs, batch_lbls);

      for (const auto &img : batch_imgs) {
        std::vector<float> out;
        model.forward(img, out);
        epoch_loss += model.loss->compute(out, img);
        model.backward(img, img);
        model.update_weights_adam(++processed); // Step count simulation
      }
    }

    auto end_epoch = std::chrono::high_resolution_clock::now();
    double dur = std::chrono::duration_cast<std::chrono::milliseconds>(
                     end_epoch - start_epoch)
                     .count() /
                 1000.0;
    float avg_loss = epoch_loss / (NUM_BATCHES * BATCH_SIZE);

    logger.log_epoch(epoch, avg_loss, dur);
    std::cout << "Epoch " << epoch << " | Loss: " << std::fixed
              << std::setprecision(5) << avg_loss << " | Time: " << dur << "s"
              << std::endl;
  }

  auto end_total = std::chrono::high_resolution_clock::now();
  double total_dur =
      std::chrono::duration_cast<std::chrono::seconds>(end_total - start_total)
          .count();
  std::cout << "Total CPU Time: " << total_dur << "s" << std::endl;

  model.save_weights("autoencoder_weights_cpu.bin");

  // Extract & SVM
  extract_cpu(model, dataset.get_train_images(), dataset.get_train_labels(),
              "train", 2000);
  extract_cpu(model, dataset.get_test_images(), dataset.get_test_labels(),
              "test", 400);
  run_svm_script();

  // Visualization
  std::cout << "Saving CPU visualization..." << std::endl;
  for (int i = 0; i < 5; ++i) {
    std::vector<float> out;
    model.forward(dataset.get_test_images()[i], out);
    save_image_ppm(dataset.get_test_images()[i],
                   "viz_cpu_orig_" + std::to_string(i) + ".ppm");
    save_image_ppm(out, "viz_cpu_recon_" + std::to_string(i) + ".ppm");
  }
}

// =================================================================================
// PHẦN 2: GPU PIPELINE (Unified V1 & V2)
// =================================================================================

void extract_gpu(GPUAutoencoder &model,
                 const std::vector<std::vector<float>> &images,
                 const std::vector<uint8_t> &labels, const std::string &prefix,
                 size_t limit, bool use_v2 = false) { // [THÊM HAM SỐ use_v2]

  std::string feat_file = prefix + "_features.bin";
  std::string lbl_file = prefix + "_labels.bin";
  std::ofstream out_feat(feat_file, std::ios::binary);
  std::ofstream out_lbl(lbl_file, std::ios::binary);

  size_t N = std::min(images.size(), limit);

  if (use_v2)
    std::cout << "Extracting " << prefix << " (" << N
              << ") on GPU V2 (Optimized)..." << std::endl;
  else
    std::cout << "Extracting " << prefix << " (" << N
              << ") on GPU V1 (Naive)..." << std::endl;

  for (size_t i = 0; i < N; ++i) {
    std::vector<float> latent;

    // [GỌI HÀM TƯƠNG ỨNG]
    if (use_v2)
      model.extract_features_v2(images[i], latent);
    else
      model.extract_features(images[i], latent);

    out_feat.write(reinterpret_cast<char *>(latent.data()),
                   latent.size() * sizeof(float));
    out_lbl.write(reinterpret_cast<const char *>(&labels[i]), sizeof(uint8_t));

    if ((i + 1) % 1000 == 0)
      std::cout << "\r" << (i + 1) << "/" << N << std::flush;
  }
  std::cout << "\nDone." << std::endl;
}

void run_gpu_pipeline(CIFAR10Dataset &dataset, bool use_v2 = false,
                      int epochs = 20, int n_images = 50000) {
  if (use_v2)
    std::cout << "\n========== STARTING GPU V2 MODE (OPTIMIZED) =========="
              << std::endl;
  else
    std::cout << "\n========== STARTING GPU V1 MODE (NAIVE) =========="
              << std::endl;

  // Config: GPU mạnh nên chạy Full Dataset
  const int EPOCHS = epochs;
  const int NUM_TRAIN = n_images;
  const int BATCH_SIZE = 64;
  const int TOTAL_BATCHES = (NUM_TRAIN + BATCH_SIZE - 1) / BATCH_SIZE;

  GPUAutoencoder model;
  model.initialize();

  std::string log_name =
      use_v2 ? "training_log_gpu_v2.csv" : "training_log_gpu_v1.csv";
  Logger logger(log_name);

  auto start_total = std::chrono::high_resolution_clock::now();
  int step = 0;

  for (int epoch = 1; epoch <= EPOCHS; ++epoch) {
    float epoch_loss = 0.0f;
    auto start_epoch = std::chrono::high_resolution_clock::now();

    dataset.shuffle(42);

    // Batch Loop
    for (int i = 0; i < NUM_TRAIN; i += BATCH_SIZE) {
      int batch_idx = i / BATCH_SIZE;
      step++;

      model.zero_grad(); // Reset gradient

      int current_batch = std::min(BATCH_SIZE, NUM_TRAIN - i);
      float batch_loss_accum = 0.0f;

      for (int j = 0; j < current_batch; ++j) {
        model.set_input(dataset.get_train_images()[i + j]);

        // [SWITCH V1 / V2]
        if (use_v2)
          model.forward_v2();
        else
          model.forward();

        float loss = model.get_loss();
        batch_loss_accum += loss;
        epoch_loss += loss;

        if (use_v2)
          model.backward_v2();
        else
          model.backward();
      }
      model.update_weights(step);

      // Print Progress Bar
      if ((batch_idx + 1) % 5 == 0 || i + BATCH_SIZE >= NUM_TRAIN) {
        float progress = (float)(i + current_batch) / NUM_TRAIN * 100.0f;
        auto now = std::chrono::high_resolution_clock::now();
        double elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                             now - start_epoch)
                             .count() /
                         1000.0;
        double speed = (i + current_batch) / (elapsed + 1e-5);

        std::cout << "\r[Epoch " << epoch << "] "
                  << "Batch " << (batch_idx + 1) << "/" << TOTAL_BATCHES << " ("
                  << std::fixed << std::setprecision(1) << progress << "%)"
                  << " | Loss: " << std::setprecision(4)
                  << (epoch_loss / (i + current_batch))
                  << " | Speed: " << std::setprecision(1) << speed
                  << " img/s   " << std::flush;
      }
    }

    auto end_epoch = std::chrono::high_resolution_clock::now();
    double dur = std::chrono::duration_cast<std::chrono::milliseconds>(
                     end_epoch - start_epoch)
                     .count() /
                 1000.0;
    float avg_loss = epoch_loss / NUM_TRAIN;

    std::cout << std::endl; // Newline after progress bar
    logger.log_epoch(epoch, avg_loss, dur);
  }

  auto end_total = std::chrono::high_resolution_clock::now();
  double total_dur =
      std::chrono::duration_cast<std::chrono::seconds>(end_total - start_total)
          .count();
  std::cout << "\nTotal GPU Time: " << total_dur << "s" << std::endl;

  std::string weight_file = use_v2 ? "autoencoder_weights_gpu_v2.bin"
                                   : "autoencoder_weights_gpu_v1.bin";
  model.save_weights(weight_file);

  // Extract Full Data & SVM
  extract_gpu(model, dataset.get_train_images(), dataset.get_train_labels(),
              "train", 50000, use_v2);
  extract_gpu(model, dataset.get_test_images(), dataset.get_test_labels(),
              "test", 10000, use_v2);
  run_svm_script();
}

// =================================================================================
// MAIN ENTRY
// =================================================================================

int main(int argc, char *argv[]) {
  if (argc < 3) {
    std::cerr << "Usage: " << argv[0] << " <mode> <data_path>" << std::endl;
    std::cerr << "Mode: --cpu, --gpu (v1), --gpu-v2 (optimized)" << std::endl;
    return 1;
  }

  std::string mode = argv[1];
  std::string data_path = argv[2];

  int epochs = 20;
  if (argc >= 4)
    epochs = std::stoi(argv[3]);

  int n_images = 50000;
  if (argc >= 5)
    n_images = std::stoi(argv[4]);

  try {
    std::cout << "Loading dataset from: " << data_path << std::endl;
    CIFAR10Dataset dataset(data_path);
    dataset.load();
    std::cout << "Dataset loaded. Train: " << dataset.num_train()
              << ", Test: " << dataset.num_test() << std::endl;

    if (mode == "--cpu") {
      run_cpu_pipeline(dataset, epochs);
    } else if (mode == "--gpu") {
      run_gpu_pipeline(dataset, false, epochs, n_images); // V1
    } else if (mode == "--gpu-v2") {
      run_gpu_pipeline(dataset, true, epochs, n_images); // V2
    } else {
      std::cerr << "Invalid mode. Use --cpu, --gpu, or --gpu-v2" << std::endl;
      return 1;
    }

  } catch (const std::exception &e) {
    std::cerr << "Error: " << e.what() << std::endl;
    return 1;
  }

  return 0;
}