// metal/src/engine/scr_state.hpp - the Metal window's session state (window.cpp) as bytes, and the SCR
// relocation of one QSA layer's cells (qsa.metal's cache layout).  Header-only: shared by the window patch and the
// CPU probe, so the probe tests the code the engine would run.
//
// window.cpp's state, per its init (L554-567) and checkpoint_save (L1688):
//   36 GDN layers (l % 4 != 3): gst_ f32 [128 i][48 h][128 j] = 3 MiB, ghist_[hp_] f32 [10240 ch][3] = 120 KiB
//   pleHist_ f32 [10240 ch][9] = 360 KiB, and ss.ple_prev[2] (the PLE n-gram token window)
//   12 QSA layers (l % 4 == 3), per cell c and kv head h (2):
//     kq_/vq_ int8 [c][h][256], ksc_/vsc_ half [c][h][4] (one scale per 64 dims)
//     pooled_ f32 [b][128] for complete blocks b = c / 4, tail_ f32 [4][128] (raw keys, slot c % 4), dead_ f32 [128]
// qsa_prep (qsa.metal L104-147) applies NEOX RoPE at position = cell (pos_base is 0 on every path) to dims 0..63
// only, pairs (i, i + 32), i < 32, angle cell * base^(-2i/64); k is rotated BEFORE its int8 quantization, so the
// rotated dims are exactly scale group 0 of each row.  pooled[b] is rotated at position 4 b, also dims 0..63.
#pragma once
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>

namespace scr {

constexpr int NGDN = 36, NQSA = 12, QKV = 2, QD = 256, QROT = 64, ID = 128, NG_HIST = 9, CH = 10240;
constexpr size_t GST_BYTES = (size_t) 4 * 128 * 48 * 128;       // 3,145,728
constexpr size_t GHIST_BYTES = (size_t) 4 * CH * 3;              // 122,880
constexpr size_t PLE_BYTES = (size_t) 4 * NG_HIST * CH;          // 368,640
constexpr size_t TAIL_BYTES = (size_t) 4 * 4 * ID, DEAD_BYTES = (size_t) 4 * ID;
constexpr size_t KQ_CELL = (size_t) QKV * QD;                     // int8 bytes per cell (kq or vq)
constexpr size_t SC_CELL = (size_t) QKV * 4 * 2;                  // half bytes per cell (ksc or vsc)
constexpr size_t POOL_BLOCK = (size_t) 4 * ID;                    // f32 bytes per pooled block (4 cells)

// bytes of a saved session: the running state, then the positional cells [0, n) of every QSA layer
constexpr size_t fixed_bytes() { return NGDN * (GST_BYTES + GHIST_BYTES) + PLE_BYTES + NQSA * (TAIL_BYTES + DEAD_BYTES); }
constexpr size_t cell_bytes(size_t n) { return NQSA * (n * (2 * KQ_CELL + 2 * SC_CELL) + (n / 4) * POOL_BLOCK); }

// the file: this header, the token ids, then the sections in a fixed order, each starting on a 16 KiB boundary
// (the page size: F_NOCACHE reads into page-aligned MTLBuffer storage go straight to the device)
constexpr size_t ALIGN = 16384;
inline size_t align_up(size_t x) { return (x + ALIGN - 1) / ALIGN * ALIGN; }
struct SessionHeader {
    char magic[8] = {'S', 'T', 'R', 'S', 'E', 'S', 'S', '1'};
    uint32_t version = 1;
    uint32_t flags = 0;          // bit 0: PLE on, bit 1: MTP sections follow
    int64_t n_cells = 0;         // positions [0, n_cells) are in the cache; the next token runs at n_cells
    int64_t n_ids = 0;           // the token ids of those positions (the harness's prefix check)
    uint64_t model_bytes = 0;    // the model shard's size and its tensor count: a session is one model's
    uint64_t model_tensors = 0;
    double rope_base = 0;
    int32_t ple_prev[2] = {-1, -1};
    uint64_t payload_bytes = 0;  // everything after the header and ids, for the size check before any copy
};
static_assert(sizeof(SessionHeader) % 8 == 0, "packed as written");

// ---- fp16 and the qsa_prep int8 rule (L136-147): per 64 dims sh = half(amax / 127), q = clamp(rint(x / sf))
inline float h2f(uint16_t h) { _Float16 v; std::memcpy(&v, &h, 2); return (float) v; }
inline uint16_t f2h(float f) { _Float16 v = (_Float16) f; uint16_t h; std::memcpy(&h, &v, 2); return h; }
inline void quant64(const float* x, int8_t* q, uint16_t* sc) {
    float m = 0.0f;
    for (int i = 0; i < 64; ++i) m = std::max(m, std::fabs(x[i]));
    *sc = f2h(m / 127.0f);
    const float sf = h2f(*sc);
    for (int i = 0; i < 64; ++i) q[i] = (int8_t) (sf > 0.0f ? std::clamp((int) std::nearbyint(x[i] / sf), -127, 127) : 0);
}

// the rotation by delta positions: cos/sin of delta * theta_i in double, as window.cpp's table (L545-551)
struct Rot {
    float c[32], s[32];
    Rot(double base, int64_t delta) {
        for (int i = 0; i < 32; ++i) {
            const double th = (double) delta * std::pow(base, -2.0 * i / 64.0);
            c[i] = (float) std::cos(th);
            s[i] = (float) std::sin(th);
        }
    }
};
// the first 64 dims of a row rotated in place (norm_rope's pairs: a = x[i], b = x[i + 32])
inline void rotate64(float* x, const Rot& r) {
    for (int i = 0; i < 32; ++i) {
        const float a = x[i], b = x[i + 32];
        x[i] = a * r.c[i] - b * r.s[i];
        x[i + 32] = a * r.s[i] + b * r.c[i];
    }
}

// One QSA layer's span of n cells moved from c_src to c_dst = c_src + delta, keys re-rotated by delta.
//   kq: only group 0 (dims 0..63, scale ksc[.][.][0]) is rotated: dequant, rotate, requant with a fresh scale.
//       dims 64..255 and scales 1..3 move as bytes.  vq / vsc carry no position: bytes.
//   pooled: the complete blocks inside the span, b -> b + delta / 4, rotated by delta (4 b -> 4 b + delta).  delta must
//       be a multiple of 4, else every block of the span would mix different cells and the raw keys aren't kept.
//       The junction block (the one holding c_dst when c_dst % 4 != 0) is not moved: it mixes the new span before it,
//       so the caller recomputes the span's first (4 - c_dst % 4) % 4 cells or more (scr_min_junction).
// Cells are walked away from the overlap (ascending when moving down), so source rows are read before they're hit.
inline int64_t scr_min_junction(int64_t c_dst) { return (4 - (c_dst & 3)) & 3; }
inline bool relocate_layer(int8_t* kq, int8_t* vq, uint16_t* ksc, uint16_t* vsc, float* pooled, int64_t c_src,
                           int64_t c_dst, int64_t n, const Rot& rot) {
    const int64_t delta = c_dst - c_src;
    if (delta == 0 || n <= 0) return true;
    if ((delta & 3) != 0 || c_dst < 0) return false;
    auto cell = [&](int64_t k) {
        const int64_t s = c_src + k, d = c_dst + k;
        for (int h = 0; h < QKV; ++h) {
            const size_t rs = (size_t) s * QKV + h, rd = (size_t) d * QKV + h;
            float x[64];
            const float sf = h2f(ksc[rs * 4]);
            for (int i = 0; i < 64; ++i) x[i] = (float) kq[rs * QD + i] * sf;
            rotate64(x, rot);
            std::memmove(kq + rd * QD + 64, kq + rs * QD + 64, QD - 64);
            std::memmove(ksc + rd * 4 + 1, ksc + rs * 4 + 1, 3 * 2);
            quant64(x, kq + rd * QD, ksc + rd * 4);
            std::memmove(vq + rd * QD, vq + rs * QD, QD);
            std::memmove(vsc + rd * 4, vsc + rs * 4, 4 * 2);
        }
    };
    const int64_t b0 = (c_src + 3) / 4, b1 = (c_src + n) / 4, db = delta / 4;   // complete blocks [b0, b1)
    auto block = [&](int64_t b) {
        float x[ID];
        std::memcpy(x, pooled + (size_t) b * ID, sizeof x);
        rotate64(x, rot);
        std::memcpy(pooled + (size_t) (b + db) * ID, x, sizeof x);
    };
    if (delta < 0) {
        for (int64_t k = 0; k < n; ++k) cell(k);
        for (int64_t b = b0; b < b1; ++b) block(b);
    } else {
        for (int64_t k = n - 1; k >= 0; --k) cell(k);
        for (int64_t b = b1 - 1; b >= b0; --b) block(b);
    }
    return true;
}

// ---- GDN staleness, the part with a closed form.  gdn.metal (L89-119): S <- g S + k delta^T, delta = (v - g S^T k) b,
// i.e. S' = g (I - b k k^T) S + b k v^T with |k| <= 1 (k = cv / sqrt(sum cv^2 + eps)) and b = sigmoid in (0, 1).  Two
// runs over the same C inputs from states S_1, S_2 differ by D' = g (I - b k k^T) D, and |I - b k k^T|_2 = 1, so
//     |D_end|_F <= exp(sum_t log g_t) |D_0|_F,   log g_t = softplus(alpha_t + dt_h) * ssm_a_h  (<= 0)
// per value head.  That bounds the direct term of keeping S_ABC instead of S_AB'C; the indirect term (C's own k, v,
// gate inputs computed under B) has no closed form and is what the splice test measures.
inline double log_gate(float alpha, float dt, float ssm_a) {
    const double v = (double) alpha + dt;
    const double sp = v > 20.0 ? v : std::log1p(std::exp(v));
    return sp * (double) ssm_a;
}

}  // namespace scr