// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include "core/autograd.h"

#include <string>
#include <utility>
#include <vector>

namespace forge::train {

// Two optimizer kinds behind one interface:
//  - "adamw": AdamW with decoupled weight decay on every parameter
//    (llm.c/PyTorch convention: eps outside sqrt, wd only on dim>=2 params).
//  - "muon": Keller Jordan's Muon on 2-D hidden matrices — momentum
//    orthogonalized by 5 Newton-Schulz iterations — with AdamW kept for
//    embeddings / lm_head / 1-D params (the standard Muon split). The Muon
//    group's LR is lr * muon_lr_ratio so both groups follow one schedule.
// Deduplicates tied parameters by Var::id(). m_ doubles as Muon's momentum
// buffer; v_ stays zero for Muon params so the checkpoint format is
// identical across kinds.
class Optimizer {
public:
    struct Options {
        std::string kind = "adamw";  // "adamw" | "muon"
        float beta1 = 0.9f;
        float beta2 = 0.95f;
        float eps = 1e-8f;
        float weight_decay = 0.1f;
        float muon_momentum = 0.95f;
        float muon_lr_ratio = 1.0f;  // muon lr = lr * ratio
    };

    Optimizer(const std::vector<std::pair<std::string, Var>>& named_params, Options opts);

    // Global-norm gradient clip: returns the pre-clip norm and scales all
    // grads by min(1, max_norm/norm). No-op when max_norm <= 0. (CPU path.)
    float clip_global_norm(float max_norm);

    void step(float lr);
    void zero_grad();

    // Backend-routed step: clip + update in one call, returns the pre-clip
    // grad norm. On Metal it expects backward() already encoded: it encodes
    // per-tensor sumsq, syncs (this is the step's loss-readback boundary),
    // folds the clip factor into the update, and syncs again at the end.
    // Muon adds one extra sync: the Newton-Schulz input must be normalized
    // by its Frobenius norm, which the CPU reads between the two phases.
    float step_with_clip(float lr, float max_norm);

    // Allocate m/v now (checkpoint loading needs the buffers to exist).
    void ensure_state();

    int64_t t() const { return t_; }
    // Checkpoint access (M4): moments in parameter order.
    std::vector<Tensor>& m() { return m_; }
    std::vector<Tensor>& v() { return v_; }
    void set_t(int64_t t) { t_ = t; }
    const std::vector<Var>& params() const { return params_; }

private:
    // AdamW update for one param (CPU path).
    void adamw_cpu(size_t pi, float lr);
    // Muon update for one param (CPU path); grads already clipped.
    void muon_cpu(size_t pi, float lr);

    Options opts_;
    std::vector<Var> params_;
    std::vector<bool> decay_;   // dim >= 2
    std::vector<bool> muon_;    // kind=="muon" && 2-D && not embedding/head
    std::vector<Tensor> m_, v_; // f32, allocated lazily at first step
    int64_t t_ = 0;
};

// Historical name; checkpoints and tests predate the Muon mode.
using AdamW = Optimizer;

} // namespace forge::train