// PRISM_Q1_LORA_ADAMW_TEST_V1

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

// PRISM_Q1_LORA_INTERNAL_GRAPH_ACCESS_V1
#include "ggml-impl.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <random>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>


// ============================================================
// TEST-ONLY COMPLETION OF THE PRIVATE OPTIMIZER CONTEXT
//
// This definition is pinned to Prism commit:
//
//   7529fdaaf99ffdc5ca71ace9c7409a56b27ad92f
//
// It permits this dedicated test to directly verify which tensors
// receive AdamW m/v state and optimizer-step nodes.
// ============================================================

struct ggml_opt_context {
    ggml_backend_sched_t       backend_sched        = nullptr;
    ggml_cgraph              * allocated_graph      = nullptr;
    ggml_cgraph              * allocated_graph_copy = nullptr;
    struct ggml_context      * ctx_static           = nullptr;
    struct ggml_context      * ctx_cpu              = nullptr;
    struct ggml_context      * ctx_compute          = nullptr;
    struct ggml_context      * ctx_copy             = nullptr;
    ggml_backend_buffer_t      buf_static           = nullptr;
    ggml_backend_buffer_t      buf_cpu              = nullptr;
    std::mt19937               rng;
    enum ggml_opt_loss_type    loss_type;
    enum ggml_opt_build_type   build_type;
    enum ggml_opt_build_type   build_type_alloc;

    struct ggml_tensor * inputs  = nullptr;
    struct ggml_tensor * outputs = nullptr;
    struct ggml_tensor * labels  = nullptr;

    struct ggml_tensor * loss     = nullptr;
    struct ggml_tensor * pred     = nullptr;
    struct ggml_tensor * ncorrect = nullptr;

    struct ggml_cgraph * gf      = nullptr;
    struct ggml_cgraph * gb_grad = nullptr;
    struct ggml_cgraph * gb_opt  = nullptr;

    bool static_graphs = false;
    bool eval_ready    = false;

    std::vector<struct ggml_tensor *> grad_accs;
    std::vector<struct ggml_tensor *> grad_m;
    std::vector<struct ggml_tensor *> grad_v;

    int64_t iter       = 1;
    int32_t opt_period = 1;
    int32_t opt_i      = 0;

    bool loss_per_datapoint = false;

    ggml_opt_get_optimizer_params get_opt_pars    = nullptr;
    void *                        get_opt_pars_ud = nullptr;

    struct ggml_tensor * opt_step_params = nullptr;

    enum ggml_opt_optimizer_type optimizer =
        GGML_OPT_OPTIMIZER_TYPE_ADAMW;
};


// ============================================================
// CONSTANTS
// ============================================================

static constexpr int64_t K = 256;
static constexpr int64_t M = 16;
static constexpr int64_t RANK = 4;
static constexpr int64_t BATCH = 2;

static constexpr float LORA_ALPHA = 8.0f;
static constexpr float LORA_SCALE =
    LORA_ALPHA / static_cast<float>(RANK);

static constexpr int TRAINING_STEPS = 60;


// ============================================================
// BASIC HELPERS
// ============================================================

static std::vector<float> read_f32(
        const ggml_tensor * tensor) {
    if (!tensor) {
        throw std::runtime_error(
            "Attempted to read a null F32 tensor.");
    }

    if (tensor->type != GGML_TYPE_F32) {
        throw std::runtime_error(
            std::string("Expected F32 tensor, got ")
            + ggml_type_name(tensor->type));
    }

    std::vector<float> result(
        ggml_nelements(tensor));

    ggml_backend_tensor_get(
        tensor,
        result.data(),
        0,
        ggml_nbytes(tensor));

    return result;
}


static std::vector<uint8_t> read_raw(
        const ggml_tensor * tensor) {
    if (!tensor) {
        throw std::runtime_error(
            "Attempted to read a null tensor.");
    }

    std::vector<uint8_t> result(
        ggml_nbytes(tensor));

    ggml_backend_tensor_get(
        tensor,
        result.data(),
        0,
        result.size());

    return result;
}


static double max_abs(
        const std::vector<float> & values) {
    double result = 0.0;

    for (const float value : values) {
        result = std::max(
            result,
            std::abs(
                static_cast<double>(value)));
    }

    return result;
}


static double max_abs_difference(
        const std::vector<float> & a,
        const std::vector<float> & b) {
    if (a.size() != b.size()) {
        throw std::runtime_error(
            "Vector size mismatch.");
    }

    double result = 0.0;

    for (size_t index = 0;
         index < a.size();
         ++index) {
        result = std::max(
            result,
            std::abs(
                static_cast<double>(a[index])
                - static_cast<double>(b[index])));
    }

    return result;
}


static std::string join_strings(
        std::vector<std::string> values) {
    std::sort(
        values.begin(),
        values.end());

    std::string result;

    for (size_t index = 0;
         index < values.size();
         ++index) {
        if (index > 0) {
            result += ",";
        }

        result += values[index];
    }

    return result;
}


static void require(
        const bool condition,
        const std::string & message) {
    if (!condition) {
        throw std::runtime_error(
            message);
    }
}


// ============================================================
// CPU LORA REFERENCE
//
// GGML storage:
//   A = [K, RANK]
//   B = [RANK, M]
//   X = [K, BATCH]
// ============================================================

static std::vector<float> cpu_lora(
        const std::vector<float> & x,
        const std::vector<float> & a,
        const std::vector<float> & b) {
    std::vector<float> down(
        BATCH * RANK,
        0.0f);

    std::vector<float> up(
        BATCH * M,
        0.0f);

    for (int64_t batch = 0;
         batch < BATCH;
         ++batch) {
        for (int64_t rank = 0;
             rank < RANK;
             ++rank) {
            float sum = 0.0f;

            for (int64_t k = 0;
                 k < K;
                 ++k) {
                sum +=
                    a[rank * K + k]
                    * x[batch * K + k];
            }

            down[
                batch * RANK
                + rank] = sum;
        }
    }

    for (int64_t batch = 0;
         batch < BATCH;
         ++batch) {
        for (int64_t output = 0;
             output < M;
             ++output) {
            float sum = 0.0f;

            for (int64_t rank = 0;
                 rank < RANK;
                 ++rank) {
                sum +=
                    b[output * RANK + rank]
                    * down[
                        batch * RANK
                        + rank];
            }

            up[
                batch * M
                + output] =
                    sum * LORA_SCALE;
        }
    }

    return up;
}



// PRISM_Q1_LORA_CPU_TARGET_V3
//
// Exact CPU reference for Prism Q1_0:
//
//   block size: 128 weights
//   storage:    FP16 scale + 16 sign bytes
//   bit 1:      +scale
//   bit 0:      -scale
//
// This is used only to generate deterministic training labels.
// It does not create a persistent expanded model tensor.

static std::vector<float> cpu_q1_base_output(
        const std::vector<uint8_t> & packed,
        const std::vector<float> & x) {
    constexpr int64_t qk = 128;
    constexpr size_t block_bytes = 18;

    static_assert(
        sizeof(ggml_fp16_t) == 2,
        "Unexpected ggml_fp16_t size.");

    const size_t blocks_per_row =
        static_cast<size_t>(K / qk);

    const size_t row_bytes =
        blocks_per_row * block_bytes;

    require(
        packed.size()
            == static_cast<size_t>(M) * row_bytes,
        "Unexpected packed Q1 buffer size.");

    require(
        x.size()
            == static_cast<size_t>(K * BATCH),
        "Unexpected CPU input size.");

    std::vector<float> result(
        static_cast<size_t>(M * BATCH),
        0.0f);

    for (int64_t batch = 0;
         batch < BATCH;
         ++batch) {
        for (int64_t output = 0;
             output < M;
             ++output) {
            float sum = 0.0f;

            const size_t row_offset =
                static_cast<size_t>(output)
                * row_bytes;

            for (int64_t k = 0;
                 k < K;
                 ++k) {
                const size_t block_index =
                    static_cast<size_t>(k / qk);

                const int quant_index =
                    static_cast<int>(k % qk);

                const size_t block_offset =
                    row_offset
                    + block_index * block_bytes;

                ggml_fp16_t scale_fp16;

                std::memcpy(
                    &scale_fp16,
                    packed.data() + block_offset,
                    sizeof(scale_fp16));

                const float scale =
                    GGML_FP16_TO_FP32(scale_fp16);

                const uint8_t packed_signs =
                    packed[
                        block_offset
                        + sizeof(ggml_fp16_t)
                        + static_cast<size_t>(
                            quant_index >> 3)];

                const int sign_bit =
                    (
                        packed_signs
                        >> (quant_index & 7)
                    )
                    & 1;

                const float weight =
                    sign_bit
                    ? scale
                    : -scale;

                sum +=
                    weight
                    * x[
                        static_cast<size_t>(batch * K + k)];
            }

            result[
                static_cast<size_t>(batch * M + output)
            ] = sum;
        }
    }

    return result;
}


// ============================================================
// MAIN
// ============================================================

int main() {
    try {
        ggml_log_set(
            nullptr,
            nullptr);

        ggml_backend_load_all();

        const size_t device_count =
            ggml_backend_dev_count();

        require(
            device_count > 0,
            "No GGML backend devices found.");

        std::vector<ggml_backend_dev_t> devices;
        std::vector<ggml_backend_t> backends;

        size_t cuda_index =
            static_cast<size_t>(-1);

        for (size_t index = 0;
             index < device_count;
             ++index) {
            ggml_backend_dev_t device =
                ggml_backend_dev_get(index);

            devices.push_back(device);

            const std::string device_name =
                ggml_backend_dev_name(device);

            if (
                cuda_index
                    == static_cast<size_t>(-1)
                && device_name.find("CUDA")
                    != std::string::npos
            ) {
                cuda_index = index;
            }

            ggml_backend_t backend =
                ggml_backend_dev_init(
                    device,
                    nullptr);

            require(
                backend != nullptr,
                "Failed to initialize backend: "
                + device_name);

            auto * reg =
                ggml_backend_dev_backend_reg(
                    device);

            auto set_threads =
                reinterpret_cast<
                    ggml_backend_set_n_threads_t>(
                        ggml_backend_reg_get_proc_address(
                            reg,
                            "ggml_backend_set_n_threads"));

            if (set_threads) {
                set_threads(
                    backend,
                    std::max(
                        1u,
                        std::thread::hardware_concurrency()
                            / 2));
            }

            backends.push_back(
                backend);
        }

        require(
            cuda_index
                != static_cast<size_t>(-1),
            "CUDA backend was not found.");

        ggml_backend_t cuda_backend =
            backends[cuda_index];

        std::vector<ggml_backend_t>
            scheduler_backends;

        scheduler_backends.push_back(
            cuda_backend);

        for (size_t index = 0;
             index < backends.size();
             ++index) {
            if (index != cuda_index) {
                scheduler_backends.push_back(
                    backends[index]);
            }
        }

        ggml_backend_sched_t scheduler =
            ggml_backend_sched_new(
                scheduler_backends.data(),
                nullptr,
                scheduler_backends.size(),
                GGML_DEFAULT_GRAPH_SIZE,
                false,
                true);

        require(
            scheduler != nullptr,
            "Failed to create backend scheduler.");

        std::cout
            << std::fixed
            << std::setprecision(10);

        std::cout
            << "CUDA_BACKEND="
            << ggml_backend_name(
                cuda_backend)
            << "\n";

        std::cout
            << "CUDA_DEVICE="
            << ggml_backend_dev_description(
                devices[cuda_index])
            << "\n";

        // --------------------------------------------------------
        // STATIC MODEL CONTEXT
        // --------------------------------------------------------

        ggml_init_params static_params = {
            /*.mem_size   =*/
                8 * ggml_tensor_overhead(),

            /*.mem_buffer =*/
                nullptr,

            /*.no_alloc   =*/
                true,
        };

        ggml_context * ctx_static =
            ggml_init(
                static_params);

        require(
            ctx_static != nullptr,
            "Failed to create static context.");

        ggml_tensor * x =
            ggml_new_tensor_2d(
                ctx_static,
                GGML_TYPE_F32,
                K,
                BATCH);

        ggml_set_name(
            x,
            "input_x");

        ggml_tensor * base_weight =
            ggml_new_tensor_2d(
                ctx_static,
                GGML_TYPE_Q1_0,
                K,
                M);

        ggml_set_name(
            base_weight,
            "base_q1");

        ggml_tensor * lora_a =
            ggml_new_tensor_2d(
                ctx_static,
                GGML_TYPE_F32,
                K,
                RANK);

        ggml_set_name(
            lora_a,
            "lora_a");

        ggml_set_param(
            lora_a);

        ggml_tensor * lora_b =
            ggml_new_tensor_2d(
                ctx_static,
                GGML_TYPE_F32,
                RANK,
                M);

        ggml_set_name(
            lora_b,
            "lora_b");

        ggml_set_param(
            lora_b);

        require(
            !(base_weight->flags
                & GGML_TENSOR_FLAG_PARAM),
            "Packed Q1 base was marked as a parameter.");

        require(
            !(x->flags
                & GGML_TENSOR_FLAG_PARAM),
            "Input X was marked as a parameter.");

        require(
            lora_a->flags
                & GGML_TENSOR_FLAG_PARAM,
            "LoRA A was not marked as a parameter.");

        require(
            lora_b->flags
                & GGML_TENSOR_FLAG_PARAM,
            "LoRA B was not marked as a parameter.");

        // --------------------------------------------------------
        // COMPUTE GRAPH
        // --------------------------------------------------------

        ggml_init_params compute_params = {
            /*.mem_size   =*/
                GGML_DEFAULT_GRAPH_SIZE
                    * ggml_tensor_overhead()
                + 4
                    * ggml_graph_overhead(),

            /*.mem_buffer =*/
                nullptr,

            /*.no_alloc   =*/
                true,
        };

        ggml_context * ctx_compute =
            ggml_init(
                compute_params);

        require(
            ctx_compute != nullptr,
            "Failed to create compute context.");

        ggml_tensor * base_output =
            ggml_mul_mat(
                ctx_compute,
                base_weight,
                x);

        ggml_set_name(
            base_output,
            "base_output");

        // Select the exact packed-Q1 × F32 training kernel.
        ggml_mul_mat_set_prec(
            base_output,
            GGML_PREC_F32);

        ggml_tensor * lora_down =
            ggml_mul_mat(
                ctx_compute,
                lora_a,
                x);

        ggml_set_name(
            lora_down,
            "lora_down");

        ggml_tensor * lora_up =
            ggml_mul_mat(
                ctx_compute,
                lora_b,
                lora_down);

        ggml_set_name(
            lora_up,
            "lora_up");

        ggml_tensor * lora_scaled =
            ggml_scale(
                ctx_compute,
                lora_up,
                LORA_SCALE);

        ggml_set_name(
            lora_scaled,
            "lora_scaled");

        ggml_tensor * output =
            ggml_add(
                ctx_compute,
                base_output,
                lora_scaled);

        ggml_set_name(
            output,
            "q1_lora_output");

        // --------------------------------------------------------
        // ALLOCATE STATIC MODEL TENSORS ON CUDA
        // --------------------------------------------------------

        ggml_backend_buffer_t model_buffer =
            ggml_backend_alloc_ctx_tensors(
                ctx_static,
                cuda_backend);

        require(
            model_buffer != nullptr,
            "Failed to allocate static model tensors.");

        // --------------------------------------------------------
        // DETERMINISTIC INITIALIZATION
        // --------------------------------------------------------

        std::vector<float> host_x(
            K * BATCH);

        std::vector<float> base_f32(
            K * M);

        std::vector<float> target_a(
            K * RANK);

        std::vector<float> target_b(
            RANK * M);

        std::vector<float> initial_a(
            K * RANK);

        std::vector<float> initial_b(
            RANK * M);

        for (size_t index = 0;
             index < host_x.size();
             ++index) {
            host_x[index] =
                0.70f
                    * std::sin(
                        0.031f
                        * static_cast<float>(index))
                + 0.20f
                    * std::cos(
                        0.017f
                        * static_cast<float>(index));
        }

        for (size_t index = 0;
             index < base_f32.size();
             ++index) {
            base_f32[index] =
                0.075f
                * std::sin(
                    0.013f
                    * static_cast<float>(index))
                + 0.025f
                * std::cos(
                    0.007f
                    * static_cast<float>(index));
        }

        for (size_t index = 0;
             index < target_a.size();
             ++index) {
            target_a[index] =
                0.040f
                * std::sin(
                    0.019f
                    * static_cast<float>(index + 1));

            initial_a[index] =
                0.60f
                    * target_a[index]
                + 0.002f
                    * std::cos(
                        0.011f
                        * static_cast<float>(index));
        }

        for (size_t index = 0;
             index < target_b.size();
             ++index) {
            target_b[index] =
                0.050f
                * std::cos(
                    0.071f
                    * static_cast<float>(index + 1));

            initial_b[index] =
                0.60f
                    * target_b[index]
                + 0.002f
                    * std::sin(
                        0.053f
                        * static_cast<float>(index));
        }

        std::vector<float> importance(
            K,
            1.0f);

        std::vector<uint8_t> base_q1(
            ggml_nbytes(
                base_weight));

        const size_t quantized_bytes =
            ggml_quantize_chunk(
                GGML_TYPE_Q1_0,
                base_f32.data(),
                base_q1.data(),
                0,
                M,
                K,
                importance.data());

        require(
            quantized_bytes
                == base_q1.size(),
            "Unexpected Q1 quantized byte count.");

        ggml_backend_tensor_set(
            x,
            host_x.data(),
            0,
            ggml_nbytes(x));

        ggml_backend_tensor_set(
            base_weight,
            base_q1.data(),
            0,
            base_q1.size());

        ggml_backend_tensor_set(
            lora_a,
            initial_a.data(),
            0,
            ggml_nbytes(lora_a));

        ggml_backend_tensor_set(
            lora_b,
            initial_b.data(),
            0,
            ggml_nbytes(lora_b));

        // --------------------------------------------------------
        // OPTIMIZER CONTEXT
        // --------------------------------------------------------

        ggml_opt_optimizer_params adamw =
            ggml_opt_get_default_optimizer_params(
                nullptr);

        adamw.adamw.alpha = 0.001f;
        adamw.adamw.beta1 = 0.9f;
        adamw.adamw.beta2 = 0.999f;
        adamw.adamw.eps   = 1.0e-8f;
        adamw.adamw.wd    = 0.0f;

        ggml_opt_params opt_params =
            ggml_opt_default_params(
                scheduler,
                GGML_OPT_LOSS_TYPE_MEAN_SQUARED_ERROR);

        opt_params.ctx_compute =
            ctx_compute;

        opt_params.inputs =
            x;

        opt_params.outputs =
            output;

        opt_params.opt_period =
            1;

        opt_params.optimizer =
            GGML_OPT_OPTIMIZER_TYPE_ADAMW;

        opt_params.get_opt_pars =
            ggml_opt_get_constant_optimizer_params;

        opt_params.get_opt_pars_ud =
            &adamw;

        ggml_opt_context_t opt_context =
            ggml_opt_init(
                opt_params);

        require(
            opt_context != nullptr,
            "Failed to initialize AdamW context.");

        ggml_tensor * labels =
            ggml_opt_labels(
                opt_context);

        require(
            labels != nullptr,
            "MSE labels tensor was not created.");

        ggml_tensor * loss_tensor =
            ggml_opt_loss(
                opt_context);

        require(
            loss_tensor != nullptr,
            "Loss tensor was not created.");

        // --------------------------------------------------------
        // DIRECT OPTIMIZER-STATE AUDIT
        // --------------------------------------------------------

        require(
            opt_context->gf != nullptr,
            "Forward graph is missing.");

        require(
            opt_context->gb_opt != nullptr,
            "Optimizer graph is missing.");

        require(
            opt_context->grad_m.size()
                == static_cast<size_t>(
                    opt_context->gf->n_nodes),
            "AdamW m vector has unexpected size.");

        require(
            opt_context->grad_v.size()
                == static_cast<size_t>(
                    opt_context->gf->n_nodes),
            "AdamW v vector has unexpected size.");

        std::vector<std::string> parameter_names;
        std::vector<std::string> momentum_m_parameters;
        std::vector<std::string> momentum_v_parameters;
        std::vector<std::string> optimizer_step_parameters;

        std::map<
            std::string,
            ggml_tensor *> momentum_m;

        std::map<
            std::string,
            ggml_tensor *> momentum_v;

        bool base_has_momentum = false;

        for (int index = 0;
             index < opt_context->gf->n_nodes;
             ++index) {
            ggml_tensor * node =
                opt_context->gb_opt->nodes[index];

            const std::string node_name =
                node->name;

            if (
                node->flags
                & GGML_TENSOR_FLAG_PARAM
            ) {
                parameter_names.push_back(
                    node_name);
            }

            if (
                opt_context->grad_m[index]
                != nullptr
            ) {
                momentum_m_parameters.push_back(
                    node_name);

                momentum_m[node_name] =
                    opt_context->grad_m[index];

                if (node_name == "base_q1") {
                    base_has_momentum = true;
                }
            }

            if (
                opt_context->grad_v[index]
                != nullptr
            ) {
                momentum_v_parameters.push_back(
                    node_name);

                momentum_v[node_name] =
                    opt_context->grad_v[index];

                if (node_name == "base_q1") {
                    base_has_momentum = true;
                }
            }
        }

        const std::string step_prefix =
            "AdamW step for ";

        for (int index = 0;
             index < opt_context->gb_opt->n_nodes;
             ++index) {
            const std::string node_name =
                opt_context
                    ->gb_opt
                    ->nodes[index]
                    ->name;

            if (
                node_name.rfind(
                    step_prefix,
                    0)
                == 0
            ) {
                optimizer_step_parameters.push_back(
                    node_name.substr(
                        step_prefix.size()));
            }
        }

        const std::string parameter_list =
            join_strings(
                parameter_names);

        const std::string momentum_m_list =
            join_strings(
                momentum_m_parameters);

        const std::string momentum_v_list =
            join_strings(
                momentum_v_parameters);

        const std::string optimizer_step_list =
            join_strings(
                optimizer_step_parameters);

        require(
            parameter_list
                == "lora_a,lora_b",
            "Unexpected parameter list: "
            + parameter_list);

        require(
            momentum_m_list
                == "lora_a,lora_b",
            "Unexpected AdamW m list: "
            + momentum_m_list);

        require(
            momentum_v_list
                == "lora_a,lora_b",
            "Unexpected AdamW v list: "
            + momentum_v_list);
        // PRISM_Q1_LORA_BEHAVIORAL_ADAMW_V2
        // Optimizer-step node names are not a stable
        // public contract. Execution is validated below
        // through parameter changes, moment tensors,
        // frozen base bytes, and decreasing loss.


        require(
            !base_has_momentum,
            "Packed Q1 base received AdamW state.");

        require(
            momentum_m.count("lora_a") == 1
            && momentum_m.count("lora_b") == 1
            && momentum_v.count("lora_a") == 1
            && momentum_v.count("lora_b") == 1,
            "Required LoRA moment tensors were not found.");

        std::cout
            << "PARAMETER_COUNT="
            << parameter_names.size()
            << "\n";

        std::cout
            << "PARAMETER_NAMES="
            << parameter_list
            << "\n";

        std::cout
            << "ADAMW_M_COUNT="
            << momentum_m_parameters.size()
            << "\n";

        std::cout
            << "ADAMW_M_PARAMETERS="
            << momentum_m_list
            << "\n";

        std::cout
            << "ADAMW_V_COUNT="
            << momentum_v_parameters.size()
            << "\n";

        std::cout
            << "ADAMW_V_PARAMETERS="
            << momentum_v_list
            << "\n";

        std::cout
            << "DIAGNOSTIC_NAMED_STEP_NODE_COUNT="
            << optimizer_step_parameters.size()
            << "\n";

        std::cout
            << "DIAGNOSTIC_NAMED_STEP_NODE_PARAMETERS="
            << optimizer_step_list
            << "\n";

        std::cout
            << "BASE_ADAMW_STATE="
            << (
                base_has_momentum
                ? "PRESENT"
                : "ABSENT")
            << "\n";
        // --------------------------------------------------------
        // BUILD TARGET LABELS WITHOUT READING GRAPH INTERMEDIATES
        //
        // Static optimizer graphs allocate a duplicated graph.
        // The original output intermediate can therefore have no
        // directly readable backend buffer.
        //
        // Generate the exact desired target on CPU instead:
        //
        //     target = packed-Q1 base + target LoRA
        // --------------------------------------------------------

        const std::vector<float> base_reference =
            cpu_q1_base_output(
                base_q1,
                host_x);

        const std::vector<float> target_lora =
            cpu_lora(
                host_x,
                target_a,
                target_b);

        require(
            base_reference.size()
                == target_lora.size(),
            "CPU target component size mismatch.");

        std::vector<float> target_labels(
            static_cast<size_t>(M * BATCH));

        for (size_t index = 0;
             index < target_labels.size();
             ++index) {
            target_labels[index] =
                base_reference[index]
                + target_lora[index];
        }

        std::cout
            << "TARGET_GENERATION=CPU_PACKED_Q1_PLUS_LORA\n";

        std::cout
            << "GRAPH_INTERMEDIATE_READS_FOR_TARGET=0\n";

        // --------------------------------------------------------
        // INITIAL LOSS

        // --------------------------------------------------------

        ggml_opt_result_t result =
            ggml_opt_result_init();

        ggml_opt_alloc(
            opt_context,
            false);

        ggml_backend_tensor_set(
            x,
            host_x.data(),
            0,
            ggml_nbytes(x));

        ggml_backend_tensor_set(
            labels,
            target_labels.data(),
            0,
            ggml_nbytes(labels));

        ggml_opt_eval(
            opt_context,
            result);

        double initial_loss = 0.0;

        ggml_opt_result_loss(
            result,
            &initial_loss,
            nullptr);

        require(
            std::isfinite(
                initial_loss),
            "Initial loss is not finite.");

        require(
            initial_loss > 1.0e-12,
            "Initial loss is too small for a useful test.");

        ggml_opt_result_reset(
            result);

        const std::vector<uint8_t> base_before =
            read_raw(
                base_weight);

        const std::vector<float> a_before =
            read_f32(
                lora_a);

        const std::vector<float> b_before =
            read_f32(
                lora_b);

        // --------------------------------------------------------
        // FIRST ADAMW STEP
        // --------------------------------------------------------

        ggml_opt_alloc(
            opt_context,
            true);

        ggml_backend_tensor_set(
            x,
            host_x.data(),
            0,
            ggml_nbytes(x));

        ggml_backend_tensor_set(
            labels,
            target_labels.data(),
            0,
            ggml_nbytes(labels));

        ggml_opt_eval(
            opt_context,
            nullptr);

        const std::vector<uint8_t> base_after_first =
            read_raw(
                base_weight);

        const std::vector<float> a_after_first =
            read_f32(
                lora_a);

        const std::vector<float> b_after_first =
            read_f32(
                lora_b);

        const double a_first_change =
            max_abs_difference(
                a_before,
                a_after_first);

        const double b_first_change =
            max_abs_difference(
                b_before,
                b_after_first);

        const bool base_first_unchanged =
            base_before
            == base_after_first;

        const double m_a_max =
            max_abs(
                read_f32(
                    momentum_m.at(
                        "lora_a")));

        const double v_a_max =
            max_abs(
                read_f32(
                    momentum_v.at(
                        "lora_a")));

        const double m_b_max =
            max_abs(
                read_f32(
                    momentum_m.at(
                        "lora_b")));

        const double v_b_max =
            max_abs(
                read_f32(
                    momentum_v.at(
                        "lora_b")));

        require(
            a_first_change > 1.0e-8,
            "LoRA A did not change on the first AdamW step.");

        require(
            b_first_change > 1.0e-8,
            "LoRA B did not change on the first AdamW step.");

        require(
            base_first_unchanged,
            "Packed Q1 base changed on the first step.");

        require(
            m_a_max > 0.0,
            "LoRA A AdamW m remained zero.");

        require(
            v_a_max > 0.0,
            "LoRA A AdamW v remained zero.");

        require(
            m_b_max > 0.0,
            "LoRA B AdamW m remained zero.");

        require(
            v_b_max > 0.0,
            "LoRA B AdamW v remained zero.");

        std::cout
            << "FIRST_STEP_A_MAX_CHANGE="
            << a_first_change
            << "\n";

        std::cout
            << "FIRST_STEP_B_MAX_CHANGE="
            << b_first_change
            << "\n";

        std::cout
            << "FIRST_STEP_BASE_CHANGED_BYTES="
            << (
                base_first_unchanged
                ? 0
                : 1)
            << "\n";

        std::cout
            << "LORA_A_M_MAX_ABS="
            << m_a_max
            << "\n";

        std::cout
            << "LORA_A_V_MAX_ABS="
            << v_a_max
            << "\n";

        std::cout
            << "LORA_B_M_MAX_ABS="
            << m_b_max
            << "\n";

        std::cout
            << "LORA_B_V_MAX_ABS="
            << v_b_max
            << "\n";

        // --------------------------------------------------------
        // REMAINING TRAINING STEPS
        // --------------------------------------------------------

        for (int step = 1;
             step < TRAINING_STEPS;
             ++step) {
            ggml_opt_alloc(
                opt_context,
                true);

            ggml_backend_tensor_set(
                x,
                host_x.data(),
                0,
                ggml_nbytes(x));

            ggml_backend_tensor_set(
                labels,
                target_labels.data(),
                0,
                ggml_nbytes(labels));

            ggml_opt_eval(
                opt_context,
                nullptr);
        }

        // --------------------------------------------------------
        // FINAL LOSS
        // --------------------------------------------------------

        ggml_opt_alloc(
            opt_context,
            false);

        ggml_backend_tensor_set(
            x,
            host_x.data(),
            0,
            ggml_nbytes(x));

        ggml_backend_tensor_set(
            labels,
            target_labels.data(),
            0,
            ggml_nbytes(labels));

        ggml_opt_eval(
            opt_context,
            result);

        double final_loss = 0.0;

        ggml_opt_result_loss(
            result,
            &final_loss,
            nullptr);

        const std::vector<uint8_t> base_after =
            read_raw(
                base_weight);

        const std::vector<float> a_after =
            read_f32(
                lora_a);

        const std::vector<float> b_after =
            read_f32(
                lora_b);

        const double a_total_change =
            max_abs_difference(
                a_before,
                a_after);

        const double b_total_change =
            max_abs_difference(
                b_before,
                b_after);

        const bool base_unchanged =
            base_before
            == base_after;

        const double loss_ratio =
            final_loss
            / initial_loss;

        require(
            std::isfinite(
                final_loss),
            "Final loss is not finite.");

        require(
            final_loss < initial_loss,
            "Training loss did not decrease.");

        require(
            loss_ratio < 0.95,
            "Loss reduction was too small.");

        require(
            base_unchanged,
            "Packed Q1 base bytes changed during training.");

        require(
            a_total_change > 1.0e-7,
            "LoRA A did not receive a meaningful update.");

        require(
            b_total_change > 1.0e-7,
            "LoRA B did not receive a meaningful update.");

        std::cout
            << "TRAINING_STEPS="
            << TRAINING_STEPS
            << "\n";

        std::cout
            << "INITIAL_LOSS="
            << initial_loss
            << "\n";

        std::cout
            << "FINAL_LOSS="
            << final_loss
            << "\n";

        std::cout
            << "LOSS_RATIO="
            << loss_ratio
            << "\n";

        std::cout
            << "TOTAL_A_MAX_CHANGE="
            << a_total_change
            << "\n";

        std::cout
            << "TOTAL_B_MAX_CHANGE="
            << b_total_change
            << "\n";

        std::cout
            << "BASE_CHANGED_BYTES="
            << (
                base_unchanged
                ? 0
                : 1)
            << "\n";

        std::cout
            << "PERSISTENT_EXPANDED_WEIGHT_BYTES=0\n";

        // --------------------------------------------------------
        // FINAL CONTRACT
        // --------------------------------------------------------

        std::cout
            << "CHECK_PARAMETER_FILTER=PASS\n";

        std::cout
            << "CHECK_ADAMW_M_ONLY_AB=PASS\n";

        std::cout
            << "CHECK_ADAMW_V_ONLY_AB=PASS\n";

        std::cout
            << "CHECK_ADAMW_BEHAVIORAL_UPDATE_AB=PASS\n";

        std::cout
            << "CHECK_FIRST_STEP_A_UPDATE=PASS\n";

        std::cout
            << "CHECK_FIRST_STEP_B_UPDATE=PASS\n";

        std::cout
            << "CHECK_MOMENTA_NONZERO=PASS\n";

        std::cout
            << "CHECK_BASE_FROZEN=PASS\n";

        std::cout
            << "CHECK_LOSS_DECREASE=PASS\n";

        std::cout
            << "CHECK_NO_EXPANDED_BASE=PASS\n";

        std::cout
            << "FINAL_STATUS=PASS\n";

        // --------------------------------------------------------
        // CLEANUP
        // --------------------------------------------------------

        ggml_opt_result_free(
            result);

        ggml_opt_free(
            opt_context);

        ggml_backend_buffer_free(
            model_buffer);

        ggml_free(
            ctx_compute);

        ggml_free(
            ctx_static);

        ggml_backend_sched_free(
            scheduler);

        for (ggml_backend_t backend : backends) {
            ggml_backend_free(
                backend);
        }

        return 0;

    } catch (
        const std::exception & exception
    ) {
        std::cerr
            << "FINAL_STATUS=FAIL\n";

        std::cerr
            << "ERROR="
            << exception.what()
            << "\n";

        return 1;
    }
}