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

#include "ops/cpu/cpu_ops.h"
#include "ops/metal/metal_ops.h"
#include "ops/ops.h"

#include <cmath>
#include <unordered_set>

namespace forge::train {

namespace {

// Newton-Schulz quintic coefficients (Keller Jordan's Muon: tuned so the
// iteration's fixed point spreads singular values toward 1 fast; 5 rounds
// suffice at bf16-level accuracy, and we run f32).
constexpr float kNsA = 3.4445f;
constexpr float kNsB = -4.7750f;
constexpr float kNsC = 2.0315f;
constexpr int kNsIters = 5;
constexpr float kNsEps = 1e-7f;

// Muon updates hidden 2-D matrices only; embeddings and the (possibly tied)
// lm_head keep AdamW, matching the reference MuonWithAuxAdam split.
bool muon_eligible(const std::string& name, const Var& p) {
    if (p.value().ndim() != 2) return false;
    if (name.find("emb") != std::string::npos) return false;     // tok_emb/pos_emb
    if (name.find("lm_head") != std::string::npos) return false;
    return true;
}

// LR adjustment for rectangular matrices (Muon reference impl).
float muon_adj(const Tensor& w) {
    const float r = float(w.size(0)), c = float(w.size(1));
    return std::sqrt(std::max(1.0f, r / c));
}

float frobenius(const Tensor& x) {
    double sq = 0.0;
    const float* p = x.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) sq += double(p[i]) * double(p[i]);
    return float(std::sqrt(sq));
}

} // namespace

Optimizer::Optimizer(const std::vector<std::pair<std::string, Var>>& named_params,
                     Options opts)
    : opts_(opts) {
    std::unordered_set<const void*> seen;
    for (const auto& [name, p] : named_params) {
        if (!p.defined() || !p.requires_grad()) continue;
        if (!seen.insert(p.id()).second) continue; // tied param, already tracked
        params_.push_back(p);
        decay_.push_back(p.value().ndim() >= 2);
        muon_.push_back(opts_.kind == "muon" && muon_eligible(name, p));
    }
}

float Optimizer::clip_global_norm(float max_norm) {
    double sq = 0.0;
    for (const Var& p : params_) {
        if (!p.has_grad()) continue;
        const float* g = p.grad().data<float>();
        for (int64_t i = 0; i < p.grad().numel(); ++i) sq += double(g[i]) * double(g[i]);
    }
    const float norm = float(std::sqrt(sq));
    if (max_norm > 0.0f && norm > max_norm) {
        const float s = max_norm / norm;
        for (const Var& p : params_) {
            if (!p.has_grad()) continue;
            float* g = p.grad().data<float>();
            for (int64_t i = 0; i < p.grad().numel(); ++i) g[i] *= s;
        }
    }
    return norm;
}

void Optimizer::ensure_state() {
    if (!m_.empty()) return;
    m_.reserve(params_.size());
    v_.reserve(params_.size());
    for (const Var& p : params_) {
        m_.push_back(Tensor::zeros(p.value().shape()));
        v_.push_back(Tensor::zeros(p.value().shape()));
    }
}

float Optimizer::step_with_clip(float lr, float max_norm) {
    if (ops::backend() == ops::Backend::CPU) {
        const float norm = clip_global_norm(max_norm);
        step(lr);
        return norm;
    }

    // Metal path. Encode per-tensor sum-of-squares into one partials buffer,
    // then sync — this is also the boundary where the step's loss becomes
    // readable.
    // These are hundreds of independent single-threadgroup reductions (one per
    // parameter tensor), which a serial encoder pointlessly serializes.
    Tensor partials = Tensor::empty({int64_t(params_.size())});
    {
        metal::ConcurrentRegion region;
        for (size_t i = 0; i < params_.size(); ++i) {
            Tensor slot = partials.slice0(int64_t(i), 1);
            metal::sumsq(params_[i].grad(), slot);
        }
    }
    metal::sync();

    double sq = 0.0;
    for (size_t i = 0; i < params_.size(); ++i) sq += double(partials.data<float>()[i]);
    const float norm = float(std::sqrt(sq));
    const float grad_scale =
        (max_norm > 0.0f && norm > max_norm) ? max_norm / norm : 1.0f;

    ensure_state(); // CPU-side zeros are safe here: the stream is closed
    ++t_;
    {
        // Each parameter's AdamW update touches only its own w/g/m/v — independent.
        metal::ConcurrentRegion region;
        for (size_t i = 0; i < params_.size(); ++i) {
            if (muon_[i]) continue;
            const float wd = decay_[i] ? opts_.weight_decay : 0.0f;
            metal::adamw_step(params_[i].value(), params_[i].grad(), m_[i], v_[i], lr,
                              opts_.beta1, opts_.beta2, t_, opts_.eps, wd, grad_scale);
        }
    }

    // Muon phase 1: momentum + nesterov direction + its Frobenius norm.
    // (Serial encoder: each param's chain has read-after-write dependencies;
    // cross-param over-ordering is harmless — these are tiny dispatches.)
    std::vector<Tensor> xs(params_.size());
    Tensor mu_partials;
    std::vector<int64_t> mu_slot(params_.size(), -1);
    int64_t n_muon = 0;
    for (size_t i = 0; i < params_.size(); ++i)
        if (muon_[i]) mu_slot[i] = n_muon++;
    if (n_muon > 0) {
        mu_partials = Tensor::empty({n_muon});
        const float beta = opts_.muon_momentum;
        for (size_t i = 0; i < params_.size(); ++i) {
            if (!muon_[i]) continue;
            const Tensor& g = params_[i].grad();
            Tensor gs = Tensor::empty(g.shape());
            xs[i] = Tensor::empty(g.shape());
            metal::scale(g, grad_scale, gs);      // g̃ = clip-scaled grad
            metal::scale(m_[i], beta, m_[i]);
            metal::accumulate(m_[i], gs);         // m = β·m + g̃
            metal::scale(m_[i], beta, xs[i]);
            metal::accumulate(xs[i], gs);         // X = g̃ + β·m (nesterov)
            Tensor slot = mu_partials.slice0(mu_slot[i], 1);
            metal::sumsq(xs[i], slot);
        }
        metal::sync();

        // Muon phase 2: normalize, orthogonalize (Newton-Schulz via the
        // existing matmul kernels), apply. For rows > cols the iteration runs
        // on the implicit transpose: A = XᵀX and X ← aX + XB, which is the
        // transpose of the tall-side recurrence — no materialized transpose.
        const float lr_m = lr * opts_.muon_lr_ratio;
        for (size_t i = 0; i < params_.size(); ++i) {
            if (!muon_[i]) continue;
            Tensor& X = xs[i];
            const int64_t r = X.size(0), c = X.size(1);
            const bool tall = r > c;
            const int64_t k = tall ? c : r;
            const float fro = std::sqrt(mu_partials.data<float>()[mu_slot[i]]);
            metal::scale(X, 1.0f / (fro + kNsEps), X);

            Tensor A = Tensor::empty({k, k});
            Tensor B = Tensor::empty({k, k});
            Tensor X2 = Tensor::empty({r, c});
            for (int it = 0; it < kNsIters; ++it) {
                if (tall) metal::matmul(X, X, A, true, false);   // A = XᵀX
                else      metal::matmul(X, X, A, false, true);   // A = XXᵀ
                metal::matmul(A, A, B);
                metal::scale(B, kNsC, B);
                metal::scale(A, kNsB, A);
                metal::accumulate(B, A);                          // B = bA + cA²
                if (tall) metal::matmul(X, B, X2);                // X2 = XB
                else      metal::matmul(B, X, X2);                // X2 = BX
                metal::scale(X, kNsA, X);
                metal::accumulate(X, X2);                         // X = aX + X2
            }
            metal::scale(params_[i].value(), 1.0f - lr_m * opts_.weight_decay,
                         params_[i].value());
            metal::scale(X, -lr_m * muon_adj(params_[i].value()), X);
            metal::accumulate(params_[i].value(), X);
        }
    }

    metal::sync(); // weights final before the next step's CPU-side zero_grad
    return norm;
}

void Optimizer::adamw_cpu(size_t pi, float lr) {
    const float bc1 = 1.0f - std::pow(opts_.beta1, float(t_));
    const float bc2 = 1.0f - std::pow(opts_.beta2, float(t_));
    Var& p = params_[pi];
    const float wd = decay_[pi] ? opts_.weight_decay : 0.0f;
    float* w = p.value().data<float>();
    const float* g = p.grad().data<float>();
    float* m = m_[pi].data<float>();
    float* v = v_[pi].data<float>();
    const int64_t n = p.value().numel();
    for (int64_t i = 0; i < n; ++i) {
        m[i] = opts_.beta1 * m[i] + (1.0f - opts_.beta1) * g[i];
        v[i] = opts_.beta2 * v[i] + (1.0f - opts_.beta2) * g[i] * g[i];
        const float mhat = m[i] / bc1;
        const float vhat = v[i] / bc2;
        w[i] -= lr * (mhat / (std::sqrt(vhat) + opts_.eps) + wd * w[i]);
    }
}

void Optimizer::muon_cpu(size_t pi, float lr) {
    Var& p = params_[pi];
    const float beta = opts_.muon_momentum;
    const float* g = p.grad().data<float>();
    float* m = m_[pi].data<float>();
    const int64_t n = p.value().numel();
    const int64_t r = p.value().size(0), c = p.value().size(1);
    const bool tall = r > c;
    const int64_t k = tall ? c : r;

    Tensor X = Tensor::empty({r, c});
    float* x = X.data<float>();
    for (int64_t i = 0; i < n; ++i) {
        m[i] = beta * m[i] + g[i];
        x[i] = g[i] + beta * m[i]; // nesterov
    }
    const float fro = frobenius(X);
    for (int64_t i = 0; i < n; ++i) x[i] /= (fro + kNsEps);

    Tensor A = Tensor::empty({k, k});
    Tensor B = Tensor::empty({k, k});
    Tensor X2 = Tensor::empty({r, c});
    for (int it = 0; it < kNsIters; ++it) {
        if (tall) cpu::matmul(X, X, A, true, false);
        else      cpu::matmul(X, X, A, false, true);
        cpu::matmul(A, A, B);
        float* pa = A.data<float>();
        float* pb = B.data<float>();
        for (int64_t i = 0; i < k * k; ++i) pb[i] = kNsB * pa[i] + kNsC * pb[i];
        if (tall) cpu::matmul(X, B, X2);
        else      cpu::matmul(B, X, X2);
        const float* p2 = X2.data<float>();
        for (int64_t i = 0; i < n; ++i) x[i] = kNsA * x[i] + p2[i];
    }

    const float lr_m = lr * opts_.muon_lr_ratio;
    const float adj = muon_adj(p.value());
    float* w = p.value().data<float>();
    for (int64_t i = 0; i < n; ++i)
        w[i] = w[i] * (1.0f - lr_m * opts_.weight_decay) - lr_m * adj * x[i];
}

void Optimizer::step(float lr) {
    ensure_state();
    ++t_; // 1-based: bias correction divides by (1 - beta^t)
    for (size_t pi = 0; pi < params_.size(); ++pi) {
        if (!params_[pi].has_grad()) continue;
        if (muon_[pi]) muon_cpu(pi, lr);
        else adamw_cpu(pi, lr);
    }
}

void Optimizer::zero_grad() {
    for (const Var& p : params_) p.zero_grad();
}

} // namespace forge::train