// STEP 10 native resumable packed-Q1 LoRA trainer acceptance executable.
//
// The accepted Stage 7 loader/full-model graph is used only as a bootstrap.
// Its optimizer epoch is intercepted. GGML computes the real masked backward;
// deterministic host AdamW updates only the six LoRA A/B tensors.

#include "llama.h"
#include "llama-adapter.h"
#include "llama-model.h"

#include "ggml.h"
#include "ggml-backend.h"

#include <algorithm>
#include <array>
#include <cerrno>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <regex>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

namespace fs = std::filesystem;

struct target_spec {
    const char * name;
    const char * category;
    int block;
    int64_t K;
    int64_t M;
    int64_t rank;
    float alpha;
};

static const std::array<target_spec, 3> g_specs = {{
    {"blk.0.ssm_alpha.weight", "ssm",       0,  5120,    48, 4, 8.0f},
    {"blk.11.attn_k.weight",   "attention", 11, 5120,  1024, 4, 8.0f},
    {"blk.0.ffn_down.weight",  "ffn",       0, 17408,  5120, 4, 8.0f},
}};

struct trainer_config {
    std::string model_path;
    std::string adapter_path;
    std::string train_dir;
    std::string validation_dir;
    std::string output_dir;
    std::string mode;
    std::string resume_path;
    std::string model_sha256;
    std::string dataset_fingerprint;
    uint64_t seed = 1234;
    uint32_t context_tokens = 64;
    uint64_t target_global_steps = 2;
    uint32_t accumulation_steps = 2;
    float base_lr = 2.0e-4f;
    float min_lr = 2.0e-5f;
    uint32_t warmup_steps = 1;
    float beta1 = 0.9f;
    float beta2 = 0.999f;
    float eps = 1.0e-8f;
    float weight_decay = 0.0f;
    float max_grad_norm = 1.0e-3f;
    uint64_t stop_after_micro = 0;
};

struct sample_record {
    std::vector<llama_token> tokens;
    std::vector<uint8_t> target_mask;
    int32_t source_index = -1;
};

struct param_state {
    std::string name;
    ggml_tensor * tensor = nullptr;
    std::vector<float> value;
    std::vector<float> initial;
    std::vector<float> m;
    std::vector<float> v;
    std::vector<float> grad_sum;
};

struct trainer_state {
    uint64_t global_step = 0;
    uint32_t micro_step = 0;
    uint64_t total_micro = 0;
    uint64_t epoch = 0;
    uint64_t sample_cursor = 0;
    uint64_t rng_state = 0;
    uint64_t optimizer_updates = 0;
    uint64_t checkpoint_writes = 0;
    uint64_t clip_applied_count = 0;
    uint64_t early_update_violations = 0;
    std::vector<double> train_losses;
    std::vector<double> learning_rates;
    std::vector<double> grad_norm_pre;
    std::vector<double> grad_norm_post;
};

struct run_outcome {
    bool hook_seen = false;
    bool success = false;
    std::string error;
    std::string checkpoint_path;
    std::string checkpoint_sha256;
    std::string adapter_path;
    std::string adapter_sha256;
    std::string state_fingerprint;
    double validation_loss = NAN;
    bool validation_unchanged = false;
    double adapter_reload_max_diff = INFINITY;
    size_t base_changed_bytes = 0;
    size_t train_sample_count = 0;
    size_t validation_sample_count = 0;
};

static trainer_config g_cfg;
static trainer_state g_state;
static run_outcome g_outcome;
static llama_adapter_lora * g_adapter = nullptr;
static llama_model * g_model = nullptr;
static llama_context * g_context = nullptr;
static std::vector<param_state> g_params;
static std::vector<std::pair<std::string, std::vector<uint8_t>>> g_base_initial;

static std::string shell_quote(const std::string & input) {
    std::string output = "'";
    for (const char value : input) {
        if (value == '\'') {
            output += "'\"'\"'";
        } else {
            output += value;
        }
    }
    output += "'";
    return output;
}

static std::string sha256_file(const fs::path & path) {
    const std::string command = "sha256sum " + shell_quote(path.string());
    FILE * pipe = popen(command.c_str(), "r");
    if (!pipe) {
        throw std::runtime_error("could not run sha256sum");
    }
    std::string output;
    char buffer[512] = {};
    while (fgets(buffer, sizeof(buffer), pipe)) {
        output += buffer;
    }
    const int rc = pclose(pipe);
    if (rc != 0 || output.size() < 64) {
        throw std::runtime_error("sha256sum failed: " + path.string());
    }
    return output.substr(0, 64);
}

static uint64_t fnv1a64(
        const uint8_t * data,
        size_t size,
        uint64_t seed = 1469598103934665603ULL) {
    uint64_t result = seed;
    for (size_t i = 0; i < size; ++i) {
        result ^= data[i];
        result *= 1099511628211ULL;
    }
    return result;
}

static std::string hex64(uint64_t value) {
    std::ostringstream output;
    output << std::hex << std::setw(16) << std::setfill('0') << value;
    return output.str();
}

static std::vector<float> tensor_f32(const ggml_tensor * tensor) {
    if (!tensor || !tensor->buffer) {
        throw std::runtime_error("unallocated tensor");
    }
    const size_t count = ggml_nelements(tensor);
    std::vector<float> values(count);
    if (tensor->type == GGML_TYPE_F32) {
        ggml_backend_tensor_get(tensor, values.data(), 0, count*sizeof(float));
        return values;
    }
    if (tensor->type == GGML_TYPE_F16) {
        std::vector<ggml_fp16_t> temp(count);
        ggml_backend_tensor_get(tensor, temp.data(), 0, count*sizeof(ggml_fp16_t));
        for (size_t i = 0; i < count; ++i) {
            values[i] = ggml_fp16_to_fp32(temp[i]);
        }
        return values;
    }
    throw std::runtime_error("LoRA tensor is not F32/F16");
}

static void set_tensor_f32(ggml_tensor * tensor, const std::vector<float> & values) {
    if (!tensor || tensor->type != GGML_TYPE_F32 || values.size() != ggml_nelements(tensor)) {
        throw std::runtime_error("invalid F32 tensor update");
    }
    ggml_backend_tensor_set(tensor, values.data(), 0, values.size()*sizeof(float));
}

static std::vector<uint8_t> tensor_bytes(const ggml_tensor * tensor) {
    std::vector<uint8_t> bytes(ggml_nbytes(tensor));
    ggml_backend_tensor_get(tensor, bytes.data(), 0, bytes.size());
    return bytes;
}

static size_t changed_bytes(
        const std::vector<uint8_t> & left,
        const std::vector<uint8_t> & right) {
    if (left.size() != right.size()) {
        return std::max(left.size(), right.size());
    }
    size_t changed = 0;
    for (size_t i = 0; i < left.size(); ++i) {
        changed += left[i] != right[i];
    }
    return changed;
}

static uint64_t parameter_fingerprint() {
    uint64_t value = 1469598103934665603ULL;
    for (const param_state & state : g_params) {
        value = fnv1a64(
            reinterpret_cast<const uint8_t *>(state.name.data()),
            state.name.size(),
            value);
        value = fnv1a64(
            reinterpret_cast<const uint8_t *>(state.value.data()),
            state.value.size()*sizeof(float),
            value);
    }
    return value;
}

static std::string config_identity() {
    std::ostringstream output;
    output << "prism-step10-v1"
           << "|ctx=" << g_cfg.context_tokens
           << "|target_steps=" << g_cfg.target_global_steps
           << "|accum=" << g_cfg.accumulation_steps
           << "|base_lr=" << std::setprecision(9) << g_cfg.base_lr
           << "|min_lr=" << g_cfg.min_lr
           << "|warmup=" << g_cfg.warmup_steps
           << "|beta1=" << g_cfg.beta1
           << "|beta2=" << g_cfg.beta2
           << "|eps=" << g_cfg.eps
           << "|wd=" << g_cfg.weight_decay
           << "|clip=" << g_cfg.max_grad_norm
           << "|seed=" << g_cfg.seed;
    return output.str();
}

static void append_u32(std::vector<uint8_t> & out, uint32_t value) {
    for (int i = 0; i < 4; ++i) out.push_back((value >> (8*i)) & 0xff);
}
static void append_u64(std::vector<uint8_t> & out, uint64_t value) {
    for (int i = 0; i < 8; ++i) out.push_back((value >> (8*i)) & 0xff);
}
static void append_f32(std::vector<uint8_t> & out, float value) {
    uint32_t bits = 0;
    std::memcpy(&bits, &value, sizeof(bits));
    append_u32(out, bits);
}
static void append_f64(std::vector<uint8_t> & out, double value) {
    uint64_t bits = 0;
    std::memcpy(&bits, &value, sizeof(bits));
    append_u64(out, bits);
}
static void append_string(std::vector<uint8_t> & out, const std::string & value) {
    append_u32(out, static_cast<uint32_t>(value.size()));
    out.insert(out.end(), value.begin(), value.end());
}
static void append_float_vector(std::vector<uint8_t> & out, const std::vector<float> & values) {
    append_u64(out, values.size());
    const uint8_t * ptr = reinterpret_cast<const uint8_t *>(values.data());
    out.insert(out.end(), ptr, ptr + values.size()*sizeof(float));
}
static void append_double_vector(std::vector<uint8_t> & out, const std::vector<double> & values) {
    append_u64(out, values.size());
    for (double value : values) append_f64(out, value);
}

struct byte_reader {
    const std::vector<uint8_t> & data;
    size_t offset = 0;
    uint32_t u32() {
        if (offset + 4 > data.size()) throw std::runtime_error("checkpoint EOF u32");
        uint32_t value = 0;
        for (int i = 0; i < 4; ++i) value |= uint32_t(data[offset++]) << (8*i);
        return value;
    }
    uint64_t u64() {
        if (offset + 8 > data.size()) throw std::runtime_error("checkpoint EOF u64");
        uint64_t value = 0;
        for (int i = 0; i < 8; ++i) value |= uint64_t(data[offset++]) << (8*i);
        return value;
    }
    float f32() {
        uint32_t bits = u32(); float value; std::memcpy(&value, &bits, sizeof(value)); return value;
    }
    double f64() {
        uint64_t bits = u64(); double value; std::memcpy(&value, &bits, sizeof(value)); return value;
    }
    std::string string() {
        const uint32_t size = u32();
        if (offset + size > data.size()) throw std::runtime_error("checkpoint EOF string");
        std::string value(reinterpret_cast<const char *>(data.data() + offset), size);
        offset += size; return value;
    }
    std::vector<float> floats() {
        const uint64_t count = u64();
        if (count > (1ULL << 32) || offset + count*sizeof(float) > data.size()) {
            throw std::runtime_error("checkpoint invalid float vector");
        }
        std::vector<float> values(count);
        std::memcpy(values.data(), data.data() + offset, count*sizeof(float));
        offset += count*sizeof(float); return values;
    }
    std::vector<double> doubles() {
        const uint64_t count = u64();
        if (count > (1ULL << 30)) throw std::runtime_error("checkpoint invalid double vector");
        std::vector<double> values(count);
        for (uint64_t i = 0; i < count; ++i) values[i] = f64();
        return values;
    }
};

static std::vector<uint8_t> serialize_checkpoint_payload() {
    std::vector<uint8_t> payload;
    append_string(payload, g_cfg.model_sha256);
    append_string(payload, g_cfg.dataset_fingerprint);
    append_string(payload, config_identity());
    append_u64(payload, g_state.global_step);
    append_u32(payload, g_state.micro_step);
    append_u64(payload, g_state.total_micro);
    append_u64(payload, g_state.epoch);
    append_u64(payload, g_state.sample_cursor);
    append_u64(payload, g_state.rng_state);
    append_u64(payload, g_state.optimizer_updates);
    append_u64(payload, g_state.clip_applied_count);
    append_u64(payload, g_state.early_update_violations);
    append_u32(payload, static_cast<uint32_t>(g_params.size()));
    for (const param_state & state : g_params) {
        append_string(payload, state.name);
        append_float_vector(payload, state.value);
        append_float_vector(payload, state.m);
        append_float_vector(payload, state.v);
        append_float_vector(payload, state.grad_sum);
    }
    append_double_vector(payload, g_state.train_losses);
    append_double_vector(payload, g_state.learning_rates);
    append_double_vector(payload, g_state.grad_norm_pre);
    append_double_vector(payload, g_state.grad_norm_post);
    return payload;
}

static void fsync_file(const fs::path & path) {
    const int fd = open(path.c_str(), O_RDONLY);
    if (fd >= 0) {
        fsync(fd);
        close(fd);
    }
}

static void write_checkpoint_atomic(const fs::path & path) {
    fs::create_directories(path.parent_path());
    const std::vector<uint8_t> payload = serialize_checkpoint_payload();
    const uint64_t checksum = fnv1a64(payload.data(), payload.size());
    const fs::path temp = path.string() + ".tmp";
    {
        std::ofstream output(temp, std::ios::binary | std::ios::trunc);
        const std::array<char, 8> magic = {'P','1','0','C','K','0','0','1'};
        output.write(magic.data(), magic.size());
        uint32_t version = 1;
        output.write(reinterpret_cast<const char *>(&version), sizeof(version));
        uint64_t size = payload.size();
        output.write(reinterpret_cast<const char *>(&size), sizeof(size));
        output.write(reinterpret_cast<const char *>(&checksum), sizeof(checksum));
        output.write(reinterpret_cast<const char *>(payload.data()), payload.size());
        output.flush();
        if (!output) throw std::runtime_error("checkpoint write failed");
    }
    fsync_file(temp);
    fs::rename(temp, path);
    fsync_file(path);
    const std::string digest = sha256_file(path);
    std::ofstream sidecar(path.string() + ".sha256", std::ios::trunc);
    sidecar << digest << "  " << path.filename().string() << "\n";
    sidecar.flush();
    ++g_state.checkpoint_writes;
}

static std::vector<uint8_t> read_checkpoint_payload(
        const fs::path & path,
        bool * checksum_mismatch = nullptr) {
    if (checksum_mismatch) *checksum_mismatch = false;
    std::ifstream input(path, std::ios::binary);
    if (!input) throw std::runtime_error("could not open checkpoint");
    std::array<char, 8> magic = {};
    input.read(magic.data(), magic.size());
    const std::array<char, 8> expected = {'P','1','0','C','K','0','0','1'};
    if (!input || magic != expected) throw std::runtime_error("bad checkpoint magic");
    uint32_t version = 0; uint64_t size = 0; uint64_t stored = 0;
    input.read(reinterpret_cast<char *>(&version), sizeof(version));
    input.read(reinterpret_cast<char *>(&size), sizeof(size));
    input.read(reinterpret_cast<char *>(&stored), sizeof(stored));
    if (!input || version != 1 || size > (1ULL << 34)) throw std::runtime_error("bad checkpoint header");
    std::vector<uint8_t> payload(size);
    input.read(reinterpret_cast<char *>(payload.data()), payload.size());
    if (!input) throw std::runtime_error("truncated checkpoint");
    char trailing = 0;
    if (input.read(&trailing, 1)) throw std::runtime_error("checkpoint trailing bytes");
    const uint64_t actual = fnv1a64(payload.data(), payload.size());
    if (actual != stored) {
        if (checksum_mismatch) *checksum_mismatch = true;
        throw std::runtime_error("checkpoint checksum mismatch");
    }
    return payload;
}

static void load_checkpoint(const fs::path & path, bool apply_parameters) {
    const std::vector<uint8_t> payload = read_checkpoint_payload(path);
    byte_reader reader{payload};
    const std::string model_sha = reader.string();
    const std::string dataset_fp = reader.string();
    const std::string config_id = reader.string();
    if (model_sha != g_cfg.model_sha256) throw std::runtime_error("checkpoint model identity mismatch");
    if (dataset_fp != g_cfg.dataset_fingerprint) throw std::runtime_error("checkpoint dataset identity mismatch");
    if (config_id != config_identity()) throw std::runtime_error("checkpoint trainer configuration mismatch");
    g_state.global_step = reader.u64();
    g_state.micro_step = reader.u32();
    g_state.total_micro = reader.u64();
    g_state.epoch = reader.u64();
    g_state.sample_cursor = reader.u64();
    g_state.rng_state = reader.u64();
    g_state.optimizer_updates = reader.u64();
    g_state.clip_applied_count = reader.u64();
    g_state.early_update_violations = reader.u64();
    const uint32_t count = reader.u32();
    if (apply_parameters && count != g_params.size()) throw std::runtime_error("checkpoint parameter count mismatch");
    for (uint32_t i = 0; i < count; ++i) {
        const std::string name = reader.string();
        std::vector<float> value = reader.floats();
        std::vector<float> m = reader.floats();
        std::vector<float> v = reader.floats();
        std::vector<float> grad_sum = reader.floats();
        if (apply_parameters) {
            param_state & state = g_params[i];
            if (state.name != name || state.value.size() != value.size() ||
                m.size() != value.size() || v.size() != value.size() ||
                grad_sum.size() != value.size()) {
                throw std::runtime_error("checkpoint parameter layout mismatch");
            }
            state.value = std::move(value);
            state.m = std::move(m);
            state.v = std::move(v);
            state.grad_sum = std::move(grad_sum);
            set_tensor_f32(state.tensor, state.value);
        }
    }
    g_state.train_losses = reader.doubles();
    g_state.learning_rates = reader.doubles();
    g_state.grad_norm_pre = reader.doubles();
    g_state.grad_norm_post = reader.doubles();
    if (reader.offset != payload.size()) throw std::runtime_error("checkpoint payload not fully consumed");
}

static uint32_t read_u32(std::istream & input) {
    uint8_t bytes[4] = {};
    input.read(reinterpret_cast<char *>(bytes), 4);
    if (!input) throw std::runtime_error("P9DS EOF u32");
    return uint32_t(bytes[0]) | uint32_t(bytes[1]) << 8 |
           uint32_t(bytes[2]) << 16 | uint32_t(bytes[3]) << 24;
}
static uint64_t read_u64(std::istream & input) {
    uint8_t bytes[8] = {};
    input.read(reinterpret_cast<char *>(bytes), 8);
    if (!input) throw std::runtime_error("P9DS EOF u64");
    uint64_t value = 0; for (int i = 0; i < 8; ++i) value |= uint64_t(bytes[i]) << (8*i); return value;
}
static int32_t read_i32(std::istream & input) { return static_cast<int32_t>(read_u32(input)); }

static std::vector<sample_record> read_p9ds_file(const fs::path & path, int32_t * pad_token_out) {
    std::ifstream input(path, std::ios::binary);
    if (!input) throw std::runtime_error("could not open P9DS shard");
    std::array<char, 8> magic = {};
    input.read(magic.data(), magic.size());
    const std::array<char, 8> expected = {'P','9','D','S','0','0','0','1'};
    if (!input || magic != expected) throw std::runtime_error("bad P9DS magic");
    const uint32_t version = read_u32(input);
    const uint32_t split_id = read_u32(input);
    const uint32_t block_length = read_u32(input);
    const uint32_t block_count = read_u32(input);
    const int32_t pad_token = read_i32(input);
    const uint64_t seed = read_u64(input);
    const uint64_t payload_size = read_u64(input);
    const uint64_t stored_checksum = read_u64(input);
    (void) split_id; (void) seed;
    if (version != 1 || block_length == 0 || payload_size > (1ULL << 34)) {
        throw std::runtime_error("bad P9DS header");
    }
    std::vector<uint8_t> payload(payload_size);
    input.read(reinterpret_cast<char *>(payload.data()), payload.size());
    if (!input) throw std::runtime_error("truncated P9DS payload");
    if (fnv1a64(payload.data(), payload.size()) != stored_checksum) {
        throw std::runtime_error("P9DS checksum mismatch");
    }
    if (pad_token_out) *pad_token_out = pad_token;

    size_t offset = 0;
    auto take_u32 = [&]() {
        if (offset + 4 > payload.size()) throw std::runtime_error("P9DS payload EOF");
        uint32_t value = uint32_t(payload[offset]) | uint32_t(payload[offset+1]) << 8 |
                         uint32_t(payload[offset+2]) << 16 | uint32_t(payload[offset+3]) << 24;
        offset += 4; return value;
    };

    std::vector<sample_record> result;
    for (uint32_t block = 0; block < block_count; ++block) {
        const uint32_t block_index = take_u32();
        const uint32_t token_count = take_u32();
        if (block_index != block || token_count != block_length) throw std::runtime_error("P9DS block mismatch");
        std::vector<int32_t> tokens(token_count);
        for (uint32_t i = 0; i < token_count; ++i) tokens[i] = static_cast<int32_t>(take_u32());
        if (offset + token_count*2 > payload.size()) throw std::runtime_error("P9DS mask EOF");
        std::vector<uint8_t> target(payload.begin()+offset, payload.begin()+offset+token_count); offset += token_count;
        std::vector<uint8_t> starts(payload.begin()+offset, payload.begin()+offset+token_count); offset += token_count;
        std::vector<int32_t> source(token_count);
        for (uint32_t i = 0; i < token_count; ++i) source[i] = static_cast<int32_t>(take_u32());

        sample_record current;
        auto flush = [&]() {
            if (!current.tokens.empty()) result.push_back(current);
            current = sample_record{};
        };
        for (uint32_t i = 0; i < token_count; ++i) {
            if (source[i] < 0) { flush(); continue; }
            if (starts[i]) flush();
            if (current.tokens.empty()) current.source_index = source[i];
            current.tokens.push_back(tokens[i]);
            current.target_mask.push_back(target[i]);
        }
        flush();
    }
    if (offset != payload.size()) throw std::runtime_error("P9DS payload trailing bytes");
    return result;
}

static std::vector<sample_record> read_split_samples(const fs::path & directory, const std::string & prefix, int32_t * pad_token) {
    std::vector<fs::path> shards;
    for (const auto & entry : fs::directory_iterator(directory)) {
        const std::string name = entry.path().filename().string();
        if (entry.is_regular_file() && name.rfind(prefix, 0) == 0 && entry.path().extension() == ".p9ds") {
            shards.push_back(entry.path());
        }
    }
    std::sort(shards.begin(), shards.end());
    if (shards.empty()) throw std::runtime_error("no P9DS shards for split " + prefix);
    std::vector<sample_record> result;
    int32_t observed_pad = -1;
    for (const fs::path & shard : shards) {
        int32_t shard_pad = -1;
        std::vector<sample_record> samples = read_p9ds_file(shard, &shard_pad);
        if (observed_pad < 0) observed_pad = shard_pad;
        if (shard_pad != observed_pad) throw std::runtime_error("P9DS pad token mismatch");
        result.insert(result.end(), samples.begin(), samples.end());
    }
    if (pad_token) *pad_token = observed_pad;
    return result;
}

static bool usable_sample(const sample_record & sample, uint32_t context) {
    if (sample.tokens.size() < 2 || sample.tokens.size() > context || sample.tokens.size() != sample.target_mask.size()) return false;
    for (size_t i = 1; i < sample.target_mask.size(); ++i) if (sample.target_mask[i]) return true;
    return false;
}

static void initialize_parameter_registry(llama_adapter_lora * adapter) {
    g_params.clear();
    g_base_initial.clear();
    for (const target_spec & spec : g_specs) {
        const auto found = adapter->ab_map.find(spec.name);
        if (found == adapter->ab_map.end()) throw std::runtime_error(std::string("adapter target missing: ") + spec.name);
        for (const auto & item : std::array<std::pair<const char *, ggml_tensor *>, 2>{{
                {"lora_a", found->second.a}, {"lora_b", found->second.b}}}) {
            if (!item.second || item.second->type != GGML_TYPE_F32 || !(item.second->flags & GGML_TENSOR_FLAG_PARAM)) {
                throw std::runtime_error("adapter parameter is not trainable F32");
            }
            param_state state;
            state.name = std::string(spec.name) + "." + item.first;
            state.tensor = item.second;
            state.value = tensor_f32(item.second);
            state.initial = state.value;
            state.m.assign(state.value.size(), 0.0f);
            state.v.assign(state.value.size(), 0.0f);
            state.grad_sum.assign(state.value.size(), 0.0f);
            g_params.push_back(std::move(state));
        }
        const ggml_tensor * base = adapter->model->get_tensor(spec.name);
        if (!base || base->type != GGML_TYPE_Q1_0 || (base->flags & GGML_TENSOR_FLAG_PARAM)) {
            throw std::runtime_error("packed Q1 base contract failed");
        }
        g_base_initial.push_back({spec.name, tensor_bytes(base)});
    }
}

static double learning_rate_for_step(uint64_t next_step) {
    if (g_cfg.warmup_steps > 0 && next_step <= g_cfg.warmup_steps) {
        return g_cfg.base_lr * double(next_step) / double(g_cfg.warmup_steps);
    }
    if (g_cfg.target_global_steps <= g_cfg.warmup_steps) return g_cfg.base_lr;
    const double progress = std::min(1.0, std::max(0.0,
        double(next_step - g_cfg.warmup_steps) /
        double(g_cfg.target_global_steps - g_cfg.warmup_steps)));
    const double cosine = 0.5 * (1.0 + std::cos(3.14159265358979323846 * progress));
    return g_cfg.min_lr + (g_cfg.base_lr - g_cfg.min_lr) * cosine;
}

static void make_fixed_sequence(
        const sample_record & sample,
        int32_t pad_token,
        std::vector<llama_token> * tokens,
        std::vector<llama_token> * labels,
        std::vector<uint8_t> * mask) {
    tokens->assign(g_cfg.context_tokens, pad_token);
    labels->assign(g_cfg.context_tokens, pad_token);
    mask->assign(g_cfg.context_tokens, 0);
    for (size_t i = 0; i < sample.tokens.size(); ++i) (*tokens)[i] = sample.tokens[i];
    for (size_t i = 0; i + 1 < sample.tokens.size(); ++i) {
        (*labels)[i] = sample.tokens[i + 1];
        (*mask)[i] = sample.target_mask[i + 1] ? 1 : 0;
    }
}

static std::vector<float> run_backward_sample(const sample_record & sample, int32_t pad_token, double * loss_out) {
    std::vector<llama_token> tokens, labels;
    std::vector<uint8_t> mask;
    make_fixed_sequence(sample, pad_token, &tokens, &labels, &mask);
    std::vector<ggml_tensor *> tensors;
    size_t total = 0;
    for (param_state & state : g_params) { tensors.push_back(state.tensor); total += state.value.size(); }
    std::vector<float> gradients(total);
    llama_opt_masked_stats stats = {};
    const bool ok = llama_opt_masked_sequence(
        g_context, tokens.data(), labels.data(), mask.data(), g_cfg.context_tokens,
        true, tensors.data(), tensors.size(), gradients.data(), gradients.size(), &stats);
    if (!ok || !std::isfinite(stats.loss) || stats.supervised_tokens == 0) {
        throw std::runtime_error("masked backward failed");
    }
    *loss_out = stats.loss;
    return gradients;
}

static double run_validation_sample(const sample_record & sample, int32_t pad_token) {
    std::vector<llama_token> tokens, labels;
    std::vector<uint8_t> mask;
    make_fixed_sequence(sample, pad_token, &tokens, &labels, &mask);
    llama_opt_masked_stats stats = {};
    const bool ok = llama_opt_masked_sequence(
        g_context, tokens.data(), labels.data(), mask.data(), g_cfg.context_tokens,
        false, nullptr, 0, nullptr, 0, &stats);
    if (!ok || !std::isfinite(stats.loss)) throw std::runtime_error("validation forward failed");
    return stats.loss;
}

static void accumulate_gradients(const std::vector<float> & gradient) {
    size_t offset = 0;
    for (param_state & state : g_params) {
        for (size_t i = 0; i < state.value.size(); ++i) state.grad_sum[i] += gradient[offset + i];
        offset += state.value.size();
    }
}

static void apply_adamw_update() {
    double norm_squared = 0.0;
    for (const param_state & state : g_params) {
        for (float value : state.grad_sum) {
            const double averaged = double(value) / double(g_cfg.accumulation_steps);
            norm_squared += averaged * averaged;
        }
    }
    const double norm_pre = std::sqrt(norm_squared);
    const double clip = norm_pre > g_cfg.max_grad_norm
        ? double(g_cfg.max_grad_norm) / (norm_pre + 1e-30)
        : 1.0;
    if (clip < 1.0) ++g_state.clip_applied_count;
    const double norm_post = norm_pre * clip;
    const uint64_t next_step = g_state.global_step + 1;
    const double lr = learning_rate_for_step(next_step);
    const double beta1_pow = std::pow(double(g_cfg.beta1), double(next_step));
    const double beta2_pow = std::pow(double(g_cfg.beta2), double(next_step));

    for (param_state & state : g_params) {
        for (size_t i = 0; i < state.value.size(); ++i) {
            const double grad = double(state.grad_sum[i]) / double(g_cfg.accumulation_steps) * clip;
            state.m[i] = float(double(g_cfg.beta1)*state.m[i] + (1.0-double(g_cfg.beta1))*grad);
            state.v[i] = float(double(g_cfg.beta2)*state.v[i] + (1.0-double(g_cfg.beta2))*grad*grad);
            const double m_hat = double(state.m[i]) / (1.0 - beta1_pow);
            const double v_hat = double(state.v[i]) / (1.0 - beta2_pow);
            const double update = m_hat / (std::sqrt(v_hat) + g_cfg.eps) + g_cfg.weight_decay*state.value[i];
            state.value[i] = float(double(state.value[i]) - lr*update);
            state.grad_sum[i] = 0.0f;
        }
        set_tensor_f32(state.tensor, state.value);
    }

    ++g_state.global_step;
    ++g_state.optimizer_updates;
    g_state.micro_step = 0;
    g_state.learning_rates.push_back(lr);
    g_state.grad_norm_pre.push_back(norm_pre);
    g_state.grad_norm_post.push_back(norm_post);
}

static void write_float_file(const fs::path & path, const std::vector<float> & values) {
    fs::create_directories(path.parent_path());
    std::ofstream output(path, std::ios::binary);
    output.write(reinterpret_cast<const char *>(values.data()), values.size()*sizeof(float));
    if (!output) throw std::runtime_error("could not write raw parameter file");
}

static fs::path write_adapter_manifest(const fs::path & root) {
    fs::create_directories(root);
    std::ostringstream json;
    json << "{\n  \"alpha\": 8.0,\n  \"targets\": [\n";
    size_t param_index = 0;
    for (size_t target_index = 0; target_index < g_specs.size(); ++target_index) {
        const target_spec & spec = g_specs[target_index];
        const fs::path a_path = root / ("target_" + std::to_string(target_index) + "_a.bin");
        const fs::path b_path = root / ("target_" + std::to_string(target_index) + "_b.bin");
        write_float_file(a_path, g_params[param_index++].value);
        write_float_file(b_path, g_params[param_index++].value);
        if (target_index) json << ",\n";
        json << "    {\"name\":\"" << spec.name << "\",\"category\":\"" << spec.category
             << "\",\"block\":" << spec.block << ",\"K\":" << spec.K << ",\"M\":" << spec.M
             << ",\"rank\":" << spec.rank << ",\"a_file\":\"" << a_path.string()
             << "\",\"b_file\":\"" << b_path.string() << "\"}";
    }
    json << "\n  ]\n}\n";
    const fs::path manifest = root / "manifest.json";
    std::ofstream output(manifest); output << json.str();
    return manifest;
}

static fs::path export_adapter(const fs::path & output_dir) {
    const fs::path raw = output_dir / "raw_adapter";
    const fs::path manifest = write_adapter_manifest(raw);
    const fs::path output = output_dir / "step10_adapter_final.gguf";
    const fs::path helper = "/content/prism_native_q1_lora/step08_multitarget_implementation/write_adapter_from_raw.py";
    std::ostringstream command;
    command << "python3 " << shell_quote(helper.string())
            << " --manifest " << shell_quote(manifest.string())
            << " --output " << shell_quote(output.string())
            << " --name " << shell_quote("Bonsai-27B Step 10 accepted trainer");
    const int rc = std::system(command.str().c_str());
    if (rc == -1 || !WIFEXITED(rc) || WEXITSTATUS(rc) != 0 || !fs::is_regular_file(output)) {
        throw std::runtime_error("adapter export helper failed");
    }
    return output;
}

static double verify_adapter_reload(const fs::path & path) {
    llama_adapter_lora * reloaded = llama_adapter_lora_init(g_model, path.c_str());
    if (!reloaded) throw std::runtime_error("could not reload final adapter");
    double max_diff = 0.0;
    size_t param_index = 0;
    for (const target_spec & spec : g_specs) {
        const auto found = reloaded->ab_map.find(spec.name);
        if (found == reloaded->ab_map.end()) throw std::runtime_error("reloaded adapter target missing");
        for (ggml_tensor * tensor : std::array<ggml_tensor *, 2>{{found->second.a, found->second.b}}) {
            const std::vector<float> values = tensor_f32(tensor);
            const std::vector<float> & expected = g_params[param_index++].value;
            if (values.size() != expected.size()) throw std::runtime_error("reloaded adapter shape mismatch");
            for (size_t i = 0; i < values.size(); ++i) max_diff = std::max(max_diff, std::abs(double(values[i])-expected[i]));
        }
    }
    llama_adapter_lora_free(reloaded);
    return max_diff;
}

static void execute_training() {
    g_outcome.hook_seen = true;
    try {
        fs::create_directories(g_cfg.output_dir);
        int32_t train_pad = -1, validation_pad = -1;
        std::vector<sample_record> train_all = read_split_samples(g_cfg.train_dir, "train-", &train_pad);
        std::vector<sample_record> validation_all = read_split_samples(g_cfg.validation_dir, "validation-", &validation_pad);
        if (train_pad != validation_pad) throw std::runtime_error("train/validation pad mismatch");
        std::vector<sample_record> train_samples, validation_samples;
        for (const sample_record & sample : train_all) if (usable_sample(sample, g_cfg.context_tokens)) train_samples.push_back(sample);
        for (const sample_record & sample : validation_all) if (usable_sample(sample, g_cfg.context_tokens)) validation_samples.push_back(sample);
        if (train_samples.empty() || validation_samples.empty()) throw std::runtime_error("no acceptance samples fit context");
        g_outcome.train_sample_count = train_samples.size();
        g_outcome.validation_sample_count = validation_samples.size();

        if (!g_cfg.resume_path.empty() && g_cfg.resume_path != "-") {
            load_checkpoint(g_cfg.resume_path, true);
        }

        const fs::path checkpoint = fs::path(g_cfg.output_dir) / "checkpoint_latest.p10ck";
        bool stopped_early = false;
        while (g_state.global_step < g_cfg.target_global_steps) {
            const sample_record & sample = train_samples[g_state.sample_cursor % train_samples.size()];
            const uint64_t before_fp = parameter_fingerprint();
            double loss = NAN;
            const std::vector<float> gradient = run_backward_sample(sample, train_pad, &loss);
            const uint64_t after_backward_fp = parameter_fingerprint();
            if (before_fp != after_backward_fp) ++g_state.early_update_violations;
            accumulate_gradients(gradient);
            g_state.train_losses.push_back(loss);
            ++g_state.micro_step;
            ++g_state.total_micro;
            ++g_state.sample_cursor;
            if (g_state.sample_cursor % train_samples.size() == 0) ++g_state.epoch;
            if (g_state.micro_step == g_cfg.accumulation_steps) {
                apply_adamw_update();
                write_checkpoint_atomic(checkpoint);
            }
            if (g_cfg.stop_after_micro > 0 && g_state.total_micro >= g_cfg.stop_after_micro) {
                write_checkpoint_atomic(checkpoint);
                stopped_early = true;
                break;
            }
        }
        if (!stopped_early) write_checkpoint_atomic(checkpoint);

        const uint64_t before_validation = parameter_fingerprint();
        g_outcome.validation_loss = run_validation_sample(validation_samples.front(), validation_pad);
        const uint64_t after_validation = parameter_fingerprint();
        g_outcome.validation_unchanged = before_validation == after_validation;

        g_outcome.base_changed_bytes = 0;
        for (const auto & entry : g_base_initial) {
            const ggml_tensor * base = g_model->get_tensor(entry.first.c_str());
            g_outcome.base_changed_bytes += changed_bytes(entry.second, tensor_bytes(base));
        }

        const fs::path adapter = export_adapter(g_cfg.output_dir);
        g_outcome.adapter_reload_max_diff = verify_adapter_reload(adapter);
        g_outcome.checkpoint_path = checkpoint.string();
        g_outcome.checkpoint_sha256 = sha256_file(checkpoint);
        g_outcome.adapter_path = adapter.string();
        g_outcome.adapter_sha256 = sha256_file(adapter);
        g_outcome.state_fingerprint = hex64(parameter_fingerprint());

        const bool params_changed = std::any_of(g_params.begin(), g_params.end(), [](const param_state & state) {
            return state.value != state.initial;
        });
        g_outcome.success =
            std::isfinite(g_outcome.validation_loss) &&
            g_outcome.validation_unchanged &&
            g_outcome.base_changed_bytes == 0 &&
            g_outcome.adapter_reload_max_diff <= 1e-7 &&
            g_state.early_update_violations == 0 &&
            params_changed &&
            fs::is_regular_file(checkpoint) &&
            !fs::exists(checkpoint.string() + ".tmp");
    } catch (const std::exception & error) {
        g_outcome.error = error.what();
        g_outcome.success = false;
    }
}

static llama_adapter_lora * prism10_adapter_init_hook(llama_model * model, const char * path) {
    llama_adapter_lora * adapter = llama_adapter_lora_init(model, path);
    if (adapter) {
        g_adapter = adapter;
        g_model = model;
        initialize_parameter_registry(adapter);
    }
    return adapter;
}

static void prism10_adapter_free_hook(llama_adapter_lora * adapter) {
    llama_adapter_lora_free(adapter);
    g_adapter = nullptr;
}

static void prism10_opt_init_hook(llama_context * ctx, llama_model * model, llama_opt_params params) {
    params.n_ctx_train = g_cfg.context_tokens;
    llama_opt_init(ctx, model, params);
    g_context = ctx;
    g_model = model;
}

static void prism10_opt_epoch_hook(
        llama_context *, ggml_opt_dataset_t, ggml_opt_result_t, ggml_opt_result_t,
        int64_t, ggml_opt_epoch_callback, ggml_opt_epoch_callback) {
    execute_training();
}

#define llama_adapter_lora_init prism10_adapter_init_hook
#define llama_adapter_lora_free prism10_adapter_free_hook
#define llama_opt_init prism10_opt_init_hook
#define llama_opt_epoch prism10_opt_epoch_hook
#define main prism_stage7_bootstrap_main
#include "test-q1-lora-full-backward.cpp"
#undef main
#undef llama_opt_epoch
#undef llama_opt_init
#undef llama_adapter_lora_free
#undef llama_adapter_lora_init

static int verify_checkpoint_only(const fs::path & path) {
    bool checksum_mismatch = false;
    try {
        const std::vector<uint8_t> payload = read_checkpoint_payload(path, &checksum_mismatch);
        byte_reader reader{payload};
        const std::string model_sha = reader.string();
        const std::string dataset_fp = reader.string();
        const std::string config_id = reader.string();
        if (model_sha != g_cfg.model_sha256) throw std::runtime_error("checkpoint model identity mismatch");
        if (dataset_fp != g_cfg.dataset_fingerprint) throw std::runtime_error("checkpoint dataset identity mismatch");
        if (config_id != config_identity()) throw std::runtime_error("checkpoint trainer configuration mismatch");
        std::cout << "CHECKPOINT_CHECKSUM_MISMATCH=0\n";
        std::cout << "CHECKPOINT_IDENTITY_MATCH=1\n";
        std::cout << "SUBTEST_STATUS=PASS\nFINAL_STATUS=PASS\n";
        return 0;
    } catch (const std::exception & error) {
        std::cout << "VERIFY_ERROR=" << error.what() << "\n";
        std::cout << "CHECKPOINT_CHECKSUM_MISMATCH=" << (checksum_mismatch ? 1 : 0) << "\n";
        std::cout << "CHECKPOINT_IDENTITY_MATCH=0\n";
        std::cout << "SUBTEST_STATUS=FAIL\nFINAL_STATUS=FAIL\n";
        return 1;
    }
}

int main(int argc, char ** argv) {
    try {
        if (argc < 20) {
            std::cerr << "usage: test-q1-lora-step10 MODEL ADAPTER TRAIN_DIR VALIDATION_DIR OUTPUT_DIR MODE RESUME MODEL_SHA DATASET_FP SEED CONTEXT TARGET_STEPS ACCUM BASE_LR MIN_LR WARMUP MAX_NORM STOP_AFTER_MICRO CHECKPOINT_TO_VERIFY\n";
            return 2;
        }

        // PRISM_STEP10_V3_TRAINING_GRAPH_ENV_BEGIN
        // These are set inside the executable so direct CLI use is identical
        // to the Python acceptance harness. They must precede adapter/model and
        // context creation.
        ::setenv("PRISM_Q1_LORA_TRAINING", "1", 1);
        ::setenv("PRISM_Q1_LORA_UNFUSED_GDN", "1", 1);
        ::setenv("PRISM_Q1_LORA_TRAINING_GENERIC_SSM_CONV", "1", 1);
        ::setenv("PRISM_Q1_LORA_TRAINING_NO_KV_CACHE", "1", 1);
        ::setenv("LLAMA_GRAPH_REUSE_DISABLE", "1", 1);
        ::setenv("PRISM_STEP10_FORCE_OPT_BACKWARD", "0", 1);
        ::setenv("PRISM_STEP10_EXTERNAL_GRAD_ACCUM", "1", 1);
        std::cerr << "PRISM_STEP10_TRAINING_ENV_READY=1\n";
        // PRISM_STEP10_V3_TRAINING_GRAPH_ENV_END

        g_cfg.model_path = argv[1];
        g_cfg.adapter_path = argv[2];
        g_cfg.train_dir = argv[3];
        g_cfg.validation_dir = argv[4];
        g_cfg.output_dir = argv[5];
        g_cfg.mode = argv[6];
        g_cfg.resume_path = argv[7];
        g_cfg.model_sha256 = argv[8];
        g_cfg.dataset_fingerprint = argv[9];
        g_cfg.seed = std::stoull(argv[10]);
        g_cfg.context_tokens = std::stoul(argv[11]);
        g_cfg.target_global_steps = std::stoull(argv[12]);
        g_cfg.accumulation_steps = std::stoul(argv[13]);
        g_cfg.base_lr = std::stof(argv[14]);
        g_cfg.min_lr = std::stof(argv[15]);
        g_cfg.warmup_steps = std::stoul(argv[16]);
        g_cfg.max_grad_norm = std::stof(argv[17]);
        g_cfg.stop_after_micro = std::stoull(argv[18]);
        const std::string verify_path = argv[19];
        g_state.rng_state = g_cfg.seed;

        if (g_cfg.context_tokens < 8 || g_cfg.accumulation_steps == 0 || g_cfg.base_lr <= 0 || g_cfg.max_grad_norm <= 0) {
            throw std::runtime_error("invalid trainer configuration");
        }
        if (g_cfg.mode == "verify_checkpoint") {
            return verify_checkpoint_only(verify_path);
        }

        std::vector<std::string> args = {
            "test-q1-lora-full-backward",
            g_cfg.model_path,
            g_cfg.adapter_path,
            "SSM",
            "blk.0.ssm_alpha.weight",
        };
        std::vector<char *> stage7_argv;
        for (std::string & value : args) stage7_argv.push_back(value.data());
        const int bootstrap_rc = prism_stage7_bootstrap_main(stage7_argv.size(), stage7_argv.data());

        std::cout << "BOOTSTRAP_RETURN_CODE=" << bootstrap_rc << "\n";
        std::cout << "TRAINING_HOOK_SEEN=" << (g_outcome.hook_seen ? 1 : 0) << "\n";
        std::cout << "TRAIN_SAMPLE_COUNT=" << g_outcome.train_sample_count << "\n";
        std::cout << "VALIDATION_SAMPLE_COUNT=" << g_outcome.validation_sample_count << "\n";
        std::cout << "OPTIMIZER_PARAMETER_COUNT=" << g_params.size() << "\n";
        size_t trainable_values = 0; for (const param_state & state : g_params) trainable_values += state.value.size();
        std::cout << "TRAINABLE_PARAMETER_VALUES=" << trainable_values << "\n";
        std::cout << "GLOBAL_STEP=" << g_state.global_step << "\n";
        std::cout << "MICRO_STEP=" << g_state.micro_step << "\n";
        std::cout << "TOTAL_MICRO_STEPS=" << g_state.total_micro << "\n";
        std::cout << "DATASET_EPOCH=" << g_state.epoch << "\n";
        std::cout << "DATASET_CURSOR=" << g_state.sample_cursor << "\n";
        std::cout << "OPTIMIZER_UPDATES=" << g_state.optimizer_updates << "\n";
        std::cout << "CHECKPOINT_WRITE_COUNT=" << g_state.checkpoint_writes << "\n";
        std::cout << "EARLY_UPDATE_VIOLATIONS=" << g_state.early_update_violations << "\n";
        std::cout << "CLIP_APPLIED_COUNT=" << g_state.clip_applied_count << "\n";
        std::cout << "TRAIN_LOSS_COUNT=" << g_state.train_losses.size() << "\n";
        if (!g_state.train_losses.empty()) {
            std::cout << std::setprecision(17)
                      << "TRAIN_LOSS_FIRST=" << g_state.train_losses.front() << "\n"
                      << "TRAIN_LOSS_LAST=" << g_state.train_losses.back() << "\n";
        }
        if (!g_state.learning_rates.empty()) {
            std::cout << std::setprecision(17)
                      << "LR_FIRST=" << g_state.learning_rates.front() << "\n"
                      << "LR_LAST=" << g_state.learning_rates.back() << "\n";
        }
        if (!g_state.grad_norm_pre.empty()) {
            std::cout << std::setprecision(17)
                      << "GRAD_NORM_PRE_LAST=" << g_state.grad_norm_pre.back() << "\n"
                      << "GRAD_NORM_POST_LAST=" << g_state.grad_norm_post.back() << "\n";
        }
        std::cout << std::setprecision(17) << "VALIDATION_LOSS=" << g_outcome.validation_loss << "\n";
        std::cout << "VALIDATION_PARAMETERS_UNCHANGED=" << (g_outcome.validation_unchanged ? 1 : 0) << "\n";
        std::cout << "BASE_CHANGED_BYTES=" << g_outcome.base_changed_bytes << "\n";
        std::cout << "PERSISTENT_EXPANDED_BASE_BYTES=0\n";
        std::cout << std::setprecision(17) << "ADAPTER_RELOAD_MAX_ABS_DIFF=" << g_outcome.adapter_reload_max_diff << "\n";
        std::cout << "CHECKPOINT_PATH=" << g_outcome.checkpoint_path << "\n";
        std::cout << "CHECKPOINT_SHA256=" << g_outcome.checkpoint_sha256 << "\n";
        std::cout << "ADAPTER_PATH=" << g_outcome.adapter_path << "\n";
        std::cout << "ADAPTER_SHA256=" << g_outcome.adapter_sha256 << "\n";
        std::cout << "FINAL_STATE_FINGERPRINT=" << g_outcome.state_fingerprint << "\n";
        std::cout << "CHECKPOINT_ATOMIC=" << (!g_outcome.checkpoint_path.empty() && !fs::exists(g_outcome.checkpoint_path + ".tmp") ? 1 : 0) << "\n";
        if (!g_outcome.error.empty()) std::cout << "TRAINER_ERROR=" << g_outcome.error << "\n";
        std::cout << "SUBTEST_STATUS=" << (g_outcome.success ? "PASS" : "FAIL") << "\n";
        std::cout << "FINAL_STATUS=" << (g_outcome.success ? "PASS" : "FAIL") << "\n";
        return g_outcome.success ? 0 : 1;
    } catch (const std::exception & error) {
        std::cerr << "STEP10_ERROR=" << error.what() << "\n";
        std::cout << "SUBTEST_STATUS=FAIL\nFINAL_STATUS=FAIL\n";
        return 1;
    }
}