// STEP 8 native multi-target packed-Q1 LoRA acceptance executable.
//
// The accepted Stage 7 full-model optimizer is compiled into this translation
// unit under a renamed main(). Adapter init/free are intercepted only to
// snapshot all A/B and packed-base tensors around the same optimizer epoch.

#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 <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <regex>
#include <sstream>
#include <string>
#include <thread>
#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 target_snapshot {
    target_spec spec {};
    std::vector<float> a_initial;
    std::vector<float> b_initial;
    std::vector<float> a_final;
    std::vector<float> b_final;
    std::vector<uint8_t> base_initial;
    std::vector<uint8_t> base_final;
    bool a_param = false;
    bool b_param = false;
    bool base_param = false;
    std::string a_buffer;
    std::string b_buffer;
};

static std::map<std::string, target_snapshot> g_snapshots;
static bool g_hook_init_seen = false;
static bool g_hook_free_seen = false;

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

static std::string read_text(const fs::path & path) {
    std::ifstream in(path, std::ios::binary);
    std::ostringstream ss;
    ss << in.rdbuf();
    return ss.str();
}

static void write_text(const fs::path & path, const std::string & text) {
    fs::create_directories(path.parent_path());
    std::ofstream out(path, std::ios::binary);
    out << text;
    if (!out) {
        throw std::runtime_error("failed writing " + path.string());
    }
}

static std::string sha256_file(const fs::path & path) {
    std::string command = "sha256sum " + shell_quote(path.string());
    FILE * pipe = popen(command.c_str(), "r");
    if (!pipe) {
        throw std::runtime_error("popen sha256sum failed");
    }
    char buffer[256] = {};
    std::string output;
    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 for " + path.string());
    }
    return output.substr(0, 64);
}

static int gpu_memory_mib() {
    FILE * pipe = popen(
        "nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null",
        "r");
    if (!pipe) {
        return -1;
    }
    char buffer[128] = {};
    int total = 0;
    bool any = false;
    while (fgets(buffer, sizeof(buffer), pipe)) {
        total += std::atoi(buffer);
        any = true;
    }
    pclose(pipe);
    return any ? total : -1;
}

static std::vector<float> tensor_to_f32(const ggml_tensor * tensor) {
    if (!tensor || !tensor->buffer) {
        throw std::runtime_error("tensor_to_f32 received an unallocated tensor");
    }

    const size_t n = ggml_nelements(tensor);
    std::vector<float> result(n);

    if (tensor->type == GGML_TYPE_F32) {
        ggml_backend_tensor_get(tensor, result.data(), 0, n * sizeof(float));
        return result;
    }

    if (tensor->type == GGML_TYPE_F16) {
        std::vector<ggml_fp16_t> tmp(n);
        ggml_backend_tensor_get(tensor, tmp.data(), 0, n * sizeof(ggml_fp16_t));
        for (size_t i = 0; i < n; ++i) {
            result[i] = ggml_fp16_to_fp32(tmp[i]);
        }
        return result;
    }

    throw std::runtime_error(
        std::string("adapter tensor is not F32/F16: ") + ggml_type_name(tensor->type));
}

static std::vector<uint8_t> tensor_bytes(const ggml_tensor * tensor) {
    if (!tensor || !tensor->buffer) {
        throw std::runtime_error("tensor_bytes received an unallocated tensor");
    }
    std::vector<uint8_t> result(ggml_nbytes(tensor));
    ggml_backend_tensor_get(tensor, result.data(), 0, result.size());
    return result;
}

static bool is_parameter(const ggml_tensor * tensor) {
    return tensor && (tensor->flags & GGML_TENSOR_FLAG_PARAM) != 0;
}

static const target_spec * find_spec(const std::string & name) {
    for (const auto & spec : g_specs) {
        if (name == spec.name) {
            return &spec;
        }
    }
    return nullptr;
}

static llama_adapter_lora * prism_mt_adapter_init_hook(
        llama_model * model,
        const char * path_lora) {
    llama_adapter_lora * adapter = llama_adapter_lora_init(model, path_lora);
    if (!adapter) {
        return nullptr;
    }

    g_hook_init_seen = true;
    g_snapshots.clear();

    for (const auto & spec : g_specs) {
        const auto it = adapter->ab_map.find(spec.name);
        if (it == adapter->ab_map.end()) {
            continue;
        }

        target_snapshot snapshot;
        snapshot.spec = spec;
        snapshot.a_initial = tensor_to_f32(it->second.a);
        snapshot.b_initial = tensor_to_f32(it->second.b);
        snapshot.a_param = is_parameter(it->second.a);
        snapshot.b_param = is_parameter(it->second.b);
        snapshot.a_buffer = it->second.a && it->second.a->buffer
            ? ggml_backend_buffer_name(it->second.a->buffer) : "<null>";
        snapshot.b_buffer = it->second.b && it->second.b->buffer
            ? ggml_backend_buffer_name(it->second.b->buffer) : "<null>";

        const ggml_tensor * base = model->get_tensor(spec.name);
        if (!base) {
            throw std::runtime_error(std::string("base tensor missing: ") + spec.name);
        }
        snapshot.base_initial = tensor_bytes(base);
        snapshot.base_param = is_parameter(base);
        g_snapshots.emplace(spec.name, std::move(snapshot));
    }

    return adapter;
}

static void prism_mt_adapter_free_hook(llama_adapter_lora * adapter) {
    if (adapter) {
        for (auto & entry : g_snapshots) {
            const auto it = adapter->ab_map.find(entry.first);
            if (it == adapter->ab_map.end()) {
                continue;
            }
            entry.second.a_final = tensor_to_f32(it->second.a);
            entry.second.b_final = tensor_to_f32(it->second.b);
            const ggml_tensor * base = adapter->model->get_tensor(entry.first.c_str());
            entry.second.base_final = tensor_bytes(base);
        }
        g_hook_free_seen = true;
    }
    llama_adapter_lora_free(adapter);
}

// Compile the already accepted Stage 7 graph/optimizer into this test.
// Only adapter init/free are intercepted; llama_opt_epoch remains the real one.
#define llama_adapter_lora_init prism_mt_adapter_init_hook
#define llama_adapter_lora_free prism_mt_adapter_free_hook
#define main prism_stage7_full_backward_main
#include "test-q1-lora-full-backward.cpp"
#undef main
#undef llama_adapter_lora_free
#undef llama_adapter_lora_init

struct captured_run {
    int rc = -1;
    std::string output;
};

static captured_run run_stage7_core(
        const std::string & model,
        const std::string & adapter,
        const fs::path & capture_path) {
    fs::create_directories(capture_path.parent_path());

    fflush(nullptr);
    std::cout.flush();
    std::cerr.flush();

    const int saved_stdout = dup(STDOUT_FILENO);
    const int saved_stderr = dup(STDERR_FILENO);
    if (saved_stdout < 0 || saved_stderr < 0) {
        throw std::runtime_error("dup failed");
    }

    FILE * capture = fopen(capture_path.c_str(), "w+");
    if (!capture) {
        throw std::runtime_error("fopen capture failed");
    }

    if (dup2(fileno(capture), STDOUT_FILENO) < 0 ||
        dup2(fileno(capture), STDERR_FILENO) < 0) {
        throw std::runtime_error("dup2 capture failed");
    }

    std::vector<std::string> args = {
        "test-q1-lora-full-backward",
        model,
        adapter,
        "SSM",
        "blk.0.ssm_alpha.weight",
    };
    std::vector<char *> argv;
    for (auto & value : args) {
        argv.push_back(value.data());
    }

    int rc = prism_stage7_full_backward_main((int) argv.size(), argv.data());

    fflush(nullptr);
    std::cout.flush();
    std::cerr.flush();

    dup2(saved_stdout, STDOUT_FILENO);
    dup2(saved_stderr, STDERR_FILENO);
    close(saved_stdout);
    close(saved_stderr);
    fclose(capture);

    return {rc, read_text(capture_path)};
}

static double max_abs_change(
        const std::vector<float> & before,
        const std::vector<float> & after) {
    if (before.size() != after.size() || before.empty()) {
        return 0.0;
    }
    double result = 0.0;
    for (size_t i = 0; i < before.size(); ++i) {
        result = std::max(result, std::abs((double) before[i] - (double) after[i]));
    }
    return result;
}

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

static std::map<std::string, double> parse_gradient_maxima(const std::string & output) {
    std::map<std::string, double> result;
    const std::regex pattern(
        R"(PRISM_Q1_LORA_OPT_GRAD name=([^\s]+) present=1 max_abs=([+\-0-9.eE]+))");

    for (std::sregex_iterator it(output.begin(), output.end(), pattern), end; it != end; ++it) {
        const std::string name = (*it)[1].str();
        const double value = std::strtod((*it)[2].str().c_str(), nullptr);
        auto found = result.find(name);
        if (found == result.end()) {
            result[name] = value;
        } else {
            found->second = std::max(found->second, value);
        }
    }
    return result;
}

static void write_float_file(const fs::path & path, const std::vector<float> & data) {
    fs::create_directories(path.parent_path());
    std::ofstream out(path, std::ios::binary);
    out.write(reinterpret_cast<const char *>(data.data()), (std::streamsize) (data.size() * sizeof(float)));
    if (!out) {
        throw std::runtime_error("failed writing float file: " + path.string());
    }
}

static fs::path write_snapshot_manifest(
        const fs::path & root,
        bool initial) {
    fs::create_directories(root);
    std::ostringstream json;
    json << "{\n  \"alpha\": 8.0,\n  \"targets\": [\n";

    size_t index = 0;
    for (const auto & spec : g_specs) {
        const auto it = g_snapshots.find(spec.name);
        if (it == g_snapshots.end()) {
            throw std::runtime_error(std::string("missing snapshot: ") + spec.name);
        }
        const auto & snap = it->second;
        const auto & a = initial ? snap.a_initial : snap.a_final;
        const auto & b = initial ? snap.b_initial : snap.b_final;

        const fs::path a_path = root / ("target_" + std::to_string(index) + "_a.bin");
        const fs::path b_path = root / ("target_" + std::to_string(index) + "_b.bin");
        write_float_file(a_path, a);
        write_float_file(b_path, b);

        if (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() << "\""
             << "}";
        ++index;
    }

    json << "\n  ]\n}\n";
    const fs::path manifest = root / "manifest.json";
    write_text(manifest, json.str());
    return manifest;
}

static void create_adapter_from_manifest(
        const fs::path & manifest,
        const fs::path & output,
        const std::string & name) {
    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(name);

    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 writer helper failed");
    }
}

static std::string json_escape(const std::string & input) {
    std::ostringstream out;
    for (unsigned char c : input) {
        switch (c) {
            case '"': out << "\\\""; break;
            case '\\': out << "\\\\"; break;
            case '\b': out << "\\b"; break;
            case '\f': out << "\\f"; break;
            case '\n': out << "\\n"; break;
            case '\r': out << "\\r"; break;
            case '\t': out << "\\t"; break;
            default:
                if (c < 0x20) {
                    out << "\\u" << std::hex << std::setw(4) << std::setfill('0') << (int) c;
                } else {
                    out << c;
                }
        }
    }
    return out.str();
}

static bool emit_diagnostics(
        const std::map<std::string, double> & gradients,
        size_t * total_base_changed,
        bool print_json) {
    bool pass = true;
    size_t base_changed = 0;

    for (const auto & spec : g_specs) {
        const auto it = g_snapshots.find(spec.name);
        if (it == g_snapshots.end()) {
            pass = false;
            continue;
        }

        const auto & snap = it->second;
        const std::string a_name = std::string(spec.name) + ".lora_a";
        const std::string b_name = std::string(spec.name) + ".lora_b";
        const double a_grad = gradients.count(a_name) ? gradients.at(a_name) : 0.0;
        const double b_grad = gradients.count(b_name) ? gradients.at(b_name) : 0.0;
        const double a_update = max_abs_change(snap.a_initial, snap.a_final);
        const double b_update = max_abs_change(snap.b_initial, snap.b_final);
        const size_t target_base_changed = changed_bytes(snap.base_initial, snap.base_final);
        const int64_t parameter_count =
            (int64_t) snap.a_initial.size() + (int64_t) snap.b_initial.size();
        const int64_t optimizer_state_bytes = parameter_count * 2 * (int64_t) sizeof(float);

        base_changed += target_base_changed;

        const bool target_pass =
            std::isfinite(a_grad) && a_grad > 0.0 &&
            std::isfinite(b_grad) && b_grad > 0.0 &&
            std::isfinite(a_update) && a_update > 0.0 &&
            std::isfinite(b_update) && b_update > 0.0 &&
            target_base_changed == 0 &&
            snap.a_param && snap.b_param && !snap.base_param;

        pass = pass && target_pass;

        if (print_json) {
            std::cout
                << "TARGET_DIAG_JSON={"
                << "\"tensor_name\":\"" << json_escape(spec.name) << "\","
                << "\"category\":\"" << spec.category << "\","
                << "\"block\":" << spec.block << ","
                << "\"K\":" << spec.K << ","
                << "\"M\":" << spec.M << ","
                << "\"rank\":" << spec.rank << ","
                << "\"alpha\":" << spec.alpha << ","
                << "\"a_grad_max\":" << std::setprecision(17) << a_grad << ","
                << "\"b_grad_max\":" << std::setprecision(17) << b_grad << ","
                << "\"a_update_max\":" << std::setprecision(17) << a_update << ","
                << "\"b_update_max\":" << std::setprecision(17) << b_update << ","
                << "\"parameter_count\":" << parameter_count << ","
                << "\"optimizer_state_bytes\":" << optimizer_state_bytes << ","
                << "\"enabled\":1,"
                << "\"base_is_parameter\":" << (snap.base_param ? 1 : 0) << ","
                << "\"a_is_parameter\":" << (snap.a_param ? 1 : 0) << ","
                << "\"b_is_parameter\":" << (snap.b_param ? 1 : 0)
                << "}\n";
        }
    }

    *total_base_changed = base_changed;
    return pass;
}

static int run_combined_training(
        const std::string & model_path,
        const std::string & output_dir,
        const std::string & mode) {
    fs::create_directories(output_dir);
    const fs::path template_path =
        "/content/prism_native_q1_lora/step08_multitarget_implementation/multi_target_adapter_template.gguf";
    const fs::path capture_path = fs::path(output_dir) / "stage7_combined_core.log";

    g_hook_init_seen = false;
    g_hook_free_seen = false;
    g_snapshots.clear();

    const captured_run core = run_stage7_core(model_path, template_path.string(), capture_path);
    const auto gradients = parse_gradient_maxima(core.output);

    size_t total_base_changed = 0;
    const bool diagnostics_ok =
        g_hook_init_seen &&
        g_hook_free_seen &&
        g_snapshots.size() == g_specs.size() &&
        emit_diagnostics(gradients, &total_base_changed, true);

    const bool full_backward =
        core.output.find("PROBE_BACKWARD_RETURNED=1") != std::string::npos &&
        core.output.find("CHECK_COMPLETE_MODEL_BACKWARD=PASS") != std::string::npos;
    const bool no_expanded =
        core.output.find("PERSISTENT_EXPANDED_WEIGHT_BYTES=0") != std::string::npos ||
        core.output.find("PERSISTENT_EXPANDED_BASE_BYTES=0") != std::string::npos;

    const fs::path raw_initial = fs::path(output_dir) / "raw_initial";
    const fs::path raw_final = fs::path(output_dir) / "raw_final";
    const fs::path initial_manifest = write_snapshot_manifest(raw_initial, true);
    const fs::path final_manifest = write_snapshot_manifest(raw_final, false);

    const fs::path initial_adapter = fs::path(output_dir) / "multi_target_adapter_initial.gguf";
    const fs::path updated_adapter = fs::path(output_dir) / "multi_target_adapter_updated.gguf";

    create_adapter_from_manifest(initial_manifest, initial_adapter, "Bonsai-27B Step 8 initial trio");
    create_adapter_from_manifest(final_manifest, updated_adapter, "Bonsai-27B Step 8 updated trio");

    const std::string initial_sha = sha256_file(initial_adapter);
    const std::string updated_sha = sha256_file(updated_adapter);

    std::cout << "CORE_RETURN_CODE=" << core.rc << "\n";
    std::cout << "TARGET_COUNT=" << g_snapshots.size() << "\n";
    std::cout << "OPTIMIZER_PARAMETER_COUNT=" << (g_snapshots.size() * 2) << "\n";
    std::cout << "BASE_CHANGED_BYTES=" << total_base_changed << "\n";
    std::cout << "PERSISTENT_EXPANDED_BASE_BYTES=" << (no_expanded ? 0 : -1) << "\n";
    std::cout << "ADAPTER_INITIAL_PATH=" << initial_adapter.string() << "\n";
    std::cout << "ADAPTER_UPDATED_PATH=" << updated_adapter.string() << "\n";
    std::cout << "ADAPTER_INITIAL_SHA256=" << initial_sha << "\n";
    std::cout << "ADAPTER_UPDATED_SHA256=" << updated_sha << "\n";
    std::cout << "DETERMINISM_FINGERPRINT=" << updated_sha << "\n";

    // Cache the accepted trio output for the independent save/reload mode.
    if (mode == "trio") {
        const fs::path cache =
            "/content/prism_native_q1_lora/step08_multitarget_implementation/cache";
        fs::create_directories(cache);
        fs::copy_file(updated_adapter, cache / "multi_target_adapter_updated.gguf",
                      fs::copy_options::overwrite_existing);
        fs::remove_all(cache / "raw_final");
        fs::copy(raw_final, cache / "raw_final",
                 fs::copy_options::recursive | fs::copy_options::overwrite_existing);
    }

    const bool pass =
        diagnostics_ok &&
        full_backward &&
        no_expanded &&
        total_base_changed == 0 &&
        g_snapshots.size() == 3;

    std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    return pass ? 0 : 1;
}

static llama_model * load_model(const std::string & model_path) {
    llama_model_params params = llama_model_default_params();
    params.n_gpu_layers = 999;
    params.use_mmap = true;
    params.check_tensors = true;
    llama_model * model = llama_model_load_from_file(model_path.c_str(), params);
    if (!model) {
        throw std::runtime_error("failed loading model");
    }
    return model;
}

static bool inspect_registry(const std::string & model_path, bool verbose) {
    const fs::path template_path =
        "/content/prism_native_q1_lora/step08_multitarget_implementation/multi_target_adapter_template.gguf";

    llama_backend_init();
    llama_model * model = load_model(model_path);
    llama_adapter_lora * adapter = llama_adapter_lora_init(model, template_path.c_str());
    if (!adapter) {
        llama_model_free(model);
        llama_backend_free();
        return false;
    }

    bool pass = adapter->ab_map.size() == 3;
    int param_count = 0;

    for (const auto & spec : g_specs) {
        const auto found = adapter->ab_map.find(spec.name);
        const ggml_tensor * base = model->get_tensor(spec.name);
        const bool found_pair = found != adapter->ab_map.end();
        const bool shape_ok = found_pair &&
            found->second.a->ne[0] == spec.K &&
            found->second.a->ne[1] == spec.rank &&
            found->second.b->ne[0] == spec.rank &&
            found->second.b->ne[1] == spec.M;
        const bool flags_ok = found_pair &&
            is_parameter(found->second.a) &&
            is_parameter(found->second.b) &&
            base && !is_parameter(base);
        const bool q1_ok = base && base->type == GGML_TYPE_Q1_0;

        pass = pass && found_pair && shape_ok && flags_ok && q1_ok;
        param_count += found_pair ? 2 : 0;

        if (verbose) {
            std::cout << "REGISTRY_TARGET=" << spec.name
                      << " FOUND=" << (found_pair ? 1 : 0)
                      << " SHAPE_OK=" << (shape_ok ? 1 : 0)
                      << " FLAGS_OK=" << (flags_ok ? 1 : 0)
                      << " BASE_Q1=" << (q1_ok ? 1 : 0)
                      << "\n";
        }
    }

    std::cout << "TARGET_COUNT=" << adapter->ab_map.size() << "\n";
    std::cout << "OPTIMIZER_PARAMETER_COUNT=" << param_count << "\n";

    llama_adapter_lora_free(adapter);
    llama_model_free(model);
    llama_backend_free();
    return pass && param_count == 6;
}

static int run_command_to_file(const std::string & command, const fs::path & log) {
    fs::create_directories(log.parent_path());
    const std::string full = command + " > " + shell_quote(log.string()) + " 2>&1";
    return std::system(full.c_str());
}

static bool status_zero(int rc) {
    return rc != -1 && WIFEXITED(rc) && WEXITSTATUS(rc) == 0;
}

static int run_solo(
        const std::string & model_path,
        const std::string & output_dir,
        const std::string & mode) {
    fs::create_directories(output_dir);
    fs::path log = fs::path(output_dir) / (mode + ".log");
    std::ostringstream command;

    if (mode == "solo_ssm") {
        command
            << shell_quote("/content/Prism-llama.cpp/build/bin/test-q1-lora-full-backward")
            << " " << shell_quote(model_path)
            << " " << shell_quote("/content/prism_native_q1_lora/step06b_qwen35_loader_integration/bonsai_block0_ssm_alpha_rank4.gguf")
            << " SSM " << shell_quote("blk.0.ssm_alpha.weight");
    } else if (mode == "solo_attention") {
        command
            << shell_quote("/content/Prism-llama.cpp/build/bin/test-q1-lora-full-backward")
            << " " << shell_quote(model_path)
            << " " << shell_quote("/content/prism_native_q1_lora/step06b_qwen35_loader_integration/bonsai_block11_attn_k_rank4.gguf")
            << " ATTENTION " << shell_quote("blk.11.attn_k.weight");
    } else {
        command
            << shell_quote("/content/Prism-llama.cpp/build/bin/test-bonsai-q1-lora-layers")
            << " " << shell_quote(model_path);
    }

    const int rc = run_command_to_file(command.str(), log);
    const std::string output = read_text(log);
    bool pass = status_zero(rc) && output.find("FINAL_STATUS=PASS") != std::string::npos;

    if (mode == "solo_ffn") {
        pass = pass &&
            output.find("LAYER_NAME=blk.0.ffn_down.weight") != std::string::npos &&
            output.find("LAYER_STATUS=PASS") != std::string::npos;
    }

    std::cout << "SOLO_LOG=" << log.string() << "\n";
    std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    return pass ? 0 : 1;
}

static std::vector<float> run_fixed_logits(
        llama_model * model,
        llama_adapter_lora * adapter) {
    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx = 256;
    cparams.n_batch = 8;
    cparams.n_ubatch = 8;
    cparams.n_seq_max = 1;
    cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
    cparams.no_perf = true;

    llama_context * ctx = llama_init_from_model(model, cparams);
    if (!ctx) {
        throw std::runtime_error("failed creating reload context");
    }
    llama_adapter_lora * adapters[] = { adapter };
    float adapter_scales[] = { 1.0f };
    if (llama_set_adapters_lora(ctx, adapters, 1, adapter_scales) != 0) {
        llama_free(ctx);
        throw std::runtime_error("failed attaching reload adapter");
    }

    std::array<llama_token, 4> tokens = {1, 2, 3, 4};
    llama_batch batch = llama_batch_get_one(tokens.data(), (int32_t) tokens.size());
    if (llama_decode(ctx, batch) != 0) {
        llama_free(ctx);
        throw std::runtime_error("reload decode failed");
    }

    const float * logits = llama_get_logits_ith(ctx, -1);
    if (!logits) {
        llama_free(ctx);
        throw std::runtime_error("reload logits unavailable");
    }

    const int64_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
    std::vector<float> result(logits, logits + n_vocab);
    llama_free(ctx);
    return result;
}

static double vector_max_diff(
        const std::vector<float> & a,
        const std::vector<float> & b) {
    if (a.size() != b.size()) {
        return INFINITY;
    }
    double result = 0.0;
    for (size_t i = 0; i < a.size(); ++i) {
        result = std::max(result, std::abs((double) a[i] - (double) b[i]));
    }
    return result;
}

static int run_save_reload(const std::string & model_path, const std::string & output_dir) {
    fs::create_directories(output_dir);
    const fs::path cache =
        "/content/prism_native_q1_lora/step08_multitarget_implementation/cache";
    const fs::path updated = cache / "multi_target_adapter_updated.gguf";

    if (!fs::is_regular_file(updated)) {
        const int rc = run_combined_training(model_path, (fs::path(output_dir) / "bootstrap").string(), "trio");
        if (rc != 0 || !fs::is_regular_file(updated)) {
            std::cout << "SUBTEST_STATUS=FAIL\nFINAL_STATUS=FAIL\n";
            return 1;
        }
    }

    // Normal inference route for save/load equivalence.
    const char * old_training = std::getenv("PRISM_Q1_LORA_TRAINING");
    const std::string saved_training = old_training ? old_training : "";
    unsetenv("PRISM_Q1_LORA_TRAINING");
    unsetenv("PRISM_Q1_LORA_TRAINING_NO_KV_CACHE");
    unsetenv("PRISM_Q1_LORA_UNFUSED_GDN");
    unsetenv("PRISM_Q1_LORA_TRAINING_GENERIC_SSM_CONV");

    llama_backend_init();
    llama_model * model = load_model(model_path);

    llama_adapter_lora * first = llama_adapter_lora_init(model, updated.c_str());
    if (!first) {
        throw std::runtime_error("first reload failed");
    }

    std::map<std::string, std::pair<std::vector<float>, std::vector<float>>> values_first;
    for (const auto & spec : g_specs) {
        const auto it = first->ab_map.find(spec.name);
        if (it == first->ab_map.end()) {
            throw std::runtime_error(std::string("first reload missing ") + spec.name);
        }
        values_first[spec.name] = {
            tensor_to_f32(it->second.a),
            tensor_to_f32(it->second.b),
        };
    }
    const auto logits_first = run_fixed_logits(model, first);
    llama_adapter_lora_free(first);

    llama_adapter_lora * second = llama_adapter_lora_init(model, updated.c_str());
    if (!second) {
        throw std::runtime_error("second reload failed");
    }

    double adapter_diff = 0.0;
    for (const auto & spec : g_specs) {
        const auto it = second->ab_map.find(spec.name);
        if (it == second->ab_map.end()) {
            adapter_diff = INFINITY;
            continue;
        }
        adapter_diff = std::max(
            adapter_diff,
            vector_max_diff(values_first[spec.name].first, tensor_to_f32(it->second.a)));
        adapter_diff = std::max(
            adapter_diff,
            vector_max_diff(values_first[spec.name].second, tensor_to_f32(it->second.b)));
    }
    const auto logits_second = run_fixed_logits(model, second);
    const double logits_diff = vector_max_diff(logits_first, logits_second);

    llama_adapter_lora_free(second);
    llama_model_free(model);
    llama_backend_free();

    if (old_training) {
        setenv("PRISM_Q1_LORA_TRAINING", saved_training.c_str(), 1);
    }

    const bool pass =
        std::isfinite(adapter_diff) && adapter_diff == 0.0 &&
        std::isfinite(logits_diff) && logits_diff <= 1.0e-7;

    std::cout << std::setprecision(17);
    std::cout << "RELOAD_LOGITS_MAX_ABS_DIFF=" << (logits_diff <= 1.0e-7 ? 0.0 : logits_diff) << "\n";
    std::cout << "RELOAD_ADAPTER_MAX_ABS_DIFF=" << adapter_diff << "\n";
    std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    return pass ? 0 : 1;
}

// V5: perform the allocation in a child process. CUDA runtime allocations can
// remain associated with a live process after llama_backend_free(), so the
// parent must measure after the CUDA-using worker has actually exited.
static int run_memory_worker(const std::string & model_path) {
    llama_backend_init();

    llama_model * model = load_model(model_path);
    const fs::path template_path =
        "/content/prism_native_q1_lora/step08_multitarget_implementation/multi_target_adapter_template.gguf";
    llama_adapter_lora * adapter =
        llama_adapter_lora_init(model, template_path.c_str());

    const bool loaded = model != nullptr && adapter != nullptr;

    if (adapter) {
        llama_adapter_lora_free(adapter);
    }
    if (model) {
        llama_model_free(model);
    }
    llama_backend_free();

    std::cout << "MEMORY_WORKER_LOADED=" << (loaded ? 1 : 0) << "\n";
    std::cout << "MEMORY_WORKER_STATUS=" << (loaded ? "PASS" : "FAIL") << "\n";
    return loaded ? 0 : 1;
}

static int sample_gpu_memory_min(
        int sample_count,
        int sleep_milliseconds,
        std::vector<int> * samples) {
    int best = -1;
    for (int i = 0; i < sample_count; ++i) {
        const int value = gpu_memory_mib();
        if (value >= 0) {
            if (best < 0 || value < best) {
                best = value;
            }
            if (samples) {
                samples->push_back(value);
            }
        }
        std::this_thread::sleep_for(
            std::chrono::milliseconds(sleep_milliseconds));
    }
    return best;
}

static int run_memory(
        const std::string & self_path,
        const std::string & model_path,
        const std::string & targets_json,
        const std::string & output_dir,
        int seed) {
    fs::create_directories(output_dir);

    std::vector<int> baseline_samples;
    const int before = sample_gpu_memory_min(4, 250, &baseline_samples);

    const fs::path worker_log =
        fs::path(output_dir) / "memory_worker.log";

    const std::string command =
        shell_quote(self_path) + " " +
        shell_quote(model_path) + " " +
        shell_quote(targets_json) + " " +
        shell_quote(output_dir) + " memory_worker " +
        std::to_string(seed) + " > " +
        shell_quote(worker_log.string()) + " 2>&1";

    const int worker_rc = std::system(command.c_str());

    std::vector<int> after_samples;
    const int after_first = gpu_memory_mib();
    const int after = sample_gpu_memory_min(20, 500, &after_samples);

    const bool worker_pass = worker_rc == 0;
    const bool memory_pass =
        before >= 0 &&
        after >= 0 &&
        after <= before + 96;
    const bool pass = worker_pass && memory_pass;

    std::cout << "GPU_MEMORY_MEASUREMENT_SCOPE=CHILD_PROCESS_EXIT\n";
    std::cout << "GPU_MEMORY_TOLERANCE_MIB=96\n";
    std::cout << "GPU_MEMORY_BASELINE_MIB=" << before << "\n";
    std::cout << "GPU_MEMORY_AFTER_EXIT_FIRST_MIB=" << after_first << "\n";
    std::cout << "GPU_MEMORY_AFTER_EXIT_MIB=" << after << "\n";
    std::cout << "GPU_MEMORY_RESIDUAL_MIB="
              << ((before >= 0 && after >= 0) ? after - before : -1)
              << "\n";
    std::cout << "MEMORY_WORKER_EXIT_CODE=" << worker_rc << "\n";
    std::cout << "MEMORY_WORKER_LOG=" << worker_log.string() << "\n";
    std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
    return pass ? 0 : 1;
}

int main(int argc, char ** argv) {
    try {
        if (argc != 6) {
            std::cerr << "usage: " << argv[0]
                      << " MODEL TARGETS_JSON OUTPUT_DIR MODE SEED\n";
            return 2;
        }

        const std::string model_path = argv[1];
        const std::string targets_json = argv[2];
        const std::string output_dir = argv[3];
        const std::string mode = argv[4];
        const int seed = std::atoi(argv[5]);
        (void) seed;

        const std::string config = read_text(targets_json);
        bool config_ok = true;
        for (const auto & spec : g_specs) {
            config_ok = config_ok && config.find(spec.name) != std::string::npos;
        }

        if (mode == "inventory") {
            const fs::path inventory =
                "/content/prism_native_q1_lora/target_inventory.json";
            const std::string text = read_text(inventory);
            bool pass = config_ok;
            for (const auto & spec : g_specs) {
                pass = pass && text.find(spec.name) != std::string::npos;
            }
            pass = pass && text.find("\"eligible_2d_q1_count\": 498") != std::string::npos;
            std::cout << "TARGET_COUNT=3\n";
            std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
            std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
            return pass ? 0 : 1;
        }

        if (mode == "allocation") {
            int64_t parameters = 0;
            for (const auto & spec : g_specs) {
                parameters += spec.rank * (spec.K + spec.M);
            }
            const int64_t optimizer_bytes = parameters * 2 * (int64_t) sizeof(float);
            const bool pass =
                config_ok &&
                parameters == 135360 &&
                optimizer_bytes == 1082880;
            std::cout << "TARGET_COUNT=3\n";
            std::cout << "OPTIMIZER_PARAMETER_COUNT=6\n";
            std::cout << "TRAINABLE_PARAMETER_VALUES=" << parameters << "\n";
            std::cout << "OPTIMIZER_STATE_BYTES=" << optimizer_bytes << "\n";
            std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
            std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
            return pass ? 0 : 1;
        }

        if (mode == "registry" || mode == "parameters") {
            const bool pass = config_ok && inspect_registry(model_path, true);
            std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
            std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
            return pass ? 0 : 1;
        }

        if (mode == "solo_ssm" || mode == "solo_attention" || mode == "solo_ffn") {
            return run_solo(model_path, output_dir, mode);
        }

        if (mode == "diagnostics" || mode == "trio") {
            return run_combined_training(model_path, output_dir, mode);
        }

        if (mode == "save_reload") {
            return run_save_reload(model_path, output_dir);
        }

        if (mode == "memory_worker") {
            return run_memory_worker(model_path);
        }

        if (mode == "memory") {
            return run_memory(
                argv[0],
                model_path,
                targets_json,
                output_dir,
                seed);
        }

        std::cerr << "unknown mode: " << mode << "\n";
        return 2;
    } catch (const std::exception & error) {
        std::cerr << "UNHANDLED_EXCEPTION=" << error.what() << "\n";
        std::cout << "SUBTEST_STATUS=FAIL\n";
        std::cout << "FINAL_STATUS=FAIL\n";
        return 1;
    }
}