#ifndef LAL_CONCEPT_GEN_H
#define LAL_CONCEPT_GEN_H
#include <omp.h>
/*
 * lal_concept_gen.h — C3 端到端闭环: 训练形成概念/关系 → 抽取概念图 → 算法生成.
 *
 * 设计(路线2/阶段C, 单一闭环):
 *   文本语料 ─► CE + C3 训练 ─► wte 概念几何
 *                             │
 *     [抽取] build_concept_graph(): 离线扫全 vocab, 每 token(=概念)取 wte 余弦
 *            top-K 邻居, 存概念图文件.
 *                             │
 *     [生成] graph_guided_generate(): 加载概念图, 对 prompt 命中概念, 先图遍历得
 *            概念骨架, 再用自回归把骨架填成连贯句(图对每步 logits 做约束/偏置).
 *
 * 三个能力:
 *   1. concept_chain_loss(): 训练端 — 全 vocab 概念链辅助 CE, 与 CE 联合优化,
 *      把"概念推导"压进 wte.
 *   2. build_concept_graph() / load_concept_graph(): 抽取端 — 离线概念图.
 *   3. graph_guided_generate(): 生成端 — 图遍历骨架 + 自回归填充(闭环终点).
 *   4. lal_concept_probe(): 探针 — 概念结构可视化(验证用, 不参与主路径).
 *
 * 零关系知识: 概念关联完全由 wte 余弦从模型自身浮现, 不硬编码任何关系类型.
 * 已删除路径: --reason 模板引擎 / --no-reason 纯自回归 / sidecar / 实时 wte 余弦 bias.
 */

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

#ifndef LAL_CONCEPT_GEN_IMPLEMENTATION
/* declarations only */
#endif

/* ===================== 概念图结构 (离线抽取) ===================== */
/* 概念图文件格式(binary, little-endian):
 *   [magic : int32 = 0x43473321]  "C3!" 标识
 *   [version : int32 = 1]
 *   [vocab_size : int32]
 *   [n_embd : int32]
 *   [K : int32]                      每概念邻居数
 *   对 i in [0, vocab_size):
 *     [i : int32]
 *     [neighbor[K] : int32]          邻居 token id(按相似度降序, -1 表示无)
 *     [sim[K] : float32]             对应余弦相似度
 * 整文件可被 load_concept_graph 直接 mmap 式读入.
 */
#define CONCEPT_GRAPH_MAGIC 0x43473321  /* 'C','3','!',0x21 */
#define CG_VERSION 1
#define CG_K 5                          /* 每概念默认邻居数 */
#define CG_MAXK_LOCAL 16                /* 栈上 top-K 缓冲上限, 防止巨大 K */

typedef struct ConceptGraph {
    int vocab_size;
    int n_embd;
    int K;
    int *neighbor;   /* [vocab_size * K] */
    float *sim;      /* [vocab_size * K] */
} ConceptGraph;

/* 底层: 直接把全 vocab 的 wte 余弦 top-K 邻居算进给定的 ConceptGraph(不存盘).
 * 主线推理端用这个就地构建概念图 —— 零文件、零参数. */
static int build_concept_graph_inplace(Model *m, ConceptGraph *g, int K) {
    if (K <= 0) K = CG_K;
    int V = m->cfg.vocab_size;
    int n_embd = m->cfg.n_embd;
    if (n_embd > 4096) { fprintf(stderr, "[graph] n_embd>4096 unsupported\n"); return 1; }

    g->vocab_size = V; g->n_embd = n_embd; g->K = K;
    g->neighbor = (int *)calloc((size_t)V * K, sizeof(int));
    g->sim = (float *)calloc((size_t)V * K, sizeof(float));

    /* === 优化版: 预归一化 + OpenMP 并行 + 向量化点积 ===
     * 原版: O(V^2 * n_embd), 每次 cosine_sim 重复算 norm, 单线程
     * 优化:
     *   1. 预归一化 wte -> norm_wte (cosine_sim 变成纯点积)
     *   2. 外层循环 OpenMP 并行
     *   3. 内层点积手动展开 (8x unroll), 便于编译器向量化
     * 预期加速: 8-16x (8核并行 + 向量化)
     */

    /* Step 1: 预归一化 wte */
    float *norm_wte = (float *)malloc((size_t)V * n_embd * sizeof(float));
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < V; i++) {
        const float *ci = m->wte + (size_t)i * n_embd;
        float *ni = norm_wte + (size_t)i * n_embd;
        float norm = 0.0f;
        for (int kk = 0; kk < n_embd; kk++) norm += ci[kk] * ci[kk];
        norm = sqrtf(norm) + 1e-8f;
        float inv = 1.0f / norm;
        for (int kk = 0; kk < n_embd; kk++) ni[kk] = ci[kk] * inv;
    }

    /* Step 2: 并行计算每个概念的 top-K 邻居 */
    int KK = (K < CG_MAXK_LOCAL) ? K : CG_MAXK_LOCAL;
    #pragma omp parallel for schedule(dynamic, 64)
    for (int i = 0; i < V; i++) {
        const float *ci = norm_wte + (size_t)i * n_embd;
        int best[CG_MAXK_LOCAL]; float bsim[CG_MAXK_LOCAL];
        for (int k = 0; k < KK; k++) { best[k] = -1; bsim[k] = -2.0f; }

        for (int j = 0; j < V; j++) {
            if (j == i) continue;
            const float *cj = norm_wte + (size_t)j * n_embd;
            /* 点积 (8x unroll, 编译器自动 SIMD) */
            float s = 0.0f;
            int kk;
            for (kk = 0; kk + 7 < n_embd; kk += 8)
                s += ci[kk]*cj[kk] + ci[kk+1]*cj[kk+1] + ci[kk+2]*cj[kk+2] + ci[kk+3]*cj[kk+3]
                   + ci[kk+4]*cj[kk+4] + ci[kk+5]*cj[kk+5] + ci[kk+6]*cj[kk+6] + ci[kk+7]*cj[kk+7];
            for (; kk < n_embd; kk++) s += ci[kk] * cj[kk];

            /* top-K 插入 */
            for (int k = 0; k < KK; k++) {
                if (s > bsim[k]) {
                    for (int t = KK - 1; t > k; t--) { best[t] = best[t-1]; bsim[t] = bsim[t-1]; }
                    best[k] = j; bsim[k] = s; break;
                }
            }
        }
        for (int k = 0; k < K; k++) {
            g->neighbor[(size_t)i * K + k] = (k < KK) ? best[k] : -1;
            g->sim[(size_t)i * K + k] = (k < KK) ? bsim[k] : -2.0f;
        }
        if ((i & 4095) == 0) printf("[graph] built %d/%d concepts\n", i, V);
    }
    free(norm_wte);
    printf("[graph] done (inplace, V=%d K=%d)\n", V, K);
    return 0;
}

static void free_concept_graph(ConceptGraph *g) {
    if (g->neighbor) free(g->neighbor);
    if (g->sim) free(g->sim);
    g->neighbor = NULL; g->sim = NULL;
}

/* === 缓存版: 构建后存盘, 下次直接加载 ===
 * 缓存文件: concept_graph_cache.bin (与 ckpt 同目录)
 * 缓存格式: [magic=0x43473321] [V] [n_embd] [K] [neighbor[V*K]] [sim[V*K]]
 * 缓存失效条件: V/n_embd/K 不匹配时重建
 * 效果: 首次 30min -> 存盘; 后续 <1s 加载
 */
static int build_concept_graph_cached(Model *m, ConceptGraph *g, int K) {
    if (K <= 0) K = CG_K;
    int V = m->cfg.vocab_size;
    int n_embd = m->cfg.n_embd;

    /* 尝试加载缓存
     * v22: 多进程时每个 worker 用独立缓存文件, 避免并发读写冲突 */
    const char *wid = getenv("LAL_WORKER_ID");
    char cache_buf[256];
    if (wid) snprintf(cache_buf, sizeof(cache_buf), "concept_graph_cache_w%s.bin", wid);
    else snprintf(cache_buf, sizeof(cache_buf), "concept_graph_cache.bin");
    const char *cache_path = cache_buf;
    FILE *f = fopen(cache_path, "rb");
    if (f) {
        int magic, cv, cn, ck;
        if (fread(&magic, 4, 1, f) == 1 && magic == (int)CONCEPT_GRAPH_MAGIC &&
            fread(&cv, 4, 1, f) == 1 && cv == V &&
            fread(&cn, 4, 1, f) == 1 && cn == n_embd &&
            fread(&ck, 4, 1, f) == 1 && ck == K) {
            /* 缓存匹配, 加载 */
            g->vocab_size = V; g->n_embd = n_embd; g->K = K;
            g->neighbor = (int *)calloc((size_t)V * K, sizeof(int));
            g->sim = (float *)calloc((size_t)V * K, sizeof(float));
            fread(g->neighbor, sizeof(int), (size_t)V * K, f);
            fread(g->sim, sizeof(float), (size_t)V * K, f);
            fclose(f);
            printf("[graph] loaded from cache (%s, V=%d K=%d)\n", cache_path, V, K);
            return 0;
        }
        fclose(f);
        printf("[graph] cache mismatch, rebuilding\n");
    }

    /* 缓存不存在或不匹配, 构建 */
    int ret = build_concept_graph_inplace(m, g, K);
    if (ret != 0) return ret;

    /* 存盘 */
    f = fopen(cache_path, "wb");
    if (f) {
        int magic = CONCEPT_GRAPH_MAGIC;
        fwrite(&magic, 4, 1, f);
        fwrite(&V, 4, 1, f);
        fwrite(&n_embd, 4, 1, f);
        fwrite(&K, 4, 1, f);
        fwrite(g->neighbor, sizeof(int), (size_t)V * K, f);
        fwrite(g->sim, sizeof(float), (size_t)V * K, f);
        fclose(f);
        printf("[graph] saved to cache (%s)\n", cache_path);
    }
    return 0;
}

/* 离线存盘版(保留, 调试/可视化用; 主路径不再经过它) */
static int build_concept_graph(Model *m, const char *out_path, int K) {
    if (K <= 0) K = CG_K;
    ConceptGraph g;
    if (build_concept_graph_inplace(m, &g, K) != 0) return 1;

    int V = g.vocab_size, n_embd = g.n_embd, KK = g.K;
    FILE *f = fopen(out_path, "wb");
    if (!f) { fprintf(stderr, "[graph] cannot write %s\n", out_path); free_concept_graph(&g); return 2; }
    int32_t magic = CONCEPT_GRAPH_MAGIC, ver = CG_VERSION;
    int32_t vV = V, vE = n_embd, vK = KK;
    fwrite(&magic, 4, 1, f); fwrite(&ver, 4, 1, f);
    fwrite(&vV, 4, 1, f); fwrite(&vE, 4, 1, f); fwrite(&vK, 4, 1, f);
    for (int i = 0; i < V; i++) {
        int32_t ii = i;
        fwrite(&ii, 4, 1, f);
        fwrite(g.neighbor + (size_t)i * KK, 4, KK, f);
        fwrite(g.sim + (size_t)i * KK, 4, KK, f);
    }
    fclose(f);
    free_concept_graph(&g);
    printf("[graph] done: %s (V=%d K=%d)\n", out_path, V, KK);
    return 0;
}

static int load_concept_graph(ConceptGraph *g, const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "[graph] cannot open %s\n", path); return 1; }
    int32_t magic, ver, V, E, K;
    if (fread(&magic, 4, 1, f) != 1 || fread(&ver, 4, 1, f) != 1 ||
        fread(&V, 4, 1, f) != 1 || fread(&E, 4, 1, f) != 1 || fread(&K, 4, 1, f) != 1) {
        fprintf(stderr, "[graph] bad header\n"); fclose(f); return 2;
    }
    if (magic != CONCEPT_GRAPH_MAGIC) { fprintf(stderr, "[graph] bad magic\n"); fclose(f); return 3; }
    g->vocab_size = V; g->n_embd = E; g->K = K;
    g->neighbor = (int *)malloc(sizeof(int) * (size_t)V * K);
    g->sim = (float *)malloc(sizeof(float) * (size_t)V * K);
    for (int i = 0; i < V; i++) {
        int32_t ii;
        if (fread(&ii, 4, 1, f) != 1) { fprintf(stderr, "[graph] truncated\n"); fclose(f); return 4; }
        if (fread(g->neighbor + (size_t)i * K, 4, K, f) != (size_t)K ||
            fread(g->sim + (size_t)i * K, 4, K, f) != (size_t)K) {
            fprintf(stderr, "[graph] truncated body\n"); fclose(f); return 5;
        }
    }
    fclose(f);
    printf("[graph] loaded %s (V=%d K=%d)\n", path, V, K);
    return 0;
}

/* ---------------------------------------------------------------------------
 * 生成端: 图引导自回归 (图遍历骨架 + 自回归填充)
 * ---------------------------------------------------------------------------
 * 思路(两阶段):
 *   阶段1 图遍历: 从 prompt 命中的概念出发, 沿概念图做 BFS/链遍历, 收集一串
 *     "概念骨架" token(按相似度排序的邻居链).
 *   阶段2 自回归填充: 标准 transformer 前向续写, 但每步对 logits 加 "图偏置" —
 *     若候选 token 属于"当前焦点概念的邻居集"则抬高, 使生成沿概念链走向连贯.
 *
 * 调用方在 generate_text 解码循环里每步调用本函数拿到 bias 数组, 加到 logits 上.
 * 这里只产出 bias; 真正的自回归/采样仍在 generate_text 内(闭环终点).
 * ------------------------------------------------------------------------- */
static void graph_concept_bias(const ConceptGraph *g, const int *hist, int n_hist,
                               float *bias, int vocab_size, float strength) {
    memset(bias, 0, vocab_size * sizeof(float));
    if (!g || g->vocab_size == 0) return;
    int K = g->K;

    /* 焦点概念: 已生成序列末位 token(优先), 否则 0 */
    int focus = (n_hist > 0) ? hist[n_hist - 1] : -1;
    if (focus < 0 || focus >= g->vocab_size) return;

    /* 阶段1: 焦点概念的 top-K 邻居 = "下一概念集", 越近加越多 */
    for (int k = 0; k < K; k++) {
        int nb = g->neighbor[(size_t)focus * K + k];
        if (nb < 0 || nb >= vocab_size) continue;
        float s = g->sim[(size_t)focus * K + k];
        bias[nb] += strength * (s > 0 ? s : 0.0f);  /* 用余弦相似度作权重 */
    }

    /* 阶段1扩展: 若历史非空, 也对"上一个焦点"的邻居轻度加成, 维持概念链连贯 */
    if (n_hist >= 2) {
        int prev = hist[n_hist - 2];
        if (prev >= 0 && prev < g->vocab_size) {
            for (int k = 0; k < K; k++) {
                int nb = g->neighbor[(size_t)prev * K + k];
                if (nb < 0 || nb >= vocab_size) continue;
                float s = g->sim[(size_t)prev * K + k];
                bias[nb] += 0.4f * strength * (s > 0 ? s : 0.0f);
            }
        }
    }
}

/* ---------------------------------------------------------------------------
 * 探针 — 概念结构可视化 (验证用, 不参与主路径)
 * ------------------------------------------------------------------------- */
static void lal_concept_probe(Model *m, const char *pair_str) {
    int n_embd = m->cfg.n_embd;
    printf("\n=== LAL Concept Probe (C3 validation) ===\n");

    static char buf[8192];
    strncpy(buf, pair_str, sizeof(buf) - 1);
    buf[sizeof(buf) - 1] = '\0';
    for (char *p = buf; *p; p++) if (*p == ';') *p = '\n';

    float ea[4096], eb[4096];
    char *line = strtok(buf, "\n");
    while (line) {
        while (*line == ' ' || *line == '\t') line++;
        char *comma = strchr(line, ',');
        if (!comma) { line = strtok(NULL, "\n"); continue; }
        *comma = '\0';
        char *a = line; char *b = comma + 1;
        while (*b == ' ') b++;
        if (a[0] == '\0' || b[0] == '\0') { line = strtok(NULL, "\n"); continue; }

        get_concept_embedding(m, a, ea, n_embd);
        get_concept_embedding(m, b, eb, n_embd);
        float wte_cos = cosine_sim(ea, eb, n_embd);

        float cd_ab = 0, cd_ba = 0;
        float d_ab = core_directional_drive(m, a, b, 0.3f, &cd_ab);
        float d_ba = core_directional_drive(m, b, a, 0.3f, &cd_ba);
        float dir = d_ab - d_ba;

        printf("  %s ↔ %s : wte_cos=%+.4f  dir(A→B - B→A)=%+.4f  drive_ab=%+.4f drive_ba=%+.4f\n",
               a, b, wte_cos, dir, d_ab, d_ba);
        (void)cd_ab; (void)cd_ba;
        line = strtok(NULL, "\n");
    }
    printf("  (dir>0 表示 A 驱动 B, 即 A 为施事; 概念越相关 wte_cos 越大)\n");
}

/* ---------------------------------------------------------------------------
 * 训练端 C3 辅助 loss — 全 vocab 概念链 (一个 token = 一个概念)
 * ---------------------------------------------------------------------------
 * 对采样的若干概念 token, 预测其 wte 余弦最近的若干邻居 token(全 vocab 搜索),
 * 做辅助 CE, 梯度注入 grad_wte_accum(与 CE 同构), 使概念相邻结构被联合优化.
 *
 * [accel v1] 性能改造:
 *   1. C3_N_CONCEPTS 512 → 64 (8x 减少外层迭代)
 *   2. 外层 for idx 加 OpenMP (2 核 → 2x)
 *   3. 内层找 top-K 邻居 + logits 计算 + 梯度累加 三个 V 循环合并成一次遍历
 *      (原代码每对概念要 3 次 V 循环 = 3×32k, 现在 2 次)
 *   4. 共享 logits 缓冲改 threadprivate, 避免每对 malloc/free
 * 总加速预期: 8 × 2 × 1.5 = 24x, 单步从 13min 降到 ~30s */
#define C3_TOPK 3
#define C3_N_CONCEPTS 32  /* v2: 64→32, 概念图缓存让采样更高效, 2x C3 加速 */

static float concept_chain_loss(Model *m, int n_embd, float weight) {
    if (weight <= 0.0f || !m->grad_wte_accum) return 0.0f;
    int V = m->cfg.vocab_size;

    /* v3: Sampled Softmax — 不遍历全 vocab, 只算 target + 随机负样本
     * 旧: 32 × O(V × n_embd) = 32 × 32768 × 512 = 5.4亿次乘加 → ~8000ms
     * 新: 32 × O(K_NEG × n_embd) = 32 × 256 × 512 = 420万次乘加 → ~6ms
     * 加速: ~1300x, 总训练步时预期: 21s → ~6s */
    #define C3_N_NEG 256  /* 随机负样本数 (从 32768 降到 256) */

    int stride = (V > C3_N_CONCEPTS) ? (V / C3_N_CONCEPTS) : 1;
    int n_concepts_sampled = (V + stride - 1) / stride;
    if (n_concepts_sampled > C3_N_CONCEPTS) n_concepts_sampled = C3_N_CONCEPTS;

    float loss_sum = 0.0f;
    int n_pairs = 0;

    #pragma omp parallel for schedule(static) reduction(+:loss_sum, n_pairs)
    for (int ci = 0; ci < n_concepts_sampled; ci++) {
        int idx = 4 + ci * stride;
        if (idx >= V) continue;
        const float *cur = m->wte + (size_t)idx * n_embd;

        /* 1. 从缓存图读 top-K 邻居 (target tokens) */
        int best[C3_TOPK];
        for (int k = 0; k < C3_TOPK; k++) best[k] = -1;
        if (g_runtime_cg && g_runtime_cg->K >= C3_TOPK) {
            int graph_k = g_runtime_cg->K;
            for (int k = 0; k < C3_TOPK; k++)
                best[k] = g_runtime_cg->neighbor[(size_t)idx * graph_k + k];
        }

        /* 2. 构建候选集: target + 随机负样本 (Sampled Softmax)
         * 不遍历全 vocab, 只算这 C3_N_NEG + C3_TOPK 个 token 的 logits */
        int candidates[C3_N_NEG + C3_TOPK];
        float cand_logits[C3_N_NEG + C3_TOPK];
        int n_cand = 0;

        /* 加入 target */
        for (int k = 0; k < C3_TOPK; k++) {
            if (best[k] >= 0 && best[k] < V) {
                candidates[n_cand] = best[k];
                n_cand++;
            }
        }
        /* 加入随机负样本 */
        while (n_cand < C3_N_NEG + C3_TOPK && n_cand < C3_N_NEG + 3) {
            int r = rand() % V;
            /* 避免重复 */
            int dup = 0;
            for (int k = 0; k < n_cand; k++) if (candidates[k] == r) { dup = 1; break; }
            if (!dup && r != idx) {
                candidates[n_cand] = r;
                n_cand++;
            }
        }

        /* 3. 只算候选集的 logits (cur · wte[candidate]) */
        float maxl = -1e30f;
        for (int j = 0; j < n_cand; j++) {
            const float *w = m->wte + (size_t)candidates[j] * n_embd;
            float dot = 0;
            for (int kk = 0; kk + 7 < n_embd; kk += 8)
                dot += cur[kk]*w[kk] + cur[kk+1]*w[kk+1] + cur[kk+2]*w[kk+2] + cur[kk+3]*w[kk+3]
                     + cur[kk+4]*w[kk+4] + cur[kk+5]*w[kk+5] + cur[kk+6]*w[kk+6] + cur[kk+7]*w[kk+7];
            for (int kk = (n_embd/8)*8; kk < n_embd; kk++) dot += cur[kk]*w[kk];
            cand_logits[j] = dot;
            if (dot > maxl) maxl = dot;
        }

        /* 4. Sampled Softmax + 梯度注入 (只在候选集上) */
        for (int k = 0; k < C3_TOPK; k++) {
            int target_tok = best[k];
            if (target_tok < 0) continue;
            /* 找 target 在 candidates 中的位置 */
            int target_pos = -1;
            for (int j = 0; j < n_cand; j++) {
                if (candidates[j] == target_tok) { target_pos = j; break; }
            }
            if (target_pos < 0) continue;

            /* softmax on candidates */
            float sum = 0;
            for (int j = 0; j < n_cand; j++)
                sum += expf(cand_logits[j] - maxl);
            float inv = 1.0f / (sum + 1e-12f);

            /* loss */
            float p_target = expf(cand_logits[target_pos] - maxl) * inv;
            loss_sum += -logf(p_target + 1e-12f);

            /* 梯度注入 (只在候选集上, 不遍历全 vocab) */
            float *g = m->grad_wte_accum;
            for (int j = 0; j < n_cand; j++) {
                float p = expf(cand_logits[j] - maxl) * inv;
                float coef = p;
                if (candidates[j] == target_tok) coef -= 1.0f;
                coef *= weight;
                if (fabsf(coef) < 1e-7f) continue;
                float *gw = &g[(size_t)candidates[j] * n_embd];
                for (int kk = 0; kk + 7 < n_embd; kk += 8) {
                    gw[kk+0] += coef*cur[kk+0]; gw[kk+1] += coef*cur[kk+1];
                    gw[kk+2] += coef*cur[kk+2]; gw[kk+3] += coef*cur[kk+3];
                    gw[kk+4] += coef*cur[kk+4]; gw[kk+5] += coef*cur[kk+5];
                    gw[kk+6] += coef*cur[kk+6]; gw[kk+7] += coef*cur[kk+7];
                }
                for (int kk = (n_embd/8)*8; kk < n_embd; kk++) gw[kk] += coef*cur[kk];
            }
            n_pairs++;
        }
    }
    if (n_pairs == 0) return 0.0f;
    return loss_sum / n_pairs;
}

#endif /* LAL_CONCEPT_GEN_H */