/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_attn_weights.c — HDC Memory Viability Tests 1+2
 *
 * Test 1: Attention Weight on Retrieved Entries
 *   Measures what fraction of softmax mass goes to retrieved (disk) entries
 *   vs cache entries, per head, per layer, per token.
 *
 * Test 2: Retrieval Changes Attention Distribution
 *   Compares attention entropy with and without HybridMem.
 *
 * Method:
 *   Pass A: HybridMem disabled (small cache, entries lost) — baseline entropy
 *   Pass B: HybridMem enabled (small cache, entries retrievable) — entropy + mass
 *   Both passes use identical prompt (fact + filler + question).
 *
 * Output: CSV log files + summary analysis printed to stderr.
 *
 * Usage:
 *   ./examples/test_attn_weights \
 *     --model ../models/base/RWKV-X-3.6B-64k-Base.safetensors \
 *     --vocab rwkvx/data/rwkv_vocab_v20230424.txt \
 *     --filler ../docs/architecture/rwkv-x-deep-dive.md \
 *     --cache-size 8 --hm-top-k 5 --quantize \
 *     --log-file /tmp/attn_weights.csv
 */

#include "rwkvx/rwkvx.h"
#include "tensor/tensor.h"

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

/* Read entire file into malloc'd string. */
static char* read_file(const char *path) {
    FILE *fp = fopen(path, "r");
    if (!fp) return NULL;
    fseek(fp, 0, SEEK_END);
    long len = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    char *buf = (char*)malloc(len + 1);
    size_t nread = fread(buf, 1, len, fp);
    buf[nread] = '\0';
    fclose(fp);
    return buf;
}

/* Generate tokens greedily from a prefilled state. */
static char* generate(rx_model *model, rx_tokenizer *tok, rx_state *state,
                      tm_tensor *logits, int max_tokens, rx_config cfg) {
    rx_sampling_params params = rx_sampling_defaults();
    params.temperature = 0.0f;

    uint32_t *gen_tokens = (uint32_t*)calloc(max_tokens, sizeof(uint32_t));
    int n_gen = 0;
    tm_tensor *cur_logits = logits;

    for (int i = 0; i < max_tokens; i++) {
        if (!cur_logits) break;
        uint32_t token = rx_sample(cur_logits, &params, NULL, cfg.vocab_size);
        if (i > 0) tm_free(cur_logits);
        if (token == 0) { cur_logits = NULL; break; }
        gen_tokens[n_gen++] = token;
        cur_logits = rx_forward(model, token, state);
    }
    if (cur_logits && cur_logits != logits) tm_free(cur_logits);

    char *text = rx_tokenizer_decode(tok, gen_tokens, n_gen);
    free(gen_tokens);
    return text;
}

/* Recursive cleanup */
static void rmdir_recursive(const char *path) {
    char cmd[512];
    snprintf(cmd, sizeof(cmd), "rm -rf %s", path);
    system(cmd);
}

/* Parse one CSV line: token_step,layer,head,ret_mass,cache_mass,entropy */
typedef struct {
    int token_step;
    int layer;
    int head;
    float ret_mass;
    float cache_mass;
    float entropy;
} csv_row;

/* Load CSV log into array. Returns count. */
static int load_csv(const char *path, csv_row **rows_out) {
    FILE *fp = fopen(path, "r");
    if (!fp) { *rows_out = NULL; return 0; }

    int cap = 4096, count = 0;
    csv_row *rows = (csv_row*)malloc(cap * sizeof(csv_row));
    char line[256];

    /* Skip header */
    if (fgets(line, sizeof(line), fp) == NULL) {
        fclose(fp);
        free(rows);
        *rows_out = NULL;
        return 0;
    }

    while (fgets(line, sizeof(line), fp)) {
        if (count >= cap) {
            cap *= 2;
            rows = (csv_row*)realloc(rows, cap * sizeof(csv_row));
        }
        csv_row *r = &rows[count];
        if (sscanf(line, "%d,%d,%d,%f,%f,%f",
                   &r->token_step, &r->layer, &r->head,
                   &r->ret_mass, &r->cache_mass, &r->entropy) == 6) {
            count++;
        }
    }
    fclose(fp);
    *rows_out = rows;
    return count;
}

int main(int argc, char **argv) {
    const char *model_path = NULL;
    const char *vocab_path = NULL;
    const char *filler_path = NULL;
    const char *log_file = "/tmp/attn_weights.csv";
    int cache_size = 8;
    int hm_top_k = 5;
    int max_filler_chars = 2000;
    int max_tokens = 32;
    int quantize = 0;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--model") == 0 && i + 1 < argc)
            model_path = argv[++i];
        else if (strcmp(argv[i], "--vocab") == 0 && i + 1 < argc)
            vocab_path = argv[++i];
        else if (strcmp(argv[i], "--filler") == 0 && i + 1 < argc)
            filler_path = argv[++i];
        else if (strcmp(argv[i], "--cache-size") == 0 && i + 1 < argc)
            cache_size = atoi(argv[++i]);
        else if (strcmp(argv[i], "--hm-top-k") == 0 && i + 1 < argc)
            hm_top_k = atoi(argv[++i]);
        else if (strcmp(argv[i], "--max-filler") == 0 && i + 1 < argc)
            max_filler_chars = atoi(argv[++i]);
        else if (strcmp(argv[i], "--max-tokens") == 0 && i + 1 < argc)
            max_tokens = atoi(argv[++i]);
        else if (strcmp(argv[i], "--quantize") == 0)
            quantize = 1;
        else if (strcmp(argv[i], "--log-file") == 0 && i + 1 < argc)
            log_file = argv[++i];
        else if (strcmp(argv[i], "--help") == 0) {
            fprintf(stderr,
                "Usage: %s --model PATH --vocab PATH --filler PATH [options]\n"
                "  --cache-size N     KV cache slots per attn layer (default: 8)\n"
                "  --hm-top-k N       Retrieval count (default: 5)\n"
                "  --max-filler N     Max chars of filler text (default: 2000)\n"
                "  --max-tokens N     Tokens to generate (default: 32)\n"
                "  --quantize         Quantize weights to INT8\n"
                "  --log-file PATH    CSV output path (default: /tmp/attn_weights.csv)\n",
                argv[0]);
            return 0;
        }
    }

    if (!model_path || !vocab_path || !filler_path) {
        fprintf(stderr, "Error: --model, --vocab, and --filler are required\n");
        return 1;
    }

    srand(42);

    /* Load shared resources */
    fprintf(stderr, "Loading tokenizer...\n");
    rx_tokenizer *tok = rx_tokenizer_new(vocab_path);
    if (!tok) { fprintf(stderr, "Failed to load tokenizer\n"); return 1; }

    fprintf(stderr, "Loading model%s...\n", quantize ? " (INT8)" : "");
    rx_model *model = quantize ? rx_model_load_q8(model_path)
                               : rx_model_load(model_path);
    if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; }
    rx_config cfg = rx_model_config(model);

    /* Build prompt: fact + filler + question */
    const char *fact = "The secret project codename is AURORA-7 and the lead researcher is Dr. Kimura.";
    const char *question = "\n\nQuestion: What is the secret project codename and who leads it?\nAnswer:";

    char *filler_raw = read_file(filler_path);
    if (!filler_raw) {
        fprintf(stderr, "Failed to read filler file: %s\n", filler_path);
        return 1;
    }
    int filler_len = (int)strlen(filler_raw);
    if (filler_len > max_filler_chars) {
        filler_raw[max_filler_chars] = '\0';
        filler_len = max_filler_chars;
    }

    int prompt_len = (int)(strlen(fact) + 2 + filler_len + strlen(question) + 1);
    char *prompt = (char*)malloc(prompt_len);
    snprintf(prompt, prompt_len, "%s\n\n%s%s", fact, filler_raw, question);
    free(filler_raw);

    /* Tokenize */
    uint32_t *tokens = (uint32_t*)malloc(32768 * sizeof(uint32_t));
    int n_tokens = rx_tokenizer_encode(tok, prompt, tokens, 32768);

    fprintf(stderr, "\n==========================================\n");
    fprintf(stderr, "HDC Memory Viability — Tests 1+2\n");
    fprintf(stderr, "==========================================\n");
    fprintf(stderr, "Prompt tokens:  %d\n", n_tokens);
    fprintf(stderr, "Cache size:     %d (per attn layer)\n", cache_size);
    fprintf(stderr, "HM top-k:       %d\n", hm_top_k);
    fprintf(stderr, "Generate:       %d tokens (greedy)\n", max_tokens);
    fprintf(stderr, "Attn layers:    %d\n", cfg.n_attn_layer);
    fprintf(stderr, "Heads:          %d\n", cfg.n_head);
    fprintf(stderr, "\n");

    /* ============================================================
     * Pass A: Small cache, NO HybridMem (baseline entropy)
     * ============================================================ */
    char log_a_path[512], log_b_path[512];
    snprintf(log_a_path, sizeof(log_a_path), "%s.baseline", log_file);
    snprintf(log_b_path, sizeof(log_b_path), "%s", log_file);

    fprintf(stderr, "--- Pass A: cache_size=%d, NO HybridMem (baseline) ---\n", cache_size);
    {
        rx_hybridmem_config hm_cfg = rx_hybridmem_config_defaults();
        hm_cfg.store_dir = NULL;
        hm_cfg.max_cache_len = cache_size;

        rx_state *state_a = rx_state_new_with_hybridmem(model, &hm_cfg);
        if (!state_a) { fprintf(stderr, "Failed to create state A\n"); return 1; }

        /* Enable attention logging */
        FILE *log_a = fopen(log_a_path, "w");
        fprintf(log_a, "token_step,layer,head,ret_mass,cache_mass,entropy\n");
        rx_state_set_attn_log(state_a, log_a);

        struct timespec t0, t1;
        clock_gettime(CLOCK_MONOTONIC, &t0);

        /* Prefill + generate */
        tm_tensor *logits_a = NULL;
        for (int i = 0; i < n_tokens; i++) {
            if (logits_a) tm_free(logits_a);
            logits_a = rx_forward(model, tokens[i], state_a);
        }

        clock_gettime(CLOCK_MONOTONIC, &t1);
        double ms = (t1.tv_sec - t0.tv_sec) * 1000.0 + (t1.tv_nsec - t0.tv_nsec) / 1e6;
        fprintf(stderr, "Prefill: %d tokens in %.1fms (%.1f tok/s)\n",
                n_tokens, ms, n_tokens * 1000.0 / ms);

        char *output_a = generate(model, tok, state_a, logits_a, max_tokens, cfg);
        tm_free(logits_a);

        rx_state_set_attn_log(state_a, NULL);
        fclose(log_a);
        rx_state_free(state_a);

        fprintf(stderr, "Output A: %s\n\n", output_a ? output_a : "(null)");
        free(output_a);
    }

    /* ============================================================
     * Pass B: Small cache, WITH HybridMem (entropy + mass)
     * ============================================================ */
    const char *hm_dir = "/tmp/hm_attn_test";
    rmdir_recursive(hm_dir);

    fprintf(stderr, "--- Pass B: cache_size=%d, WITH HybridMem (top_k=%d) ---\n",
            cache_size, hm_top_k);
    {
        rx_hybridmem_config hm_cfg = rx_hybridmem_config_defaults();
        hm_cfg.store_dir = hm_dir;
        hm_cfg.max_cache_len = cache_size;
        hm_cfg.top_k = hm_top_k;

        rx_state *state_b = rx_state_new_with_hybridmem(model, &hm_cfg);
        if (!state_b) { fprintf(stderr, "Failed to create state B\n"); return 1; }

        FILE *log_b = fopen(log_b_path, "w");
        fprintf(log_b, "token_step,layer,head,ret_mass,cache_mass,entropy\n");
        rx_state_set_attn_log(state_b, log_b);

        struct timespec t0, t1;
        clock_gettime(CLOCK_MONOTONIC, &t0);

        tm_tensor *logits_b = NULL;
        for (int i = 0; i < n_tokens; i++) {
            if (logits_b) tm_free(logits_b);
            logits_b = rx_forward(model, tokens[i], state_b);
        }

        clock_gettime(CLOCK_MONOTONIC, &t1);
        double ms = (t1.tv_sec - t0.tv_sec) * 1000.0 + (t1.tv_nsec - t0.tv_nsec) / 1e6;
        fprintf(stderr, "Prefill: %d tokens in %.1fms (%.1f tok/s)\n",
                n_tokens, ms, n_tokens * 1000.0 / ms);

        char *output_b = generate(model, tok, state_b, logits_b, max_tokens, cfg);
        tm_free(logits_b);

        rx_state_set_attn_log(state_b, NULL);
        fclose(log_b);
        rx_state_flush(state_b);
        rx_state_free(state_b);

        fprintf(stderr, "Output B: %s\n\n", output_b ? output_b : "(null)");
        free(output_b);
    }

    rmdir_recursive(hm_dir);

    /* ============================================================
     * Analysis
     * ============================================================ */
    fprintf(stderr, "\n==========================================\n");
    fprintf(stderr, "ANALYSIS\n");
    fprintf(stderr, "==========================================\n\n");

    /* Load both CSV logs */
    csv_row *rows_a = NULL, *rows_b = NULL;
    int count_a = load_csv(log_a_path, &rows_a);
    int count_b = load_csv(log_b_path, &rows_b);

    fprintf(stderr, "Pass A log: %d rows (%s)\n", count_a, log_a_path);
    fprintf(stderr, "Pass B log: %d rows (%s)\n\n", count_b, log_b_path);

    /* ---- Test 1: Attention weight on retrieved entries ---- */
    fprintf(stderr, "--- Test 1: Attention Weight on Retrieved Entries ---\n");
    {
        int heads_gt1 = 0, heads_gt5 = 0, total_heads = 0;
        float total_ret_mass = 0.0f;
        int generation_rows = 0;

        for (int i = 0; i < count_b; i++) {
            total_heads++;
            if (rows_b[i].ret_mass > 0.01f) heads_gt1++;
            if (rows_b[i].ret_mass > 0.05f) heads_gt5++;

            /* Count generation tokens (step > n_tokens - 1) */
            if (rows_b[i].token_step >= n_tokens) {
                total_ret_mass += rows_b[i].ret_mass;
                generation_rows++;
            }
        }

        float pct_gt1 = total_heads > 0 ? 100.0f * heads_gt1 / total_heads : 0;
        float pct_gt5 = total_heads > 0 ? 100.0f * heads_gt5 / total_heads : 0;
        float avg_ret_mass_gen = generation_rows > 0
            ? total_ret_mass / generation_rows : 0;

        fprintf(stderr, "  Heads with >1%% retrieved mass: %d/%d (%.1f%%)\n",
                heads_gt1, total_heads, pct_gt1);
        fprintf(stderr, "  Heads with >5%% retrieved mass: %d/%d (%.1f%%)\n",
                heads_gt5, total_heads, pct_gt5);
        fprintf(stderr, "  Avg retrieved mass (generation tokens): %.4f\n",
                avg_ret_mass_gen);

        /* Per-layer summary */
        fprintf(stderr, "\n  Per-layer breakdown (generation tokens only):\n");
        for (int l = 0; l < cfg.n_attn_layer; l++) {
            int layer_heads = 0, layer_gt1 = 0, layer_gt5 = 0;
            float layer_mass = 0.0f;
            for (int i = 0; i < count_b; i++) {
                if (rows_b[i].layer == l && rows_b[i].token_step >= n_tokens) {
                    layer_heads++;
                    layer_mass += rows_b[i].ret_mass;
                    if (rows_b[i].ret_mass > 0.01f) layer_gt1++;
                    if (rows_b[i].ret_mass > 0.05f) layer_gt5++;
                }
            }
            float avg = layer_heads > 0 ? layer_mass / layer_heads : 0;
            fprintf(stderr, "    Layer %d: avg_ret_mass=%.4f, >1%%=%d/%d, >5%%=%d/%d\n",
                    l, avg, layer_gt1, layer_heads, layer_gt5, layer_heads);
        }

        fprintf(stderr, "\n  VERDICT: ");
        if (pct_gt5 >= 20.0f) {
            fprintf(stderr, "PASS (%.1f%% of heads allocate >5%% to retrieved)\n", pct_gt5);
        } else if (pct_gt1 >= 5.0f) {
            fprintf(stderr, "MARGINAL (%.1f%% >1%%, %.1f%% >5%% — needs investigation)\n",
                    pct_gt1, pct_gt5);
        } else {
            fprintf(stderr, "FAIL (<1%% attention weight on retrieved entries)\n");
        }
    }

    /* ---- Test 2: Entropy shift ---- */
    fprintf(stderr, "\n--- Test 2: Retrieval Changes Attention Distribution ---\n");
    {
        /* Build per-(step, layer, head) entropy maps for both passes.
         * Simple approach: iterate rows, match by (step, layer, head). */
        int n_shifts = 0, shifts_gt01 = 0;
        float total_delta = 0.0f;

        /* For each row in B, find matching row in A by (step, layer, head) */
        for (int i = 0; i < count_b; i++) {
            for (int j = 0; j < count_a; j++) {
                if (rows_a[j].token_step == rows_b[i].token_step &&
                    rows_a[j].layer == rows_b[i].layer &&
                    rows_a[j].head == rows_b[i].head) {
                    float delta = rows_b[i].entropy - rows_a[j].entropy;
                    total_delta += delta;
                    n_shifts++;
                    if (fabsf(delta) > 0.1f) shifts_gt01++;
                    break;
                }
            }
        }

        float avg_delta = n_shifts > 0 ? total_delta / n_shifts : 0;
        float pct_shifted = n_shifts > 0 ? 100.0f * shifts_gt01 / n_shifts : 0;

        fprintf(stderr, "  Matched (step,layer,head) pairs: %d\n", n_shifts);
        fprintf(stderr, "  Avg entropy shift (B - A): %.4f nats\n", avg_delta);
        fprintf(stderr, "  Heads with |delta_H| > 0.1 nats: %d/%d (%.1f%%)\n",
                shifts_gt01, n_shifts, pct_shifted);

        fprintf(stderr, "\n  VERDICT: ");
        if (pct_shifted >= 25.0f) {
            fprintf(stderr, "PASS (%.1f%% of heads show entropy shift >0.1 nats)\n",
                    pct_shifted);
        } else if (pct_shifted >= 10.0f) {
            fprintf(stderr, "MARGINAL (%.1f%% shifted — some effect but weak)\n",
                    pct_shifted);
        } else {
            fprintf(stderr, "FAIL (retrieval doesn't measurably affect attention)\n");
        }
    }

    /* Cleanup */
    free(rows_a);
    free(rows_b);
    free(prompt);
    free(tokens);
    rx_model_free(model);
    rx_tokenizer_free(tok);

    fprintf(stderr, "\n==========================================\n");
    fprintf(stderr, "CSV logs saved to:\n");
    fprintf(stderr, "  Baseline: %s\n", log_a_path);
    fprintf(stderr, "  HybridMem: %s\n", log_b_path);
    fprintf(stderr, "==========================================\n");

    return 0;
}