// model.h — конфигурация, параметры и прямой проход трансформера
#pragma once
#include "tensor.h"
#include <string>
#include <vector>

namespace xt {

// Магия формата: 'X','N','H','1' в little-endian
constexpr uint32_t XNH_MAGIC = 0x31484E58;

struct Config {
    int   vocab_size = 256;
    int   n_layer    = 4;
    int   n_head     = 4;
    int   n_embd     = 128;
    int   block_size = 128;   // максимальный контекст
    int   ffn_dim    = 0;     // 0 => 4*n_embd
    int   rope_base  = 10000;
    float rms_eps    = 1e-5f;
    float init_std   = 0.02f;
    int   n_vocab_out = 0;    // 0 => tied (используем wte как выходную матрицу)

    int head_dim()  const { return n_embd / n_head; }
    int ffn()       const { return ffn_dim > 0 ? ffn_dim : 4 * n_embd; }
    int vocab_out() const { return n_vocab_out > 0 ? n_vocab_out : vocab_size; }
    bool valid(std::string& err) const;
};

// Все обучаемые веса. Именно этот список сериализуется в файл модели.
struct Params {
    Config cfg;

    Tensor wte;                       // [vocab, n_embd] — вход
    Tensor lm_head;                   // [n_embd, vocab_out] — только если НЕ tied
    std::vector<Tensor> wq, wk, wv, wo;
    std::vector<Tensor> rms_attn, rms_ffn;
    std::vector<Tensor> w1, w2, w3;  // SwiGLU
    Tensor rms_final;                 // [1, n_embd]

    // Градиенты (тет же порядок, что у all())
    Tensor gwte, glm_head, grms_final;
    std::vector<Tensor> gq, gk, gv, go, grms_attn, grms_ffn, g1, g2, g3;

    // Моменты Adam
    std::vector<Tensor> m, v;

    size_t n_params() const;
    bool   tied() const;
    void   alloc_grads();
    void   init(uint64_t seed = 1234);
    void   zero_grad();
    void   adam_step(float lr, float b1, float b2, float eps, float wd,
                     int step, float clip);

    std::vector<const Tensor*> all() const;      // только для чтения/сохранения
    std::vector<Tensor*> all_w();                 // для изменения весов
    std::vector<Tensor*> all_grads();
};

// Кэш прямого прохода, нужен обратному
struct LayerCache {
    Tensor x;        // вход в слой
    Tensor xb;       // нормализованный вход (attention)
    Tensor q, k, v;  // [B, n_embd], RoPE применён к q,k
    Tensor att;      // [n_head, B, B] — веса внимания
    Tensor attout;   // [B, n_embd] — конкатенация голов
    Tensor proj;     // [B, n_embd] — после wo
    Tensor x2;       // после residual attention
    Tensor x2b;      // нормализация перед FFN
    Tensor h1, h3, ha;
    Tensor fo;       // [B, n_embd] — после w2
    Tensor inv_rms, inv_rms2;
};

struct ForwardCache {
    std::vector<LayerCache> layers;
    Tensor x;        // текущий residual-поток
    Tensor xf;       // после финального RMSNorm
    Tensor logits;   // [B, vocab_out]
    Tensor inv_rms;  // финальная нормализация
    std::vector<float> rope_cos, rope_sin;
    int block    = 0;   // длина одной последовательности
    int seq_len  = 0;   // то же, явно
    int n_seq    = 1;   // сколько последовательностей в батче
};

void rope_tables(int block, int head_dim, int base,
                 std::vector<float>& cs, std::vector<float>& sn);

// Полный прямой проход.
//   n_tok   — сколько токенов всего (может быть кратно block_size: батч)
//   seq_len — длина ОДНОЙ последовательности (<= block_size).
//             Внимание и RoPE строятся независимо для каждой
//             последовательности, токены разных батчей не видят друг друга.
void forward(const Params& p, const int* tokens, int n_tok, int seq_len, ForwardCache& fc);

// Обёртка для одиночной последовательности
void forward(const Params& p, const int* tokens, int n_tok, ForwardCache& fc);

// Только последний ряд логитов
void forward_last(const Params& p, const int* tokens, int n_tok, Tensor& logits);

}  // namespace xt