#include "data_loader.hpp"
#include "training.hpp"
#include <iostream>

TrainConfig parse_cmd(int argc, char** argv) {
    TrainConfig cfg;
    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        auto next = [&](const char* name) -> std::string {
            if (i + 1 >= argc) { std::cerr << "Missing value for " << name << "\n"; std::exit(1); }
            return std::string(argv[++i]);
        };
        if      (a == "--data_dir") cfg.data_dir = next("--data_dir");
        else if (a == "--epochs")   cfg.epochs = std::stoi(next("--epochs"));
        else if (a == "--batch")    cfg.batch  = std::stoi(next("--batch"));
        else if (a == "--lr")       cfg.lr     = std::stof(next("--lr"));
        else if (a == "--fungi")    cfg.fungi_count = std::stoi(next("--fungi"));
        else if (a == "--wd")       cfg.wd = std::stof(next("--wd"));
        else if (a == "--accuracy_tolerance") cfg.accuracy_tolerance = std::stof(next("--accuracy_tolerance"));
        else if (a == "--seed")      cfg.seed = std::stoul(next("--seed")); // Original default seed was 1337u
        else std::cerr << "Unknown arg: " << a << "\n";
    }
    return cfg;
}

int main(int argc, char** argv) {
    try {
        TrainConfig cfg = parse_cmd(argc, argv);

        std::cout << "[INFO] Loading Fashion-MNIST training data from: " << cfg.data_dir << "\n";
        FashionMNISTSet train = load_fashion_mnist_data(cfg.data_dir, true);

        std::cout << "[INFO] Loading Fashion-MNIST test data from: " << cfg.data_dir << "\n";
        FashionMNISTSet test = load_fashion_mnist_data(cfg.data_dir, false);

        train_model(train, test, cfg);

        std::cout << "[DONE] Training finished.\n";
        return 0;
    } catch (const std::exception& e) {
        std::cerr << "[FATAL] " << e.what() << "\n";
        return 1;
    }
}