// PRISM_Q1_LORA_COMPLETE_BACKWARD_V1

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

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

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>


static constexpr uint32_t TRAIN_CONTEXT = 8;


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


static bool reject_all_base_parameters(
        const ggml_tensor *,
        void *) {
    // Adapter A/B tensors were explicitly marked PARAM by the
    // Step 6 training-mode adapter loader. Reject every model
    // tensor here so no FP32 base tensor can be added.
    return false;
}


static std::vector<uint8_t> read_raw(
        const ggml_tensor * tensor) {
    require(
        tensor != nullptr,
        "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 std::vector<float> read_f32(
        const ggml_tensor * tensor) {
    require(
        tensor != nullptr,
        "Attempted to read a null F32 tensor.");

    require(
        tensor->type == GGML_TYPE_F32,
        "Expected an F32 tensor.");

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

    ggml_backend_tensor_get(
        tensor,
        result.data(),
        0,
        result.size() * sizeof(float));

    return result;
}


static double max_abs_difference(
        const std::vector<float> & before,
        const std::vector<float> & after) {
    require(
        before.size() == after.size(),
        "F32 tensor size changed.");

    double maximum = 0.0;

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

    return maximum;
}


static std::vector<llama_token> make_training_tokens(
        const llama_model * model,
        uint32_t training_context) {
    const llama_vocab * vocab =
        llama_model_get_vocab(model);

    require(
        vocab != nullptr,
        "Model vocabulary is null.");

    const std::string text =
        "Native packed one bit LoRA backward "
        "validation through the complete Qwen model graph.";

    const int required =
        -llama_tokenize(
            vocab,
            text.c_str(),
            text.size(),
            nullptr,
            0,
            true,
            true);

    require(
        required > 0,
        "Could not determine token count.");

    std::vector<llama_token> tokenized(
        static_cast<size_t>(required));

    const int written =
        llama_tokenize(
            vocab,
            text.c_str(),
            text.size(),
            tokenized.data(),
            tokenized.size(),
            true,
            true);

    require(
        written == required,
        "Tokenization failed.");

    const llama_token eos =
        llama_vocab_eos(vocab);

    require(
        eos >= 0,
        "Model has no valid EOS token.");

    while (
        tokenized.size()
        < training_context + 1
    ) {
        tokenized.push_back(eos);
    }

    tokenized.resize(
        training_context + 1);

    return tokenized;
}


struct callback_state {
    int calls = 0;
    double observed_loss = 0.0;
};


static callback_state g_callback_state;


static void training_callback(
        bool,
        ggml_opt_context_t,
        ggml_opt_dataset_t,
        ggml_opt_result_t result,
        int64_t ibatch,
        int64_t ibatch_max,
        int64_t) {
    g_callback_state.calls++;

    double loss = 0.0;

    ggml_opt_result_loss(
        result,
        &loss,
        nullptr);

    g_callback_state.observed_loss =
        loss;

    std::cout
        << "CALLBACK_INDEX="
        << ibatch
        << "\n";

    std::cout
        << "CALLBACK_MAX="
        << ibatch_max
        << "\n";

    std::cout
        << "CALLBACK_LOSS="
        << loss
        << "\n";
}


int main(
        int argc,
        char ** argv) {
    try {
        require(
            argc == 5,
            "Usage: test-q1-lora-full-backward "
            "MODEL ADAPTER MODE TARGET");

#if defined(_WIN32)
        _putenv_s(
            "PRISM_Q1_LORA_TRAINING",
            "1");
#else
        setenv(
            "PRISM_Q1_LORA_TRAINING",
            "1",
            1);
#endif

        const std::string model_path =
            argv[1];

        const std::string adapter_path =
            argv[2];

        const std::string mode =
            argv[3];

        const std::string target_name =
            argv[4];

        require(
            mode == "SSM"
                || mode == "ATTENTION",
            "MODE must be SSM or ATTENTION.");

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

        std::cout
            << "PROBE_MODE="
            << mode
            << "\n";

        std::cout
            << "PROBE_SCOPE="
            << "FULL_MODEL_CROSS_ENTROPY_BACKWARD"
            << "\n";

        std::cout
            << "TARGET_TENSOR="
            << target_name
            << "\n";

        ggml_backend_load_all();

        llama_model_params model_params =
            llama_model_default_params();

        model_params.n_gpu_layers = 999;
        model_params.use_mmap = true;

        llama_model * model =
            llama_model_load_from_file(
                model_path.c_str(),
                model_params);

        require(
            model != nullptr,
            "Could not load the Bonsai model.");

        require(
            model->arch_name() == "qwen35",
            "Loaded model is not Qwen3.5.");

        require(
            model->hparams.n_layer() == 64,
            "Expected 64 model blocks.");

        const ggml_tensor * target_base =
            model->get_tensor(
                target_name.c_str());

        require(
            target_base != nullptr,
            "Target base tensor is missing.");

        require(
            target_base->type
                == GGML_TYPE_Q1_0,
            "Target base tensor is not Q1_0.");

        require(
            !(
                target_base->flags
                & GGML_TENSOR_FLAG_PARAM
            ),
            "Packed Q1 target was already trainable.");

        if (mode == "SSM") {
            require(
                model->hparams.is_recr(0),
                "Block 0 is not recurrent.");

            require(
                model->layers[0].ssm_alpha
                    == target_base,
                "SSM target loader mapping mismatch.");
        } else {
            require(
                !model->hparams.is_recr(11),
                "Block 11 is not full attention.");

            require(
                model->layers[11].wk
                    == target_base,
                "Attention target loader mapping mismatch.");
        }

        std::cout
            << "MODEL_ARCHITECTURE="
            << model->arch_name()
            << "\n";

        std::cout
            << "MODEL_BLOCK_COUNT="
            << model->hparams.n_layer()
            << "\n";

        std::cout
            << "TARGET_BASE_TYPE=Q1_0\n";

        llama_adapter_lora * adapter =
            llama_adapter_lora_init(
                model,
                adapter_path.c_str());

        require(
            adapter != nullptr,
            "Could not load adapter GGUF.");

        const auto pair_iterator =
            adapter->ab_map.find(
                target_name);

        require(
            pair_iterator
                != adapter->ab_map.end(),
            "Adapter target pair is missing.");

        llama_adapter_lora_weight & pair =
            pair_iterator->second;

        require(
            pair.a != nullptr
                && pair.b != nullptr,
            "Adapter A/B pair is incomplete.");

        require(
            pair.a->flags
                & GGML_TENSOR_FLAG_PARAM,
            "Adapter A is not PARAM.");

        require(
            pair.b->flags
                & GGML_TENSOR_FLAG_PARAM,
            "Adapter B is not PARAM.");

        require(
            pair.a->buffer != nullptr
                && pair.b->buffer != nullptr,
            "Adapter backend buffer is null.");

        std::cout
            << "ADAPTER_PARAMETER_COUNT=2\n";

        std::cout
            << "ADAPTER_A_PARAM=1\n";

        std::cout
            << "ADAPTER_B_PARAM=1\n";

        std::cout
            << "ADAPTER_A_BUFFER="
            << ggml_backend_buffer_name(
                pair.a->buffer)
            << "\n";

        std::cout
            << "ADAPTER_B_BUFFER="
            << ggml_backend_buffer_name(
                pair.b->buffer)
            << "\n";

        llama_context_params context_params =
            llama_context_default_params();

        context_params.n_ctx =
            TRAIN_CONTEXT;

        context_params.n_batch =
            TRAIN_CONTEXT;

        context_params.n_ubatch =
            TRAIN_CONTEXT;

        context_params.n_seq_max = 1;

        context_params.flash_attn_type =
            LLAMA_FLASH_ATTN_TYPE_DISABLED;

        context_params.offload_kqv = true;
        context_params.op_offload = true;
        context_params.no_perf = true;

        llama_context * context =
            llama_init_from_model(
                model,
                context_params);

        require(
            context != nullptr,
            "Could not create training context.");

        // PRISM_Q1_LORA_DATASET_CONTEXT_FIX_V1
        //
        // Hybrid/recurrent contexts may be internally aligned to
        // a value larger than the requested context. The optimizer
        // reads context->n_ctx() tokens for each dataset sequence,
        // so the dataset must use this resolved value.
        const uint32_t actual_context =
            context->n_ctx();

        require(
            actual_context >= TRAIN_CONTEXT,
            "Resolved context is smaller than requested.");

        require(
            actual_context % context->n_batch() == 0,
            "Resolved context is not divisible by n_batch.");

        std::cout
            << "CONTEXT_REQUESTED="
            << TRAIN_CONTEXT
            << "\n";

        std::cout
            << "CONTEXT_ACTUAL="
            << actual_context
            << "\n";

        std::cout
            << "CONTEXT_BATCH="
            << context->n_batch()
            << "\n";

        std::cout
            << "CONTEXT_UBATCH="
            << context->n_ubatch()
            << "\n";

        llama_adapter_lora * adapters[] = {
            adapter,
        };

        float adapter_scales[] = {
            1.0f,
        };

        const int adapter_result =
            llama_set_adapters_lora(
                context,
                adapters,
                1,
                adapter_scales);

        require(
            adapter_result == 0,
            "Could not attach adapter to context.");

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

        const std::vector<float>
            a_before =
                read_f32(pair.a);

        const std::vector<float>
            b_before =
                read_f32(pair.b);

        ggml_opt_optimizer_params adamw =
            ggml_opt_get_default_optimizer_params(
                nullptr);

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

        llama_opt_params optimizer_params = {
            /*.n_ctx_train =*/
                actual_context,

            /*.param_filter =*/
                reject_all_base_parameters,

            /*.param_filter_ud =*/
                nullptr,

            /*.get_opt_pars =*/
                ggml_opt_get_constant_optimizer_params,

            /*.get_opt_pars_ud =*/
                &adamw,

            /*.optimizer_type =*/
                GGML_OPT_OPTIMIZER_TYPE_ADAMW,
        };

        llama_opt_init(
            context,
            model,
            optimizer_params);

        size_t model_param_count = 0;

        for (
            const auto & entry
            : llama_internal_get_tensor_map(model)
        ) {
            const ggml_tensor * tensor =
                entry.second;

            if (
                tensor != nullptr
                && (
                    tensor->flags
                    & GGML_TENSOR_FLAG_PARAM
                )
            ) {
                model_param_count++;
            }
        }

        std::cout
            << "TRAINABLE_MODEL_PARAMETER_COUNT="
            << model_param_count
            << "\n";

        require(
            model_param_count == 0,
            "A base-model tensor was marked trainable.");

        const std::vector<llama_token>
            sequence =
                make_training_tokens(
                    model,
                    actual_context);

        std::vector<llama_token> inputs(
            actual_context);

        std::vector<llama_token> labels(
            actual_context);

        for (uint32_t index = 0;
             index < actual_context;
             ++index) {
            inputs[index] =
                sequence[index];

            labels[index] =
                sequence[index + 1];

            require(
                labels[index] >= 0
                && static_cast<uint32_t>(
                    labels[index])
                    < model->vocab.n_tokens(),
                "A training label is invalid.");
        }

        ggml_opt_dataset_t dataset =
            ggml_opt_dataset_init(
                GGML_TYPE_I32,
                GGML_TYPE_I32,

                // One datapoint is one complete context sequence.
                // Shape: [actual_context, 1].
                actual_context,
                actual_context,
                1,
                1);

        std::cout
            << "DATASET_SEQUENCE_LENGTH="
            << actual_context
            << "\n";

        std::cout
            << "DATASET_SEQUENCE_COUNT=1\n";

        std::cout
            << "DATASET_SHARD_COUNT=1\n";

        require(
            dataset != nullptr,
            "Could not create training dataset.");

        ggml_tensor * dataset_data =
            ggml_opt_dataset_data(
                dataset);

        ggml_tensor * dataset_labels =
            ggml_opt_dataset_labels(
                dataset);

        require(
            dataset_data != nullptr
                && dataset_labels != nullptr,
            "Dataset tensors are null.");

        ggml_backend_tensor_set(
            dataset_data,
            inputs.data(),
            0,
            inputs.size()
                * sizeof(llama_token));

        ggml_backend_tensor_set(
            dataset_labels,
            labels.data(),
            0,
            labels.size()
                * sizeof(llama_token));

        ggml_opt_result_t result =
            ggml_opt_result_init();

        require(
            result != nullptr,
            "Could not create optimizer result.");

        g_callback_state = {};

        // PRISM_Q1_LORA_CALLBACK_COUNT_FIX_V2
        const int expected_callback_count =
            static_cast<int>(
                actual_context
                / context->n_ubatch());

        std::cout
            << "OPTIMIZER_EXPECTED_CALLBACK_COUNT="
            << expected_callback_count
            << "\n";

        std::cout
            << "PROBE_BACKWARD_BEGIN=1\n";

        llama_opt_epoch(
            context,
            dataset,
            result,
            nullptr,
            1,
            training_callback,
            nullptr);

        llama_synchronize(context);

        std::cout
            << "PROBE_BACKWARD_RETURNED=1\n";

        double loss = 0.0;

        ggml_opt_result_loss(
            result,
            &loss,
            nullptr);

        require(
            std::isfinite(loss),
            "Cross-entropy loss is not finite.");

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

        const std::vector<float>
            a_after =
                read_f32(pair.a);

        const std::vector<float>
            b_after =
                read_f32(pair.b);

        const double a_change =
            max_abs_difference(
                a_before,
                a_after);

        const double b_change =
            max_abs_difference(
                b_before,
                b_after);

        size_t base_changed_bytes = 0;

        require(
            base_before.size()
                == base_after.size(),
            "Packed base size changed.");

        for (size_t index = 0;
             index < base_before.size();
             ++index) {
            if (
                base_before[index]
                != base_after[index]
            ) {
                base_changed_bytes++;
            }
        }

        std::cout
            << "LOSS_CROSS_ENTROPY="
            << loss
            << "\n";

        std::cout
            << "CALLBACK_COUNT="
            << g_callback_state.calls
            << "\n";

        std::cout
            << "CALLBACK_EXPECTED_COUNT="
            << expected_callback_count
            << "\n";

        std::cout
            << "LORA_A_MAX_CHANGE="
            << a_change
            << "\n";

        std::cout
            << "LORA_B_MAX_CHANGE="
            << b_change
            << "\n";

        std::cout
            << "BASE_CHANGED_BYTES="
            << base_changed_bytes
            << "\n";

        require(
            g_callback_state.calls
                == expected_callback_count,
            "Optimizer callback count does not match "
            "the number of physical microbatches.");

        require(
            a_change > 1.0e-12,
            "LoRA A did not update.");

        require(
            b_change > 1.0e-12,
            "LoRA B did not update.");

        require(
            base_changed_bytes == 0,
            "Packed Q1 target changed.");

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

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

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

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

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

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

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

        if (mode == "SSM") {
            std::cout
                << "CHECK_COMPLETE_SSM_BLOCK_BACKWARD=PASS\n";
        } else {
            std::cout
                << "CHECK_COMPLETE_ATTENTION_BLOCK_BACKWARD=PASS\n";
        }

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

        ggml_opt_result_free(result);
        ggml_opt_dataset_free(dataset);

        llama_free(context);
        llama_adapter_lora_free(adapter);
        llama_model_free(model);

        return 0;

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

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

        return 1;
    }
}