// PRISM_Q1_LORA_LOADER_BLOCK_TEST_V2

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

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

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


static constexpr int64_t LORA_RANK = 4;
static constexpr float LORA_ALPHA = 8.0f;

static constexpr const char * SSM_TARGET =
    "blk.0.ssm_alpha.weight";

static constexpr const char * ATTN_TARGET =
    "blk.11.attn_k.weight";


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


static std::string lowercase(
        std::string value) {
    std::transform(
        value.begin(),
        value.end(),
        value.begin(),
        [] (unsigned char character) {
            return static_cast<char>(
                std::tolower(character));
        });

    return value;
}


static float deterministic_signed(
        size_t index,
        uint32_t salt) {
    uint32_t value =
        static_cast<uint32_t>(index)
        ^ salt
        ^ 0x9e3779b9u;

    value ^= value >> 16;
    value *= 0x7feb352du;
    value ^= value >> 15;
    value *= 0x846ca68bu;
    value ^= value >> 16;

    const float unit =
        static_cast<float>(
            value & 0x00ffffffu)
        / 16777215.0f;

    return 2.0f * unit - 1.0f;
}


struct base_shape {
    int64_t k = 0;
    int64_t m = 0;
    ggml_type type = GGML_TYPE_COUNT;
};


static base_shape read_base_shape(
        const std::string & model_path,
        const std::string & tensor_name) {
    ggml_context * metadata_context =
        nullptr;

    gguf_init_params params = {
        /*.no_alloc =*/ true,
        /*.ctx      =*/ &metadata_context,
    };

    gguf_context * gguf =
        gguf_init_from_file(
            model_path.c_str(),
            params);

    require(
        gguf != nullptr,
        "Could not open model GGUF metadata.");

    require(
        metadata_context != nullptr,
        "GGUF metadata context is null.");

    const int64_t tensor_id =
        gguf_find_tensor(
            gguf,
            tensor_name.c_str());

    require(
        tensor_id >= 0,
        "Target tensor does not exist: "
            + tensor_name);

    ggml_tensor * tensor =
        ggml_get_tensor(
            metadata_context,
            tensor_name.c_str());

    require(
        tensor != nullptr,
        "Target metadata tensor is null.");

    require(
        ggml_n_dims(tensor) == 2,
        "Target tensor is not two-dimensional.");

    base_shape shape = {
        tensor->ne[0],
        tensor->ne[1],
        gguf_get_tensor_type(
            gguf,
            tensor_id),
    };

    gguf_free(gguf);
    ggml_free(metadata_context);

    return shape;
}


static void generate_adapter(
        const std::string & model_path,
        const std::string & adapter_path,
        const std::string & tensor_name,
        uint32_t salt) {
    const base_shape shape =
        read_base_shape(
            model_path,
            tensor_name);

    require(
        shape.type == GGML_TYPE_Q1_0,
        "Target is not Q1_0: "
            + tensor_name);

    require(
        shape.k > 0
            && shape.m > 0
            && shape.k % 128 == 0,
        "Target has invalid Q1 dimensions.");

    const size_t a_elements =
        static_cast<size_t>(
            shape.k * LORA_RANK);

    const size_t b_elements =
        static_cast<size_t>(
            LORA_RANK * shape.m);

    const size_t tensor_bytes =
        (
            a_elements
            + b_elements
        )
        * sizeof(float);

    ggml_init_params tensor_params = {
        /*.mem_size   =*/
            tensor_bytes
            + 32 * ggml_tensor_overhead()
            + 1024 * 1024,

        /*.mem_buffer =*/
            nullptr,

        /*.no_alloc   =*/
            false,
    };

    ggml_context * tensor_context =
        ggml_init(tensor_params);

    require(
        tensor_context != nullptr,
        "Could not create adapter tensor context.");

    ggml_tensor * tensor_a =
        ggml_new_tensor_2d(
            tensor_context,
            GGML_TYPE_F32,
            shape.k,
            LORA_RANK);

    ggml_tensor * tensor_b =
        ggml_new_tensor_2d(
            tensor_context,
            GGML_TYPE_F32,
            LORA_RANK,
            shape.m);

    const std::string name_a =
        tensor_name + ".lora_a";

    const std::string name_b =
        tensor_name + ".lora_b";

    ggml_set_name(
        tensor_a,
        name_a.c_str());

    ggml_set_name(
        tensor_b,
        name_b.c_str());

    float * data_a =
        static_cast<float *>(
            tensor_a->data);

    float * data_b =
        static_cast<float *>(
            tensor_b->data);

    require(
        data_a != nullptr
            && data_b != nullptr,
        "Adapter tensor data is null.");

    const float a_scale =
        0.050f
        / std::sqrt(
            static_cast<float>(shape.k));

    const float b_scale =
        0.050f
        / std::sqrt(
            static_cast<float>(LORA_RANK));

    for (size_t index = 0;
         index < a_elements;
         ++index) {
        data_a[index] =
            a_scale
            * deterministic_signed(
                index,
                salt);
    }

    for (size_t index = 0;
         index < b_elements;
         ++index) {
        data_b[index] =
            b_scale
            * deterministic_signed(
                index,
                salt ^ 0x68bc21ebu);
    }

    gguf_context * adapter_gguf =
        gguf_init_empty();

    require(
        adapter_gguf != nullptr,
        "Could not create adapter GGUF context.");

    gguf_set_val_str(
        adapter_gguf,
        "general.type",
        "adapter");

    gguf_set_val_str(
        adapter_gguf,
        "general.architecture",
        "qwen35");

    gguf_set_val_str(
        adapter_gguf,
        "general.name",
        "Prism native Q1 loader block test");

    gguf_set_val_str(
        adapter_gguf,
        "adapter.type",
        "lora");

    gguf_set_val_f32(
        adapter_gguf,
        "adapter.lora.alpha",
        LORA_ALPHA);

    gguf_add_tensor(
        adapter_gguf,
        tensor_a);

    gguf_add_tensor(
        adapter_gguf,
        tensor_b);

    const bool written =
        gguf_write_to_file(
            adapter_gguf,
            adapter_path.c_str(),
            false);

    require(
        written,
        "Could not write adapter GGUF.");

    gguf_free(adapter_gguf);
    ggml_free(tensor_context);

    require(
        std::filesystem::exists(
            adapter_path),
        "Adapter GGUF was not created.");

    std::cout
        << "ADAPTER_GENERATED="
        << adapter_path
        << "\n";

    std::cout
        << "ADAPTER_TARGET="
        << tensor_name
        << "\n";

    std::cout
        << "ADAPTER_K="
        << shape.k
        << "\n";

    std::cout
        << "ADAPTER_M="
        << shape.m
        << "\n";
}


static void validate_adapter(
        const llama_model * model,
        const llama_adapter_lora * adapter,
        const std::string & target_name,
        const std::string & label) {
    require(
        adapter != nullptr,
        label + " adapter is null.");

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

    require(
        iterator
            != adapter->ab_map.end(),
        label + " adapter target is missing.");

    const llama_adapter_lora_weight & pair =
        iterator->second;

    require(
        pair.a != nullptr
            && pair.b != nullptr,
        label + " LoRA pair is incomplete.");

    require(
        pair.a->flags
            & GGML_TENSOR_FLAG_PARAM,
        label + " LoRA A is not a parameter.");

    require(
        pair.b->flags
            & GGML_TENSOR_FLAG_PARAM,
        label + " LoRA B is not a parameter.");

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

    require(
        base != nullptr,
        label + " base tensor is missing.");

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

    require(
        !(
            base->flags
            & GGML_TENSOR_FLAG_PARAM
        ),
        label + " base tensor became trainable.");

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

    const std::string a_buffer =
        ggml_backend_buffer_name(
            pair.a->buffer);

    const std::string b_buffer =
        ggml_backend_buffer_name(
            pair.b->buffer);

    require(
        lowercase(a_buffer).find("cuda")
            != std::string::npos,
        label + " LoRA A is not on CUDA: "
            + a_buffer);

    require(
        lowercase(b_buffer).find("cuda")
            != std::string::npos,
        label + " LoRA B is not on CUDA: "
            + b_buffer);

    std::cout
        << label
        << "_ADAPTER_PARAMETER_COUNT=2\n";

    std::cout
        << label
        << "_LORA_A_PARAM=1\n";

    std::cout
        << label
        << "_LORA_B_PARAM=1\n";

    std::cout
        << label
        << "_BASE_PARAM=0\n";

    std::cout
        << label
        << "_BASE_TYPE=Q1_0\n";

    std::cout
        << label
        << "_LORA_A_BUFFER="
        << a_buffer
        << "\n";

    std::cout
        << label
        << "_LORA_B_BUFFER="
        << b_buffer
        << "\n";
}


static std::vector<llama_token> tokenize(
        const llama_model * model,
        const std::string & text) {
    const llama_vocab * vocab =
        llama_model_get_vocab(model);

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

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

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

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

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

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

    return tokens;
}


static std::vector<float> evaluate_logits(
        llama_model * model,
        llama_adapter_lora * adapter,
        const std::vector<llama_token> & tokens) {
    llama_context_params context_params =
        llama_context_default_params();

    context_params.n_ctx = 32;
    context_params.n_batch = 16;
    context_params.n_ubatch = 16;
    context_params.flash_attn_type =
        LLAMA_FLASH_ATTN_TYPE_DISABLED;
    context_params.no_perf = true;

    llama_context * context =
        llama_init_from_model(
            model,
            context_params);

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

    if (adapter != nullptr) {
        llama_adapter_lora * adapters[] = {
            adapter,
        };

        float scales[] = {
            1.0f,
        };

        const int set_result =
            llama_set_adapters_lora(
                context,
                adapters,
                1,
                scales);

        require(
            set_result == 0,
            "Could not attach LoRA adapter.");
    }

    llama_batch batch =
        llama_batch_get_one(
            const_cast<llama_token *>(
                tokens.data()),
            static_cast<int32_t>(
                tokens.size()));

    const int decode_result =
        llama_decode(
            context,
            batch);

    require(
        decode_result == 0,
        "llama_decode failed: "
            + std::to_string(decode_result));

    float * logits =
        llama_get_logits_ith(
            context,
            -1);

    require(
        logits != nullptr,
        "Final logits are null.");

    const int64_t n_vocab =
        model->vocab.n_tokens();

    require(
        n_vocab > 0,
        "Invalid vocabulary size.");

    std::vector<float> result(
        logits,
        logits + n_vocab);

    llama_free(context);

    return result;
}


struct difference_metrics {
    double max_abs = 0.0;
    double mean_abs = 0.0;
    double l2 = 0.0;
    double checksum = 0.0;
};


static difference_metrics compare_logits(
        const std::vector<float> & baseline,
        const std::vector<float> & adapted) {
    require(
        baseline.size() == adapted.size(),
        "Logit vector size mismatch.");

    difference_metrics metrics;

    for (size_t index = 0;
         index < baseline.size();
         ++index) {
        require(
            std::isfinite(baseline[index])
                && std::isfinite(adapted[index]),
            "Non-finite model logit.");

        const double difference =
            static_cast<double>(
                adapted[index])
            - static_cast<double>(
                baseline[index]);

        const double absolute =
            std::abs(difference);

        metrics.max_abs =
            std::max(
                metrics.max_abs,
                absolute);

        metrics.mean_abs += absolute;
        metrics.l2 += difference * difference;

        metrics.checksum +=
            difference
            * static_cast<double>(
                (index % 1009) + 1);
    }

    metrics.mean_abs /=
        static_cast<double>(
            baseline.size());

    metrics.l2 =
        std::sqrt(metrics.l2);

    return metrics;
}


int main(
        int argc,
        char ** argv) {
    try {
        require(
            argc == 4,
            "Usage: test-q1-lora-loader-blocks "
            "MODEL SSM_ADAPTER ATTN_ADAPTER");

#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 ssm_adapter_path =
            argv[2];

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

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

        generate_adapter(
            model_path,
            ssm_adapter_path,
            SSM_TARGET,
            0x10293847u);

        generate_adapter(
            model_path,
            attn_adapter_path,
            ATTN_TARGET,
            0xabcdef12u);

        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 Bonsai model.");

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

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

        require(
            model->hparams.is_recr(0),
            "Block 0 is not recurrent.");

        require(
            !model->hparams.is_recr(11),
            "Block 11 is not full attention.");

        require(
            model->layers[0].ssm_alpha
                == model->get_tensor(
                    SSM_TARGET),
            "Block 0 SSM loader mapping mismatch.");

        require(
            model->layers[11].wk
                == model->get_tensor(
                    ATTN_TARGET),
            "Block 11 attention-K loader mapping mismatch.");

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

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

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

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

        std::cout
            << "BLOCK_SSM_TARGET="
            << SSM_TARGET
            << "\n";

        std::cout
            << "BLOCK_ATTN_INDEX=11\n";

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

        std::cout
            << "BLOCK_ATTN_TARGET="
            << ATTN_TARGET
            << "\n";

        llama_adapter_lora * ssm_adapter =
            llama_adapter_lora_init(
                model,
                ssm_adapter_path.c_str());

        llama_adapter_lora * attn_adapter =
            llama_adapter_lora_init(
                model,
                attn_adapter_path.c_str());

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

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

        validate_adapter(
            model,
            ssm_adapter,
            SSM_TARGET,
            "SSM");

        validate_adapter(
            model,
            attn_adapter,
            ATTN_TARGET,
            "ATTN");

        const std::vector<llama_token> tokens =
            tokenize(
                model,
                "One bit LoRA test.");

        std::cout
            << "MODEL_PROMPT_TOKEN_COUNT="
            << tokens.size()
            << "\n";

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

        const std::vector<float> base_logits =
            evaluate_logits(
                model,
                nullptr,
                tokens);

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

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

        const std::vector<float> ssm_logits =
            evaluate_logits(
                model,
                ssm_adapter,
                tokens);

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

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

        const std::vector<float> attn_logits =
            evaluate_logits(
                model,
                attn_adapter,
                tokens);

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

        const difference_metrics ssm_difference =
            compare_logits(
                base_logits,
                ssm_logits);

        const difference_metrics attn_difference =
            compare_logits(
                base_logits,
                attn_logits);

        std::cout
            << "SSM_LOGITS_MAX_ABS_DIFF="
            << ssm_difference.max_abs
            << "\n";

        std::cout
            << "SSM_LOGITS_MEAN_ABS_DIFF="
            << ssm_difference.mean_abs
            << "\n";

        std::cout
            << "SSM_LOGITS_L2_DIFF="
            << ssm_difference.l2
            << "\n";

        std::cout
            << "SSM_LOGITS_DIFF_CHECKSUM="
            << ssm_difference.checksum
            << "\n";

        std::cout
            << "ATTN_LOGITS_MAX_ABS_DIFF="
            << attn_difference.max_abs
            << "\n";

        std::cout
            << "ATTN_LOGITS_MEAN_ABS_DIFF="
            << attn_difference.mean_abs
            << "\n";

        std::cout
            << "ATTN_LOGITS_L2_DIFF="
            << attn_difference.l2
            << "\n";

        std::cout
            << "ATTN_LOGITS_DIFF_CHECKSUM="
            << attn_difference.checksum
            << "\n";

        require(
            ssm_difference.max_abs > 1.0e-8,
            "SSM adapter did not change logits.");

        require(
            ssm_difference.l2 > 1.0e-8,
            "SSM logit L2 difference is zero.");

        require(
            attn_difference.max_abs > 1.0e-8,
            "Attention adapter did not change logits.");

        require(
            attn_difference.l2 > 1.0e-8,
            "Attention logit L2 difference is zero.");

        require(
            std::abs(
                ssm_difference.checksum)
                > 1.0e-10,
            "SSM difference checksum is zero.");

        require(
            std::abs(
                attn_difference.checksum)
                > 1.0e-10,
            "Attention difference checksum is zero.");

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

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

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

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

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

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

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

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

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

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

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

        llama_adapter_lora_free(
            ssm_adapter);

        llama_adapter_lora_free(
            attn_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;
    }
}