// train.h — обратный проход и оптимизатор
#pragma once
#include "model.h"
#include "checkpoint.h"
#include <vector>

namespace xt {

// Считает градиенты по всем параметрам и возвращает mean loss.
// x — входные токены [n-1], y — цели [n-1] (сдвиг на 1)
float backward(Params& p, const ForwardCache& fc,
               const int* x, const int* y, int n_pairs);

// Только loss по готовому форварду (для валидации, без обратного прохода)
float softmax_eval(const Params& p, const ForwardCache& fc,
                   const int* x, const int* y, int n_pairs, Tensor& per_token_loss);

struct TrainConfig {
    int   steps       = 1000;
    int   batch_size  = 8;
    int   block       = 64;      // длина обучающего окна
    float lr          = 3e-4f;
    float beta1       = 0.9f;
    float beta2       = 0.999f;
    float eps         = 1e-8f;
    float weight_decay = 0.01f;
    float clip        = 1.0f;
    int   warmup      = 100;
    float lr_min_frac = 0.1f;     // нижняя граница после затухания
    uint64_t seed     = 1337;
    int   threads     = 0;        // 0 => все ядра
    int   log_every   = 50;
    int   ckpt_every  = 0;        // 0 => без промежуточных сохранений
    int   val_every   = 0;        // 0 => без валидации
    int   val_tokens  = 20000;
};

// Кусочек корпуса: массив токенов + границы предложений
struct Dataset {
    std::vector<int> ids;
    size_t n() const { return ids.size(); }
};

struct TrainStats {
    float loss = 0;
    float val_loss = 0;
    int   step = 0;
    double tokens_seen = 0;
};

// Основной цикл обучения
TrainStats train_model(Params& p, const Tokenizer& tok, const std::string& corpus_path,
                       const TrainConfig& tc, const std::string& out_path,
                       int resume_step = 0);

}  // namespace xt