#include "sentence_model.h"

#include "json.h"
#if defined(MSIME_NEURAL_AVX2)
#include "int8_linear_avx2.h"
#if defined(_MSC_VER)
#include <intrin.h>
#endif
#endif

#include <algorithm>
#include <cmath>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <limits>
#include <stdexcept>
#include <utility>

namespace neural
{
namespace
{

// The exact GELU, matching torch.nn.functional.gelu with its default settings. The tanh
// approximation is a different function and the model was not trained with it.
float erf_approx(float x)
{
    // Abramowitz and Stegun 7.1.26; error stays below 1.5e-7 across the range.
    float sign = x < 0.0f ? -1.0f : 1.0f;
    x = std::fabs(x);
    float t = 1.0f / (1.0f + 0.327591100f * x);
    float poly = t * (0.254829600f + t * (-0.284496740f + t * (1.421413700f + t * (-1.453152000f + t * 1.061405400f))));
    return sign * (1.0f - poly * std::exp(-x * x));
}

float gelu(float x)
{
    constexpr float kInvSqrt2 = 0.70710678118654752440f;
    return 0.5f * x * (1.0f + erf_approx(x * kInvSqrt2));
}

// Four independent accumulators, not one.
//
// This is where essentially all of the model's time goes, so the shape of the loop matters. A single
// running sum makes every iteration wait on the previous add: the FP unit has a ~4-cycle add latency
// but can start one per cycle, so a serial chain leaves it idle three cycles out of four, and the
// compiler cannot fix it — reassociating floating-point addition changes the result, which /fp:precise
// forbids it from doing on its own. Splitting the sum into four lanes that only meet at the end is
// that same reassociation, written out in the source where it is allowed, and it also gives the
// vectorizer a shape it can fold into packed adds.
//
// The trade is that the result differs from a strict left-to-right sum in the last bits. That is fine
// here — these are log probabilities compared against each other, and the parity test against the
// per-character reference path has room to spare — but it does mean this is not bit-identical to a
// naive reference implementation.
float dot(const float *a, const float *b, std::size_t n)
{
    float s0 = 0.0f;
    float s1 = 0.0f;
    float s2 = 0.0f;
    float s3 = 0.0f;
    std::size_t i = 0;
    for (; i + 4 <= n; i += 4)
    {
        s0 += a[i] * b[i];
        s1 += a[i + 1] * b[i + 1];
        s2 += a[i + 2] * b[i + 2];
        s3 += a[i + 3] * b[i + 3];
    }
    float total = (s0 + s1) + (s2 + s3);
    for (; i < n; ++i)
    {
        total += a[i] * b[i];
    }
    return total;
}

float dot_i8(const float *a, const std::int8_t *b, std::size_t n)
{
    float total = 0.0f;
    for (std::size_t i = 0; i < n; ++i)
    {
        total += a[i] * static_cast<float>(b[i]);
    }
    return total;
}

// out[m][o] = dot(x[m], w[o]) + bias[o], both operands row-major.
//
// Output-major on purpose: the weights are far larger than cache (the desktop model's are ~25 MB
// total), so with the rows on the outside every input row restreams the entire matrix from memory
// and the whole forward pass runs at DRAM bandwidth. Hoisting the output loop reads each weight row
// once and reuses it across the batch, which is the reason a batch of candidates costs barely more
// than one.
void linear_into(const std::vector<float> &x, const Matrix &w, const std::vector<float> &bias, std::size_t inputs,
                 std::size_t outputs, std::vector<float> &out, std::vector<float> &scratch)
{
    std::size_t rows = inputs == 0 ? 0 : x.size() / inputs;
    out.resize(rows * outputs);
#if defined(MSIME_NEURAL_AVX2)
    static const bool use_avx2 = detail::supports_avx2_fma();
    if (w.quantized && use_avx2 && rows != 0 && inputs != 0)
    {
        detail::linear_int8_avx2(x.data(), w.ints.data(), w.scales.data(), bias.empty() ? nullptr : bias.data(), rows,
                                 inputs, outputs, out.data());
        return;
    }
#endif
    if (w.quantized && scratch.size() < inputs)
    {
        scratch.resize(inputs);
    }
    for (std::size_t o = 0; o < outputs; ++o)
    {
        // Widen the row's int8 once per output row instead of once per (output row, input row).
        // The conversion is what keeps the inner product from vectorizing, and across a batch it
        // would otherwise be repeated for every row of the batch.
        //
        // The per-row scale still multiplies the finished dot product rather than the weights, the
        // same way Matrix::dot_row does it: widening an int8 is exact, but folding the scale into
        // each weight beforehand would round every element.
        const float *weights;
        float scale = 1.0f;
        if (w.quantized)
        {
            const std::int8_t *source = &w.ints[o * inputs];
            for (std::size_t i = 0; i < inputs; ++i)
            {
                scratch[i] = static_cast<float>(source[i]);
            }
            weights = scratch.data();
            scale = w.scales[o];
        }
        else
        {
            weights = &w.floats[o * inputs];
        }

        const float b = bias.empty() ? 0.0f : bias[o];
        for (std::size_t row = 0; row < rows; ++row)
        {
            out[row * outputs + o] = dot(&x[row * inputs], weights, inputs) * scale + b;
        }
    }
}

void layer_norm_into(const float *x, std::size_t rows, const std::vector<float> &weight, const std::vector<float> &bias,
                     std::size_t width, std::vector<float> &out)
{
    out.resize(rows * width);
    for (std::size_t row = 0; row < rows; ++row)
    {
        const float *source = &x[row * width];
        float *target = &out[row * width];
        float mean = 0.0f;
        for (std::size_t i = 0; i < width; ++i)
        {
            mean += source[i];
        }
        mean /= static_cast<float>(width);
        float variance = 0.0f;
        for (std::size_t i = 0; i < width; ++i)
        {
            float d = source[i] - mean;
            variance += d * d;
        }
        variance /= static_cast<float>(width);
        float inverse = 1.0f / std::sqrt(variance + 1e-5f);
        for (std::size_t i = 0; i < width; ++i)
        {
            target[i] = (source[i] - mean) * inverse * weight[i] + bias[i];
        }
    }
}

void split_qkv(const std::vector<float> &qkv, std::size_t rows, std::size_t n_embd, std::vector<float> &q,
               std::vector<float> &k, std::vector<float> &v)
{
    q.resize(rows * n_embd);
    k.resize(rows * n_embd);
    v.resize(rows * n_embd);
    for (std::size_t row = 0; row < rows; ++row)
    {
        const float *source = &qkv[row * 3 * n_embd];
        std::memcpy(&q[row * n_embd], source, n_embd * sizeof(float));
        std::memcpy(&k[row * n_embd], source + n_embd, n_embd * sizeof(float));
        std::memcpy(&v[row * n_embd], source + 2 * n_embd, n_embd * sizeof(float));
    }
}

// IEEE 754 binary16 to binary32, with subnormal handling rather than flush-to-zero.
float half_to_f32(std::uint16_t bits)
{
    std::uint32_t sign = static_cast<std::uint32_t>(bits & 0x8000u) << 16;
    std::uint32_t exponent = static_cast<std::uint32_t>((bits >> 10) & 0x1fu);
    std::uint32_t mantissa = static_cast<std::uint32_t>(bits & 0x03ffu);
    std::uint32_t out;
    if (exponent == 0)
    {
        if (mantissa == 0)
        {
            out = sign;
        }
        else
        {
            unsigned leading = 0;
            std::uint32_t m = mantissa;
            // Count leading zeros of a 32-bit value.
            while ((m & 0x80000000u) == 0)
            {
                m <<= 1;
                ++leading;
            }
            std::uint32_t exp = 134u - leading;
            std::uint32_t mant = (mantissa << (leading - 8)) & 0x007fffffu;
            out = sign | (exp << 23) | mant;
        }
    }
    else if (exponent == 0x1f)
    {
        out = sign | 0x7f800000u | (mantissa << 13);
    }
    else
    {
        out = sign | ((exponent + 127u - 15u) << 23) | (mantissa << 13);
    }
    float result;
    std::memcpy(&result, &out, sizeof(result));
    return result;
}

// Decodes UTF-8 text into Unicode code points. Malformed bytes map to U+FFFD.
std::vector<char32_t> decode_utf8(const std::string &text)
{
    std::vector<char32_t> out;
    std::size_t i = 0;
    const std::size_t n = text.size();
    while (i < n)
    {
        unsigned char c = static_cast<unsigned char>(text[i]);
        char32_t code;
        std::size_t extra;
        if (c < 0x80)
        {
            code = c;
            extra = 0;
        }
        else if ((c & 0xE0) == 0xC0)
        {
            code = c & 0x1F;
            extra = 1;
        }
        else if ((c & 0xF0) == 0xE0)
        {
            code = c & 0x0F;
            extra = 2;
        }
        else if ((c & 0xF8) == 0xF0)
        {
            code = c & 0x07;
            extra = 3;
        }
        else
        {
            out.push_back(0xFFFD);
            ++i;
            continue;
        }
        if (i + extra >= n)
        {
            out.push_back(0xFFFD);
            break;
        }
        bool ok = true;
        for (std::size_t j = 1; j <= extra; ++j)
        {
            unsigned char cc = static_cast<unsigned char>(text[i + j]);
            if ((cc & 0xC0) != 0x80)
            {
                ok = false;
                break;
            }
            code = (code << 6) | (cc & 0x3F);
        }
        if (!ok)
        {
            out.push_back(0xFFFD);
            ++i;
            continue;
        }
        out.push_back(code);
        i += extra + 1;
    }
    return out;
}

// A tensor descriptor from the safetensors header.
struct Entry
{
    std::string dtype;
    std::vector<std::size_t> shape;
    std::size_t off0 = 0;
    std::size_t off1 = 0;
};

std::vector<float> read_f32(const std::uint8_t *data, std::size_t data_len, const Entry &entry)
{
    if (entry.off1 > data_len || entry.off0 > entry.off1)
    {
        throw std::runtime_error("tensor offsets out of range");
    }
    std::size_t bytes = entry.off1 - entry.off0;
    std::size_t count = bytes / 4;
    std::vector<float> out(count);
    const std::uint8_t *p = data + entry.off0;
    for (std::size_t i = 0; i < count; ++i)
    {
        std::memcpy(&out[i], p + i * 4, 4);
    }
    return out;
}

std::vector<float> read_f16(const std::uint8_t *data, std::size_t data_len, const Entry &entry)
{
    if (entry.off1 > data_len || entry.off0 > entry.off1)
    {
        throw std::runtime_error("tensor offsets out of range");
    }
    std::size_t bytes = entry.off1 - entry.off0;
    std::size_t count = bytes / 2;
    std::vector<float> out(count);
    const std::uint8_t *p = data + entry.off0;
    for (std::size_t i = 0; i < count; ++i)
    {
        std::uint16_t bits;
        std::memcpy(&bits, p + i * 2, 2);
        out[i] = half_to_f32(bits);
    }
    return out;
}

std::size_t product(const std::vector<std::size_t> &shape)
{
    std::size_t total = 1;
    for (std::size_t d : shape)
    {
        total *= d;
    }
    return total;
}

} // namespace

#if defined(MSIME_NEURAL_AVX2)
namespace detail
{
bool supports_avx2_fma()
{
#if defined(_MSC_VER)
    int info[4] = {};
    __cpuid(info, 0);
    if (info[0] < 7)
    {
        return false;
    }
    __cpuid(info, 1);
    constexpr int kFma = 1 << 12;
    constexpr int kOsxsave = 1 << 27;
    constexpr int kAvx = 1 << 28;
    if ((info[2] & (kFma | kOsxsave | kAvx)) != (kFma | kOsxsave | kAvx) || (_xgetbv(0) & 6) != 6)
    {
        return false;
    }
    __cpuidex(info, 7, 0);
    return (info[1] & (1 << 5)) != 0;
#else
    return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma");
#endif
}
} // namespace detail
#endif

struct SentenceModel::Workspace
{
    std::vector<float> normed;
    std::vector<float> qkv;
    std::vector<float> q;
    std::vector<float> k;
    std::vector<float> v;
    std::vector<float> attended;
    std::vector<float> projected;
    std::vector<float> ff_hidden;
    std::vector<float> ff_out;
    std::vector<float> linear_scratch;
    std::vector<float> attention_scores;
};

float Matrix::dot_row(const float *x, std::size_t row, std::size_t width) const
{
    if (quantized)
    {
        return dot_i8(x, &ints[row * width], width) * scales[row];
    }
    return dot(x, &floats[row * width], width);
}

void Matrix::row_into(std::size_t index, std::size_t width, float *out) const
{
    if (quantized)
    {
        float scale = scales[index];
        const std::int8_t *p = &ints[index * width];
        for (std::size_t i = 0; i < width; ++i)
        {
            out[i] = static_cast<float>(p[i]) * scale;
        }
    }
    else
    {
        std::memcpy(out, &floats[index * width], width * sizeof(float));
    }
}

std::unique_ptr<SentenceModel> SentenceModel::load(const std::vector<std::uint8_t> &bytes)
{
    try
    {
        if (bytes.size() < 8)
        {
            return nullptr;
        }
        std::uint64_t header_len = 0;
        std::memcpy(&header_len, bytes.data(), 8);
        std::size_t header_end = 8 + static_cast<std::size_t>(header_len);
        if (header_len == 0 || header_end > bytes.size())
        {
            return nullptr;
        }
        std::string header_text(reinterpret_cast<const char *>(bytes.data() + 8), static_cast<std::size_t>(header_len));
        json::Value header = json::parse(header_text);
        if (!header.is_object())
        {
            return nullptr;
        }
        const std::uint8_t *data = bytes.data() + header_end;
        std::size_t data_len = bytes.size() - header_end;

        // Metadata.
        const json::Value *metadata = header.find("__metadata__");
        if (metadata == nullptr || !metadata->is_object())
        {
            return nullptr;
        }
        auto metadata_string = [&](const std::string &key) -> const std::string * {
            const json::Value *value = metadata->find(key);
            if (value == nullptr || !value->is_string())
            {
                return nullptr;
            }
            return &value->string;
        };
        const std::string *config_text = metadata_string("config");
        const std::string *vocab_text = metadata_string("vocab");
        if (config_text == nullptr || vocab_text == nullptr)
        {
            return nullptr;
        }

        auto model = std::unique_ptr<SentenceModel>(new SentenceModel());

        json::Value config_json = json::parse(*config_text);
        auto config_uint = [&](const std::string &key) -> std::size_t {
            const json::Value *value = config_json.find(key);
            if (value == nullptr || !value->is_number())
            {
                throw std::runtime_error("config missing " + key);
            }
            return static_cast<std::size_t>(value->number);
        };
        model->config_.vocab = config_uint("vocab");
        model->config_.n_layer = config_uint("n_layer");
        model->config_.n_head = config_uint("n_head");
        model->config_.n_embd = config_uint("n_embd");
        model->config_.context = config_uint("context");
        const Config &config = model->config_;
        if (config.n_head == 0 || config.n_embd == 0 || config.n_embd % config.n_head != 0 || config.vocab == 0 ||
            config.n_layer == 0 || config.context == 0)
        {
            return nullptr;
        }

        json::Value vocab_json = json::parse(*vocab_text);
        if (!vocab_json.is_array() || vocab_json.array.size() != config.vocab)
        {
            return nullptr;
        }
        for (std::size_t id = 0; id < vocab_json.array.size(); ++id)
        {
            const json::Value &item = vocab_json.array[id];
            if (!item.is_string())
            {
                continue;
            }
            std::vector<char32_t> chars = decode_utf8(item.string);
            // Reserved tokens (<pad> etc.) are addressed by constant, not by text; single code
            // points are the character vocabulary.
            if (chars.size() == 1)
            {
                model->index_[chars[0]] = static_cast<std::uint32_t>(id);
            }
        }

        if (const std::string *attribution = metadata_string("attribution"))
        {
            model->attribution_ = *attribution;
        }
        if (const std::string *license = metadata_string("license"))
        {
            model->license_ = *license;
        }

        // Collect every tensor descriptor (including the .scale companions).
        std::unordered_map<std::string, Entry> entries;
        for (const auto &field : header.object)
        {
            if (field.first == "__metadata__")
            {
                continue;
            }
            const json::Value &desc = field.second;
            if (!desc.is_object())
            {
                continue;
            }
            Entry entry;
            const json::Value *dtype = desc.find("dtype");
            const json::Value *shape = desc.find("shape");
            const json::Value *offsets = desc.find("data_offsets");
            if (dtype == nullptr || !dtype->is_string() || offsets == nullptr || !offsets->is_array() ||
                offsets->array.size() != 2)
            {
                return nullptr;
            }
            entry.dtype = dtype->string;
            if (shape != nullptr && shape->is_array())
            {
                for (const json::Value &dim : shape->array)
                {
                    entry.shape.push_back(static_cast<std::size_t>(dim.number));
                }
            }
            entry.off0 = static_cast<std::size_t>(offsets->array[0].number);
            entry.off1 = static_cast<std::size_t>(offsets->array[1].number);
            entries.emplace(field.first, std::move(entry));
        }

        auto materialize = [&](const std::string &name) -> Matrix {
            auto it = entries.find(name);
            if (it == entries.end())
            {
                throw std::runtime_error("missing tensor " + name);
            }
            const Entry &entry = it->second;
            Matrix matrix;
            std::size_t expected = product(entry.shape);
            if (entry.dtype == "F32")
            {
                matrix.floats = read_f32(data, data_len, entry);
                matrix.quantized = false;
                if (matrix.floats.size() != expected)
                {
                    throw std::runtime_error(name + ": shape does not match data");
                }
            }
            else if (entry.dtype == "F16")
            {
                matrix.floats = read_f16(data, data_len, entry);
                matrix.quantized = false;
                if (matrix.floats.size() != expected)
                {
                    throw std::runtime_error(name + ": shape does not match data");
                }
            }
            else if (entry.dtype == "I8")
            {
                auto scale_it = entries.find(name + ".scale");
                if (scale_it == entries.end())
                {
                    throw std::runtime_error(name + ": missing scale");
                }
                matrix.scales = read_f32(data, data_len, scale_it->second);
                std::size_t rows = entry.shape.empty() ? 1 : entry.shape.front();
                if (rows == 0 || matrix.scales.size() != rows)
                {
                    throw std::runtime_error(name + ": scale length does not match rows");
                }
                if (entry.off1 > data_len || entry.off0 > entry.off1)
                {
                    throw std::runtime_error(name + ": offsets out of range");
                }
                std::size_t count = entry.off1 - entry.off0;
                matrix.ints.resize(count);
                std::memcpy(matrix.ints.data(), data + entry.off0, count);
                matrix.quantized = true;
                if (matrix.ints.size() != expected)
                {
                    throw std::runtime_error(name + ": shape does not match data");
                }
            }
            else
            {
                throw std::runtime_error(name + ": unsupported dtype " + entry.dtype);
            }
            return matrix;
        };

        auto materialize_floats = [&](const std::string &name) -> std::vector<float> {
            Matrix matrix = materialize(name);
            if (matrix.quantized)
            {
                throw std::runtime_error(name + " is quantized but must be float");
            }
            return std::move(matrix.floats);
        };

        model->token_ = materialize("tok.weight");
        model->position_ = materialize_floats("pos.weight");
        model->blocks_.reserve(config.n_layer);
        for (std::size_t layer = 0; layer < config.n_layer; ++layer)
        {
            std::string prefix = "blocks." + std::to_string(layer) + ".";
            Block block;
            block.ln1_weight = materialize_floats(prefix + "ln1.weight");
            block.ln1_bias = materialize_floats(prefix + "ln1.bias");
            block.qkv_weight = materialize(prefix + "qkv.weight");
            block.qkv_bias = materialize_floats(prefix + "qkv.bias");
            block.proj_weight = materialize(prefix + "proj.weight");
            block.proj_bias = materialize_floats(prefix + "proj.bias");
            block.ln2_weight = materialize_floats(prefix + "ln2.weight");
            block.ln2_bias = materialize_floats(prefix + "ln2.bias");
            block.fc_weight = materialize(prefix + "fc.weight");
            block.fc_bias = materialize_floats(prefix + "fc.bias");
            block.out_weight = materialize(prefix + "out.weight");
            block.out_bias = materialize_floats(prefix + "out.bias");
            model->blocks_.push_back(std::move(block));
        }
        model->final_weight_ = materialize_floats("ln_f.weight");
        model->final_bias_ = materialize_floats("ln_f.bias");

        // Geometry checks: any mismatch means a malformed file, so refuse it.
        const std::size_t n_embd = config.n_embd;
        if (model->token_.len() != config.vocab * n_embd || model->position_.size() != config.context * n_embd ||
            model->final_weight_.size() != n_embd || model->final_bias_.size() != n_embd)
        {
            return nullptr;
        }
        for (const Block &block : model->blocks_)
        {
            if (block.qkv_weight.len() != 3 * n_embd * n_embd || block.proj_weight.len() != n_embd * n_embd ||
                block.fc_weight.len() != 4 * n_embd * n_embd || block.out_weight.len() != 4 * n_embd * n_embd)
            {
                return nullptr;
            }
        }
        return model;
    }
    catch (const std::exception &)
    {
        return nullptr;
    }
}

std::unique_ptr<SentenceModel> SentenceModel::load_file(const std::string &path)
{
    // path is UTF-8. A narrow std::string opens through the ANSI code page on Windows, which
    // cannot name a data directory under a non-ASCII user profile.
#if defined(__cpp_lib_char8_t)
    const std::filesystem::path file_path(std::u8string(path.begin(), path.end()));
#else
    const std::filesystem::path file_path = std::filesystem::u8path(path);
#endif
    std::ifstream file(file_path, std::ios::binary);
    if (!file)
    {
        return nullptr;
    }
    file.seekg(0, std::ios::end);
    std::streamoff size = file.tellg();
    if (size <= 0)
    {
        return nullptr;
    }
    file.seekg(0, std::ios::beg);
    std::vector<std::uint8_t> bytes(static_cast<std::size_t>(size));
    if (!file.read(reinterpret_cast<char *>(bytes.data()), size))
    {
        return nullptr;
    }
    return load(bytes);
}

std::optional<std::uint32_t> SentenceModel::token(char32_t character) const
{
    auto it = index_.find(character);
    if (it == index_.end())
    {
        return std::nullopt;
    }
    return it->second;
}

std::vector<std::uint32_t> SentenceModel::encode(const std::string &text) const
{
    std::vector<char32_t> chars = decode_utf8(text);
    std::vector<std::uint32_t> out;
    out.reserve(chars.size());
    for (char32_t ch : chars)
    {
        auto it = index_.find(ch);
        out.push_back(it == index_.end() ? kUnk : it->second);
    }
    return out;
}

std::vector<float> SentenceModel::embed(const std::vector<std::uint32_t> &tokens, std::size_t width,
                                        std::size_t offset) const
{
    const std::size_t n_embd = config_.n_embd;
    std::vector<float> x(tokens.size() * n_embd, 0.0f);
    std::vector<float> token_row(n_embd, 0.0f);
    for (std::size_t step = 0; step < tokens.size(); ++step)
    {
        std::size_t token = std::min<std::size_t>(tokens[step], config_.vocab - 1);
        // Position restarts at `offset` on every batch row: each row is a separate continuation of
        // the same cached prefix, not a continuation of the row before it.
        std::size_t within = width == 0 ? step : step % width;
        std::size_t position = std::min<std::size_t>(offset + within, config_.context - 1);
        token_.row_into(token, n_embd, token_row.data());
        const float *position_row = &position_[position * n_embd];
        float *row = &x[step * n_embd];
        for (std::size_t i = 0; i < n_embd; ++i)
        {
            row[i] = token_row[i] + position_row[i];
        }
    }
    return x;
}

void SentenceModel::attention(const std::vector<float> &q, const std::vector<float> &k, const std::vector<float> &v,
                              std::size_t batch, std::size_t width, const std::vector<float> *past_k,
                              const std::vector<float> *past_v, std::size_t past_len, std::vector<float> &out,
                              std::vector<float> &scores) const
{
    const std::size_t n_embd = config_.n_embd;
    const std::size_t heads = config_.n_head;
    const std::size_t head_width = n_embd / heads;
    const float scale = 1.0f / std::sqrt(static_cast<float>(head_width));
    if (past_k == nullptr || past_v == nullptr)
    {
        past_len = 0;
    }
    out.assign(batch * width * n_embd, 0.0f);
    scores.resize(past_len + width);
    for (std::size_t row = 0; row < batch; ++row)
    {
        const std::size_t base_row = row * width;
        for (std::size_t step = 0; step < width; ++step)
        {
            // Causal, in two runs: the whole shared prefix, then this row's own tokens up to and
            // including `step`. Rows never see each other.
            const std::size_t visible_self = step + 1;
            const std::size_t visible = past_len + visible_self;
            for (std::size_t head = 0; head < heads; ++head)
            {
                const std::size_t base = head * head_width;
                const float *query = &q[(base_row + step) * n_embd + base];
                float highest = -std::numeric_limits<float>::infinity();
                for (std::size_t j = 0; j < past_len; ++j)
                {
                    float s = dot(query, &(*past_k)[j * n_embd + base], head_width) * scale;
                    scores[j] = s;
                    highest = std::max(highest, s);
                }
                for (std::size_t j = 0; j < visible_self; ++j)
                {
                    float s = dot(query, &k[(base_row + j) * n_embd + base], head_width) * scale;
                    scores[past_len + j] = s;
                    highest = std::max(highest, s);
                }
                float sum = 0.0f;
                for (std::size_t j = 0; j < visible; ++j)
                {
                    scores[j] = std::exp(scores[j] - highest);
                    sum += scores[j];
                }
                float *target = &out[(base_row + step) * n_embd + base];
                for (std::size_t j = 0; j < past_len; ++j)
                {
                    float weight = scores[j] / sum;
                    const float *value = &(*past_v)[j * n_embd + base];
                    for (std::size_t i = 0; i < head_width; ++i)
                    {
                        target[i] += weight * value[i];
                    }
                }
                for (std::size_t j = 0; j < visible_self; ++j)
                {
                    float weight = scores[past_len + j] / sum;
                    const float *value = &v[(base_row + j) * n_embd + base];
                    for (std::size_t i = 0; i < head_width; ++i)
                    {
                        target[i] += weight * value[i];
                    }
                }
            }
        }
    }
}

void SentenceModel::feed_forward(std::vector<float> &x, const Block &block, Workspace &workspace) const
{
    const std::size_t n_embd = config_.n_embd;
    std::size_t rows = x.size() / n_embd;
    layer_norm_into(x.data(), rows, block.ln2_weight, block.ln2_bias, n_embd, workspace.normed);
    linear_into(workspace.normed, block.fc_weight, block.fc_bias, n_embd, 4 * n_embd, workspace.ff_hidden,
                workspace.linear_scratch);
    for (float &value : workspace.ff_hidden)
    {
        value = gelu(value);
    }
    linear_into(workspace.ff_hidden, block.out_weight, block.out_bias, 4 * n_embd, n_embd, workspace.ff_out,
                workspace.linear_scratch);
    for (std::size_t i = 0; i < x.size(); ++i)
    {
        x[i] += workspace.ff_out[i];
    }
}

std::vector<float> SentenceModel::forward(const std::vector<std::uint32_t> &tokens, std::size_t batch,
                                          std::size_t width, const PrefixCache *past, PrefixCache *record) const
{
    const std::size_t n_embd = config_.n_embd;
    const std::size_t rows = batch * width;
    const std::size_t past_len = past == nullptr ? 0 : past->length;
    std::vector<float> x = embed(tokens, width, past_len);
    Workspace workspace;
    if (record != nullptr)
    {
        // A cache is shared by every row that later attends to it, so it can only ever describe one
        // sequence. build_prefix is the only caller and always passes batch == 1 from position 0.
        record->keys.assign(blocks_.size(), {});
        record->values.assign(blocks_.size(), {});
        record->length = past_len + width;
    }
    for (std::size_t layer = 0; layer < blocks_.size(); ++layer)
    {
        const Block &block = blocks_[layer];
        layer_norm_into(x.data(), rows, block.ln1_weight, block.ln1_bias, n_embd, workspace.normed);
        linear_into(workspace.normed, block.qkv_weight, block.qkv_bias, n_embd, 3 * n_embd, workspace.qkv,
                    workspace.linear_scratch);
        split_qkv(workspace.qkv, rows, n_embd, workspace.q, workspace.k, workspace.v);
        const std::vector<float> *past_k = past_len == 0 ? nullptr : &past->keys[layer];
        const std::vector<float> *past_v = past_len == 0 ? nullptr : &past->values[layer];
        attention(workspace.q, workspace.k, workspace.v, batch, width, past_k, past_v, past_len, workspace.attended,
                  workspace.attention_scores);
        linear_into(workspace.attended, block.proj_weight, block.proj_bias, n_embd, n_embd, workspace.projected,
                    workspace.linear_scratch);
        for (std::size_t i = 0; i < x.size(); ++i)
        {
            x[i] += workspace.projected[i];
        }
        feed_forward(x, block, workspace);
        if (record != nullptr)
        {
            record->keys[layer] = workspace.k;
            record->values[layer] = workspace.v;
        }
    }
    layer_norm_into(x.data(), rows, final_weight_, final_bias_, n_embd, workspace.normed);
    return std::move(workspace.normed);
}

PrefixCache SentenceModel::build_prefix(const std::vector<std::uint32_t> &tokens) const
{
    PrefixCache cache;
    if (tokens.empty())
    {
        return cache;
    }
    forward(tokens, 1, tokens.size(), nullptr, &cache);
    return cache;
}

std::vector<float> SentenceModel::distribution(const std::vector<float> &hidden) const
{
    const std::size_t n_embd = config_.n_embd;
    const std::size_t vocab = config_.vocab;
    std::vector<float> logits(vocab, 0.0f);
    float highest = -std::numeric_limits<float>::infinity();
    for (std::size_t row = 0; row < vocab; ++row)
    {
        float logit = token_.dot_row(hidden.data(), row, n_embd);
        logits[row] = logit;
        highest = std::max(highest, logit);
    }
    float sum = 0.0f;
    for (float logit : logits)
    {
        sum += std::exp(logit - highest);
    }
    float offset = highest + std::log(sum);
    for (float &logit : logits)
    {
        logit -= offset;
    }
    return logits;
}

std::vector<float> SentenceModel::target_log_probs(const std::vector<float> &hidden, std::size_t rows,
                                                   const std::vector<std::uint32_t> &targets) const
{
    const std::size_t n_embd = config_.n_embd;
    const std::size_t vocab = config_.vocab;
    std::vector<std::size_t> active_rows;
    active_rows.reserve(rows);
    for (std::size_t row = 0; row < rows; ++row)
    {
        if (targets[row] != kNoTarget)
        {
            active_rows.push_back(row);
        }
    }
    std::vector<float> out(rows, 0.0f);
    if (active_rows.empty())
    {
        return out;
    }

    std::vector<float> logits(active_rows.size() * vocab, 0.0f);
#if defined(MSIME_NEURAL_AVX2)
    static const bool use_avx2 = detail::supports_avx2_fma();
    if (token_.quantized && use_avx2)
    {
        std::vector<float> active_hidden(active_rows.size() * n_embd);
        for (std::size_t slot = 0; slot < active_rows.size(); ++slot)
        {
            std::memcpy(&active_hidden[slot * n_embd], &hidden[active_rows[slot] * n_embd], n_embd * sizeof(float));
        }
        detail::linear_int8_avx2(active_hidden.data(), token_.ints.data(), token_.scales.data(), nullptr,
                                 active_rows.size(), n_embd, vocab, logits.data());
    }
    else
#endif
    {
        // The scalar path widens one vocabulary row once for all active rows.
        std::vector<float> scratch(token_.quantized ? n_embd : 0, 0.0f);
        for (std::size_t id = 0; id < vocab; ++id)
        {
            const float *weights;
            float scale = 1.0f;
            if (token_.quantized)
            {
                const std::int8_t *source = &token_.ints[id * n_embd];
                for (std::size_t i = 0; i < n_embd; ++i)
                {
                    scratch[i] = static_cast<float>(source[i]);
                }
                weights = scratch.data();
                scale = token_.scales[id];
            }
            else
            {
                weights = &token_.floats[id * n_embd];
            }
            for (std::size_t slot = 0; slot < active_rows.size(); ++slot)
            {
                const std::size_t row = active_rows[slot];
                logits[slot * vocab + id] = dot(&hidden[row * n_embd], weights, n_embd) * scale;
            }
        }
    }

    for (std::size_t slot = 0; slot < active_rows.size(); ++slot)
    {
        const std::size_t row = active_rows[slot];
        const float *source = &logits[slot * vocab];
        float highest = -std::numeric_limits<float>::infinity();
        for (std::size_t id = 0; id < vocab; ++id)
        {
            highest = std::max(highest, source[id]);
        }
        float sum = 0.0f;
        for (std::size_t id = 0; id < vocab; ++id)
        {
            sum += std::exp(source[id] - highest);
        }
        out[row] = source[targets[row]] - (highest + std::log(sum));
    }
    return out;
}

std::vector<float> SentenceModel::next_log_probabilities(const std::string &text) const
{
    std::vector<std::uint32_t> tokens;
    tokens.push_back(kBos);
    std::vector<std::uint32_t> encoded = encode(text);
    std::size_t room = config_.context - 1;
    std::size_t start = encoded.size() > room ? encoded.size() - room : 0;
    tokens.insert(tokens.end(), encoded.begin() + static_cast<std::ptrdiff_t>(start), encoded.end());
    std::vector<float> hidden = forward(tokens, 1, tokens.size(), nullptr, nullptr);
    const std::size_t n_embd = config_.n_embd;
    std::vector<float> last(hidden.end() - static_cast<std::ptrdiff_t>(n_embd), hidden.end());
    return distribution(last);
}

std::vector<double> SentenceModel::score_sentences(const std::string &context,
                                                   const std::vector<std::string> &texts) const
{
    if (texts.empty())
    {
        return {};
    }
    const std::size_t limit = config_.context;
    if (limit < 2)
    {
        return std::vector<double>(texts.size(), 0.0);
    }

    // The sequence each candidate is scored under is <bos> + context + candidate. Every candidate
    // shares all of that but its own tail, so the shared part is split at its last token: that token
    // leads each batch row, because the distribution it produces is what scores the candidate's
    // first character. Only what precedes it goes in the cache, since only that stays fixed.
    std::vector<std::uint32_t> full;
    full.push_back(kBos);
    for (std::uint32_t id : encode(context))
    {
        full.push_back(id);
    }

    std::vector<std::vector<std::uint32_t>> tails;
    tails.reserve(texts.size());
    std::size_t longest = 0;
    for (const std::string &text : texts)
    {
        std::vector<std::uint32_t> ids = encode(text);
        // A candidate longer than the window minus the leading token keeps its tail, which is the
        // part the remaining context can still inform.
        if (ids.size() > limit - 1)
        {
            ids.erase(ids.begin(), ids.begin() + static_cast<std::ptrdiff_t>(ids.size() - (limit - 1)));
        }
        longest = std::max(longest, ids.size());
        tails.push_back(std::move(ids));
    }
    if (longest == 0)
    {
        return std::vector<double>(texts.size(), 0.0);
    }
    // The last candidate token is a target only. Its hidden state predicts no scored character.
    const std::size_t width = longest;

    const std::uint32_t last = full.back();
    std::vector<std::uint32_t> head(full.begin(), full.end() - 1);
    // Keep the old context limit so removing the unused final token does not change scores.
    const std::size_t room = limit > longest + 1 ? limit - (longest + 1) : 0;
    if (head.size() > room)
    {
        head.erase(head.begin(), head.begin() + static_cast<std::ptrdiff_t>(head.size() - room));
    }

    const std::size_t batch = tails.size();
    // Short rows are padded with token 0; their positions get kNoTarget and are never read.
    std::vector<std::uint32_t> flat(batch * width, 0u);
    std::vector<std::uint32_t> targets(batch * width, kNoTarget);
    for (std::size_t row = 0; row < batch; ++row)
    {
        flat[row * width] = last;
        for (std::size_t i = 0; i < tails[row].size(); ++i)
        {
            // Position i predicts token i+1, so the leading token's slot carries the first character.
            targets[row * width + i] = tails[row][i];
            if (i + 1 < tails[row].size())
            {
                flat[row * width + 1 + i] = tails[row][i];
            }
        }
    }

    std::vector<float> scored;
    {
        std::lock_guard<std::mutex> lock(cache_mutex_);
        if (cached_tokens_ != head)
        {
            cached_prefix_ = build_prefix(head);
            cached_tokens_ = head;
        }
        std::vector<float> hidden = forward(flat, batch, width, &cached_prefix_, nullptr);
        scored = target_log_probs(hidden, batch * width, targets);
    }

    std::vector<double> out(batch, 0.0);
    for (std::size_t row = 0; row < batch; ++row)
    {
        double total = 0.0;
        for (std::size_t i = 0; i < tails[row].size(); ++i)
        {
            total += static_cast<double>(scored[row * width + i]);
        }
        out[row] = total;
    }
    return out;
}

} // namespace neural