// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "train/trainer.h"

#include "core/fmodel.h"
#include "ops/metal/metal_ops.h"
#include "ops/ops.h"
#include "train/checkpoint.h"
#include "train/scheduler.h"

#include <Foundation/Foundation.hpp>

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <filesystem>

namespace forge::train {

Trainer::Trainer(Config cfg, const std::string& data_dir, const std::string& out_dir,
                 const std::string& config_json)
    : cfg_(cfg), out_dir_(out_dir), config_json_(config_json) {
    std::filesystem::create_directories(out_dir_);
    model_ = std::make_unique<nn::Transformer>(cfg_.model, cfg_.train.seed);

    Optimizer::Options opts;
    opts.kind = cfg_.train.optimizer;
    opts.beta1 = cfg_.train.beta1;
    opts.beta2 = cfg_.train.beta2;
    opts.eps = cfg_.train.eps;
    opts.weight_decay = cfg_.train.weight_decay;
    opts.muon_momentum = cfg_.train.muon_momentum;
    opts.muon_lr_ratio = cfg_.train.muon_lr / cfg_.train.lr;
    opt_ = std::make_unique<Optimizer>(model_->named_parameters(), opts);

    train_data_ = std::make_unique<DataLoader>(data_dir + "/train.bin",
                                               cfg_.model.context_length,
                                               cfg_.train.seed);
    const std::string val_path = data_dir + "/val.bin";
    if (std::filesystem::exists(val_path)) {
        val_data_ = std::make_unique<DataLoader>(val_path, cfg_.model.context_length,
                                                 cfg_.train.seed + 1);
    }
}

float Trainer::eval_loss() {
    if (!val_data_) return -1.0f;
    NoGrad ng;
    const int64_t B = cfg_.train.batch_size, T = cfg_.model.context_length;
    double total = 0.0;
    for (int64_t i = 0; i < cfg_.train.eval_batches; ++i) {
        Tensor ids = Tensor::empty({B, T}, DType::I32);
        Tensor targets = Tensor::empty({B * T}, DType::I32);
        val_data_->seq_batch(i, ids, targets);
        Var loss = model_->loss(ids, targets);
        if (ops::backend() == ops::Backend::Metal) metal::sync();
        total += double(loss.value().data<float>()[0]);
    }
    return float(total / double(cfg_.train.eval_batches));
}

void Trainer::save(int64_t step) {
    CheckpointData meta;
    meta.step = step;
    meta.rng_state = uint64_t(step); // dataloader reseeded from step on resume
    meta.config_json = config_json_;
    char name[64];
    std::snprintf(name, sizeof(name), "/ckpt_%06lld.bin", static_cast<long long>(step));
    save_checkpoint(out_dir_ + name, model_->named_parameters(), opt_.get(), meta);
    save_checkpoint(out_dir_ + "/ckpt_latest.bin", model_->named_parameters(), opt_.get(),
                    meta);
    std::printf("checkpoint saved: %s\n", (out_dir_ + name).c_str());

    // Native .forge commit: the run's whole weight history lives in one
    // git-style repo; content addressing means an unchanged tensor (frozen,
    // masked, tied) is never rewritten.
    if (cfg_.train.forge_save) {
        fmodel::SaveOptions fo;
        fo.dtype = cfg_.train.forge_dtype == "f16" ? DType::F16
                 : cfg_.train.forge_dtype == "bf16" ? DType::BF16
                                                    : DType::F32;
        char tag[32];
        std::snprintf(tag, sizeof(tag), "step-%06lld", static_cast<long long>(step));
        fo.tag = tag;
        fo.step = step;
        fmodel::save(out_dir_ + "/model.forge", config_json_,
                     model_->named_parameters(), fo);
    }
}

void Trainer::train(const std::string& resume_from) {
    const auto& tc = cfg_.train;
    const int64_t B = tc.batch_size, T = cfg_.model.context_length;

    int64_t start_step = 0;
    if (!resume_from.empty()) {
        CheckpointData meta =
            load_checkpoint(resume_from, model_->named_parameters(), opt_.get());
        start_step = meta.step;
        std::printf("resumed from %s at step %lld\n", resume_from.c_str(),
                    static_cast<long long>(start_step));
    }

    FILE* csv = std::fopen((out_dir_ + "/log.csv").c_str(),
                           start_step > 0 ? "ab" : "wb");
    if (csv && start_step == 0)
        std::fprintf(csv, "step,loss,lr,grad_norm,tokens_per_sec,val_loss,elapsed_s\n");
    const auto run_t0 = std::chrono::steady_clock::now();

    const int64_t min_lr_steps = tc.max_steps;
    const float min_lr = tc.lr * tc.min_lr_ratio;
    const int64_t tokens_per_step = B * T * tc.grad_accum_steps;
    const bool wsd = tc.schedule == "wsd";
    const int64_t wsd_decay_steps =
        std::max<int64_t>(1, int64_t(float(tc.max_steps) * tc.wsd_decay_frac));

    std::printf("training %s: %lld params, %lld steps, %lld tokens/step, backend=%s, "
                "opt=%s, sched=%s\n",
                cfg_.model.name.c_str(),
                static_cast<long long>(cfg_.model.num_params()),
                static_cast<long long>(tc.max_steps),
                static_cast<long long>(tokens_per_step),
                ops::backend() == ops::Backend::Metal ? "metal" : "cpu",
                tc.optimizer.c_str(), tc.schedule.c_str());

    for (int64_t step = start_step; step < tc.max_steps; ++step) {
        NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();
        const auto t0 = std::chrono::steady_clock::now();
        const float lr = wsd
            ? lr_wsd(step, tc.lr, min_lr, tc.warmup_steps, tc.max_steps, wsd_decay_steps)
            : lr_at(step, tc.lr, min_lr, tc.warmup_steps, min_lr_steps);

        opt_->zero_grad();
        const bool on_gpu = ops::backend() == ops::Backend::Metal;
        double loss_val = 0.0;
        for (int64_t micro = 0; micro < tc.grad_accum_steps; ++micro) {
            Tensor ids = Tensor::empty({B, T}, DType::I32);
            Tensor targets = Tensor::empty({B * T}, DType::I32);
            train_data_->next_batch(ids, targets);
            Var loss = model_->loss(ids, targets);
            // scale so accumulated grads average over micro-batches
            Var scaled = ops::scale(loss, 1.0f / float(tc.grad_accum_steps));
            Tape::get().backward(scaled);
            // Sync per micro-batch, not per step: pooled buffers freed while a
            // command buffer is open sit on the allocator's retire list until
            // the next sync, so without this every micro-batch's activations
            // stay resident and peak memory scales with grad_accum_steps.
            if (on_gpu) metal::sync();
            loss_val += double(loss.value().data<float>()[0]) /
                        double(tc.grad_accum_steps);
        }

        const float grad_norm = opt_->step_with_clip(lr, tc.grad_clip);
        if (cfg_.model.n_experts > 0 && cfg_.model.moe_bias_gamma > 0.0f)
            model_->update_moe_bias(cfg_.model.moe_bias_gamma);

        const auto t1 = std::chrono::steady_clock::now();
        const double dt = std::chrono::duration<double>(t1 - t0).count();
        const double tps = double(tokens_per_step) / dt;

        float val = -1.0f;
        if (tc.eval_every > 0 && (step + 1) % tc.eval_every == 0) val = eval_loss();

        if (step < 10 || step % 10 == 0 || val >= 0.0f) {
            std::printf("step %6lld | loss %.4f | lr %.2e | gnorm %.3f | %.0f tok/s",
                        static_cast<long long>(step), loss_val, double(lr),
                        double(grad_norm), tps);
            if (val >= 0.0f) std::printf(" | val %.4f", double(val));
            std::printf("\n");
            std::fflush(stdout);
        }
        if (csv) {
            const double elapsed =
                std::chrono::duration<double>(t1 - run_t0).count();
            std::fprintf(csv, "%lld,%.6f,%.6e,%.6f,%.1f,%.6f,%.3f\n",
                         static_cast<long long>(step), loss_val, double(lr),
                         double(grad_norm), tps, double(val), elapsed);
            std::fflush(csv);
        }

        if (tc.checkpoint_every > 0 && (step + 1) % tc.checkpoint_every == 0)
            save(step + 1);

        pool->drain();
    }

    save(tc.max_steps);
    if (csv) std::fclose(csv);
}

} // namespace forge::train