/* ========================================================================
 * 高级语义探针系统 — 多层级概念结构监控
 * ========================================================================
 * 设计理念: 让训练从黑箱变透明, 像"认知发展心理学"一样观察模型学到了什么
 *
 * 5 层探针:
 *   L1. 基础概念层: 模型是否学会区分基本概念 (反义词/同义词/无关词)
 *   L2. 概念聚类层: 概念是否自动形成语义簇 (温度/大小/方向/颜色/情感)
 *   L3. 概念组合层: 复杂概念是否由基础概念组成 (火=热+亮+危险)
 *   L4. 层间传递层: 概念在 10 层中如何传递和变换
 *   L5. 生成质量层: 实际生成文本的质量评估
 * ======================================================================== */

/* L1: 基础概念区分 — 扩展探针对, 覆盖更多语义类别 */
static ConceptPair probe_pairs_extended[] = {
    /* 原有 7 对反义词 */
    {"\xe7\x83\xad", "\xe7\x83\xad", "\xe5\x86\xb7", "\xe5\x86\xb7", "temperature"},
    {"\xe5\xa4\xa7", "\xe5\xa4\xa7", "\xe5\xb0\x8f", "\xe5\xb0\x8f", "size"},
    {"\xe4\xb8\x8a", "\xe4\xb8\x8a", "\xe4\xb8\x8b", "\xe4\xb8\x8b", "position"},
    {"\xe4\xba\xae", "\xe4\xba\xae", "\xe6\x9a\x97", "\xe6\x9a\x97", "light"},
    {"\xe9\x87\x8d", "\xe9\x87\x8d", "\xe8\xbd\xbb", "\xe8\xbd\xbb", "weight"},
    {"\xe5\xbf\xab", "\xe5\xbf\xab", "\xe6\x85\xa2", "\xe6\x85\xa2", "speed"},
    {"\xe6\xb9\xbf", "\xe6\xb9\xbf", "\xe5\xb9\xb2", "\xe5\xb9\xb2", "moisture"},
    /* 新增: 同义词对 (应该高相似度) */
    {"\xe7\x83\xad", "\xe7\x83\xad", "\xe6\x9a\x96", "\xe6\x9a\x96", "syn_hot"},     /* 热-暖 */
    {"\xe5\xa4\xa7", "\xe5\xa4\xa7", "\xe5\xb7\xa8", "\xe5\xb7\xa8", "syn_big"},     /* 大-巨 */
    {"\xe5\xbf\xab", "\xe5\xbf\xab", "\xe9\x80\x9f", "\xe9\x80\x9f", "syn_fast"},    /* 快-速 */
};

/* L2: 概念聚类 — 同类概念应该聚在一起, 不同类应该分开 */
typedef struct { const char *name; const char *concepts[6]; int n; } ConceptCluster;
static ConceptCluster probe_clusters[] = {
    {"temp", {"热","冷","暖","凉","温","冰"}, 6},
    {"size", {"大","小","巨","微","长","短"}, 6},
    {"color", {"红","黄","蓝","绿","白","黑"}, 6},
    {"emotion", {"喜","怒","哀","乐","爱","恨"}, 6},
    {"nature", {"水","火","土","木","金","风"}, 6},
};
#define N_CLUSTERS (sizeof(probe_clusters) / sizeof(probe_clusters[0]))

/* L3: 概念组合 — 复杂概念 = 基础概念的加权组合
 * 火 ≈ 热 + 亮 + 红 (颜色) + 危险
 * 水 ≈ 湿 + 冷 + 透明 + 流动
 * 测量: complex_concept 与 basic_concepts 的余弦相似度排名 */
typedef struct {
    const char *name;
    const char *complex_concept;  /* 复杂概念 */
    const char *basic_concepts[6]; /* 基础概念 */
    int n_basics;
} CompositionProbe;
static CompositionProbe composition_probes[] = {
    {"fire", "火", {"热","亮","红","危险","烧","光"}, 6},
    {"water", "水", {"湿","冷","清","流","喝","冰"}, 6},
    {"sun", "太阳", {"热","亮","光","圆","远","黄"}, 6},
    {"tree", "树", {"木","绿","叶","根","高","植物"}, 6},
};
#define N_COMPOSITIONS (sizeof(composition_probes) / sizeof(CompositionProbe[0]))

/* ========================================================================
 * L2: 概念聚类分析 — 测量同类内聚 vs 异类分离
 * ======================================================================== */
static void probe_concept_clusters(Model *m) {
    int n_embd = m->cfg.n_embd;
    int n_c = (int)N_CLUSTERS;

    /* 堆分配避免栈溢出 (5×6×4096×4 = 480KB) */
    float *embs = (float *)malloc(5 * 6 * n_embd * sizeof(float));
    for (int c = 0; c < n_c; c++) {
        for (int i = 0; i < probe_clusters[c].n; i++) {
            get_concept_embedding(m, probe_clusters[c].concepts[i],
                                  embs + ((size_t)c * 6 + i) * n_embd, n_embd);
        }
    }

    float intra_sim = 0, inter_sim = 0;
    int n_intra = 0, n_inter = 0;

    for (int c = 0; c < n_c; c++) {
        for (int i = 0; i < probe_clusters[c].n; i++) {
            for (int j = i + 1; j < probe_clusters[c].n; j++) {
                intra_sim += cosine_sim(embs + ((size_t)c * 6 + i) * n_embd,
                                        embs + ((size_t)c * 6 + j) * n_embd, n_embd);
                n_intra++;
            }
        }
    }

    for (int c1 = 0; c1 < n_c; c1++) {
        for (int c2 = c1 + 1; c2 < n_c; c2++) {
            for (int i = 0; i < probe_clusters[c1].n; i++) {
                for (int j = 0; j < probe_clusters[c2].n; j++) {
                    inter_sim += cosine_sim(embs + ((size_t)c1 * 6 + i) * n_embd,
                                            embs + ((size_t)c2 * 6 + j) * n_embd, n_embd);
                    n_inter++;
                }
            }
        }
    }

    intra_sim = n_intra > 0 ? intra_sim / n_intra : 0;
    inter_sim = n_inter > 0 ? inter_sim / n_inter : 0;
    float cluster_score = intra_sim - inter_sim;  /* 越正越好 */

    printf("  [L2] cluster: intra=%.3f inter=%.3f score=%+.3f (intra应>inter)\n",
           intra_sim, inter_sim, cluster_score);
    free(embs);
}

/* ========================================================================
 * L3: 概念组合分析 — 复杂概念由哪些基础概念组成
 * ======================================================================== */
static void probe_composition(Model *m) {
    int n_embd = m->cfg.n_embd;
    float *complex_emb = (float *)malloc(n_embd * sizeof(float));
    float *basic_emb = (float *)malloc(n_embd * sizeof(float));

    for (int p = 0; p < (int)N_COMPOSITIONS; p++) {
        CompositionProbe *cp = &composition_probes[p];
        get_concept_embedding(m, cp->complex_concept, complex_emb, n_embd);

        /* 算复杂概念与每个基础概念的相似度, 排序 */
        float sims[6];
        for (int i = 0; i < cp->n_basics; i++) {
            get_concept_embedding(m, cp->basic_concepts[i], basic_emb, n_embd);
            sims[i] = cosine_sim(complex_emb, basic_emb, n_embd);
        }

        /* 找 top-3 最相关基础概念 */
        int top3[3] = {-1, -1, -1};
        float top3_sim[3] = {-2, -2, -2};
        for (int i = 0; i < cp->n_basics; i++) {
            for (int k = 0; k < 3; k++) {
                if (sims[i] > top3_sim[k]) {
                    for (int t = 2; t > k; t--) {
                        top3[t] = top3[t-1]; top3_sim[t] = top3_sim[t-1];
                    }
                    top3[k] = i; top3_sim[k] = sims[i]; break;
                }
            }
        }

        printf("  [L3] %s(%s) ≈ %s(%.2f) + %s(%.2f) + %s(%.2f)\n",
               cp->name, cp->complex_concept,
               cp->basic_concepts[top3[0]], top3_sim[0],
               cp->basic_concepts[top3[1]], top3_sim[1],
               cp->basic_concepts[top3[2]], top3_sim[2]);
    }
    free(complex_emb);
    free(basic_emb);
}

/* ========================================================================
 * L4: 层间概念传递 — 概念在每层的激活模式
 * ========================================================================
 */
static void probe_layer_propagation(Model *m) {
    int n_embd = m->cfg.n_embd;
    int mlp_dim = m->cfg.mlp_dim;
    int n_layer = m->cfg.n_layer;

    /* 取"热"和"冷"的概念, 看每层 CORE 激活差异 */
    float emb_a[4096], emb_b[4096];
    get_concept_embedding(m, "\xe7\x83\xad", emb_a, n_embd);  /* 热 */
    get_concept_embedding(m, "\xe5\x86\xb7", emb_b, n_embd);  /* 冷 */

    float *act_a = (float *)malloc(mlp_dim * sizeof(float));
    float *act_b = (float *)malloc(mlp_dim * sizeof(float));
    float gate_a[4096], gate_b[4096];

    printf("  [L4] 热/冷 层间 CORE 差异: ");
    for (int l = 0; l < n_layer; l++) {
        compute_gate_input(m, emb_a, l, gate_a, n_embd);
        compute_gate_input(m, emb_b, l, gate_b, n_embd);
        simulate_activation(m, gate_a, l, act_a, mlp_dim);
        simulate_activation(m, gate_b, l, act_b, mlp_dim);

        uint8_t *mask = m->layers[l].mlp_gate.logic_mask;
        if (!mask) continue;
        float cd = 0; int nc = 0;
        for (int j = 0; j < mlp_dim; j++) {
            if (mask[j] == 0) { cd += fabsf(act_a[j] - act_b[j]); nc++; }
        }
        cd = nc > 0 ? cd / nc : 0;
        printf("L%d=%.3f ", l, cd);
    }
    printf("\n");
    free(act_a); free(act_b);
}

/* ========================================================================
 * L5: 概念最近邻 — 模型认为哪些概念最相似
 * ========================================================================
 */
static void probe_nearest_neighbors(Model *m) {
    int n_embd = m->cfg.n_embd;
    int V = m->cfg.vocab_size;

    /* 探针概念: 看"火"的 top-5 最近邻是什么 */
    const char *probes[] = {"火", "水", "大", "热"};
    int n_probes = 4;

    for (int p = 0; p < n_probes; p++) {
        float emb[4096];
        get_concept_embedding(m, probes[p], emb, n_embd);

        /* 找 top-5 最近邻 (只搜前 1000 个 token, 避免太慢) */
        int best[5]; float bsim[5];
        for (int k = 0; k < 5; k++) { best[k] = -1; bsim[k] = -2; }
        int search_range = V < 1000 ? V : 1000;
        for (int j = 0; j < search_range; j++) {
            const float *w = m->wte + (size_t)j * n_embd;
            float s = cosine_sim(emb, w, n_embd);
            for (int k = 0; k < 5; k++) {
                if (s > bsim[k]) {
                    for (int t = 4; t > k; t--) { best[t] = best[t-1]; bsim[t] = bsim[t-1]; }
                    best[k] = j; bsim[k] = s; break;
                }
            }
        }

        printf("  [L5] \"%s\" 最近邻: ", probes[p]);
        for (int k = 0; k < 5; k++) {
            printf("id%d(%.2f) ", best[k], bsim[k]);
        }
        printf("\n");
    }
}

/* ========================================================================
 * L6: 语法探针 — 检测模型是否学会语法结构
 * ========================================================================
 * 6 个子探针:
 *   L6a 词序敏感性: "大火" vs "火大" — 词序不同意思不同, 模型应该区分
 *   L6b 搭配准确性: "喝+水" 高频搭配 vs "喝+火" 不搭配 — 模型应该偏好高频搭配
 *   L6c 语法角色: "红的苹果"(修饰) vs "苹果的红"(中心) — 模型应该区分修饰/中心
 *   L6d 句子连贯性: 正确语序 vs 打乱语序 — 模型对正确语序应该 loss 更低
 *   L6e 否定理解: "不热" vs "热" — 模型应该区分否定
 *   L6f 量词搭配: "一只猫" 正确 vs "一条猫" 错误 — 模型应该偏好正确量词
 * ======================================================================== */

/* L6a: 词序敏感性 — 比较大小关系: "大火"(形容火势) vs "火大"(形容脾气) */
static void probe_word_order(Model *m) {
    int n_embd = m->cfg.n_embd;
    /* 大火 = 大 + 火 (wte 求和), 火大 = 火 + 大 (同样求和但顺序不同)
     * wte 求和是交换的, 但模型 forward 不是 (attention 有位置编码)
     * 所以看 wte 层面应该相同, forward 层面应该不同 */
    float e_da[4096], e_hd[4096], e_d[4096], e_h[4096];
    get_concept_embedding(m, "\xe5\xa4\xa7\xe7\x83\xad", e_da, n_embd);  /* 大火 */
    get_concept_embedding(m, "\xe7\x83\xad\xe5\xa4\xa7", e_hd, n_embd);  /* 火大 */
    get_concept_embedding(m, "\xe5\xa4\xa7", e_d, n_embd);               /* 大 */
    get_concept_embedding(m, "\xe7\x83\xad", e_h, n_embd);               /* 火 */

    float sim_order = cosine_sim(e_da, e_hd, n_embd);  /* wte 层: 应该 ~1.0 (求和可交换) */
    float sim_da_d = cosine_sim(e_da, e_d, n_embd);    /* 大火 vs 大 */
    float sim_da_h = cosine_sim(e_da, e_h, n_embd);    /* 大火 vs 火 */
    /* 如果大火更接近火 → 模型学到"大火"中心词是"火" */
    /* 如果大火更接近大 → 模型还没学会中心词 */

    printf("  [L6a] 词序: 大火≈火(%.2f) vs 大火≈大(%.2f)  wte交换(%.2f) 中心词=%s\n",
           sim_da_h, sim_da_d, sim_order,
           sim_da_h > sim_da_d ? "火(正确)" : "大(错误)");
}

/* L6b: 搭配准确性 — 比较"喝水的概率" vs "喝火的概率"
 * 用 wte 组合后的余弦相似度判断搭配合理性 */
static void probe_collocation(Model *m) {
    int n_embd = m->cfg.n_embd;

    /* 搭配对: (动词, 正确宾语, 错误宾语) */
    struct { const char *verb; const char *correct; const char *wrong; const char *name; } pairs[] = {
        {"\xe5\x96\x9d", "\xe6\xb0\xb4", "\xe7\x81\xab", "喝"},
        {"\xe5\x90\x83", "\xe9\xa5\xad", "\xe7\x9f\xb3", "吃"},
        {"\xe7\x9c\x8b", "\xe4\xb9\xa6", "\xe9\xa3\x8e", "看"},
    };
    int n = sizeof(pairs) / sizeof(pairs[0]);

    float *e_verb = (float *)malloc(n_embd * sizeof(float));
    float *e_correct = (float *)malloc(n_embd * sizeof(float));
    float *e_wrong = (float *)malloc(n_embd * sizeof(float));

    printf("  [L6b] 搭配: ");
    int n_correct = 0;
    for (int i = 0; i < n; i++) {
        get_concept_embedding(m, pairs[i].verb, e_verb, n_embd);
        get_concept_embedding(m, pairs[i].correct, e_correct, n_embd);
        get_concept_embedding(m, pairs[i].wrong, e_wrong, n_embd);

        float sim_correct = cosine_sim(e_verb, e_correct, n_embd);
        float sim_wrong = cosine_sim(e_verb, e_wrong, n_embd);
        /* 正确搭配: 动词和宾语应该有适度相似 (语义相关但不完全相同) */
        /* 错误搭配: 动词和宾语应该不太相似 */
        int ok = sim_correct > sim_wrong;
        if (ok) n_correct++;
        printf("%s(%s:%.2f>%s:%.2f %s) ",
               pairs[i].name, pairs[i].correct, sim_correct, pairs[i].wrong, sim_wrong,
               ok ? "✓" : "✗");
    }
    printf("  %d/%d正确\n", n_correct, n);
    free(e_verb); free(e_correct); free(e_wrong);
}

/* L6c: 否定理解 — "不热" 应该离 "冷" 更近, 离 "热" 更远 */
static void probe_negation(Model *m) {
    int n_embd = m->cfg.n_embd;
    float *e_not_hot = (float *)malloc(n_embd * sizeof(float));
    float *e_hot = (float *)malloc(n_embd * sizeof(float));
    float *e_cold = (float *)malloc(n_embd * sizeof(float));
    float *e_not = (float *)malloc(n_embd * sizeof(float));

    /* 不热 = 不 + 热 */
    get_concept_embedding(m, "\xe4\xb8\x8d\xe7\x83\xad", e_not_hot, n_embd);  /* 不热 */
    get_concept_embedding(m, "\xe7\x83\xad", e_hot, n_embd);                  /* 热 */
    get_concept_embedding(m, "\xe5\x86\xb7", e_cold, n_embd);                 /* 冷 */
    get_concept_embedding(m, "\xe4\xb8\x8d", e_not, n_embd);                  /* 不 */

    float sim_to_hot = cosine_sim(e_not_hot, e_hot, n_embd);
    float sim_to_cold = cosine_sim(e_not_hot, e_cold, n_embd);
    float sim_to_not = cosine_sim(e_not_hot, e_not, n_embd);

    /* "不热"如果学会了否定, 应该更接近"冷"而非"热" */
    printf("  [L6c] 否定: 不热≈热(%.2f) 不热≈冷(%.2f) 不热≈不(%.2f) → %s\n",
           sim_to_hot, sim_to_cold, sim_to_not,
           sim_to_cold > sim_to_hot ? "理解否定(✓)" : "未理解(✗)");
    free(e_not_hot); free(e_hot); free(e_cold); free(e_not);
}

/* L6d: 修饰关系 — "红的花" vs "花" vs "红"
 * 修饰语+中心词组合, 应该更接近中心词 */
static void probe_modifier(Model *m) {
    int n_embd = m->cfg.n_embd;
    struct { const char *phrase; const char *modifier; const char *head; const char *name; } items[] = {
        {"\xe7\xba\xa2\xe7\x9a\x84\xe8\x8a\xb1", "\xe7\xba\xa2", "\xe8\x8a\xb1", "红的花"},  /* 红的花, 红, 花 */
        {"\xe5\xa4\xa7\xe7\x9a\x84\xe7\x8b\x97", "\xe5\xa4\xa7", "\xe7\x8b\x97", "大的狗"},  /* 大的狗, 大, 狗 */
        {"\xe7\x83\xad\xe7\x9a\x84\xe6\xb0\xb4", "\xe7\x83\xad", "\xe6\xb0\xb4", "热的水"},  /* 热的水, 热, 水 */
    };
    int n = sizeof(items) / sizeof(items[0]);

    printf("  [L6d] 修饰: ");
    int n_correct = 0;
    for (int i = 0; i < n; i++) {
        float e_phrase[4096], e_mod[4096], e_head[4096];
        get_concept_embedding(m, items[i].phrase, e_phrase, n_embd);
        get_concept_embedding(m, items[i].modifier, e_mod, n_embd);
        get_concept_embedding(m, items[i].head, e_head, n_embd);

        float sim_head = cosine_sim(e_phrase, e_head, n_embd);
        float sim_mod = cosine_sim(e_phrase, e_mod, n_embd);
        /* 短语应该更接近中心词 */
        int ok = sim_head > sim_mod;
        if (ok) n_correct++;
        printf("%s(中心%.2f>修饰%.2f %s) ", items[i].name, sim_head, sim_mod, ok ? "✓" : "✗");
    }
    printf(" %d/%d正确\n", n_correct, n);
}

/* L6e: 语法类别 — 名词 vs 动词 vs 形容词
 * 同类词应该聚在一起, 异类应该分开 */
static void probe_pos_categories(Model *m) {
    int n_embd = m->cfg.n_embd;

    const char *nouns[] = {"\xe5\xa4\xaa\xe9\x98\xb3", "\xe6\xb0\xb4", "\xe6\xa0\x91", "\xe7\x8c\xab", "\xe4\xb9\xa6", "\xe7\x81\xab\xe5\xb1\xb1"};
    const char *verbs[] = {"\xe8\xb7\x91", "\xe5\x90\x83", "\xe7\x9c\x8b", "\xe8\xaf\xb4", "\xe9\xa3\x9e", "\xe7\x9d\xa1"};
    const char *adjs[] = {"\xe5\xa4\xa7", "\xe7\x83\xad", "\xe7\xba\xa2", "\xe5\xbf\xab", "\xe9\xab\x98", "\xe7\xbe\x8e"};

    /* 堆分配避免栈溢出 (6×4096×4×3 = 288KB) */
    float *en = (float *)malloc(6 * n_embd * sizeof(float));
    float *ev = (float *)malloc(6 * n_embd * sizeof(float));
    float *ea = (float *)malloc(6 * n_embd * sizeof(float));
    for (int i = 0; i < 6; i++) {
        get_concept_embedding(m, nouns[i], en + (size_t)i * n_embd, n_embd);
        get_concept_embedding(m, verbs[i], ev + (size_t)i * n_embd, n_embd);
        get_concept_embedding(m, adjs[i], ea + (size_t)i * n_embd, n_embd);
    }

    /* 算类内平均相似度 */
    float noun_intra = 0, verb_intra = 0, adj_intra = 0;
    for (int i = 0; i < 6; i++)
        for (int j = i + 1; j < 6; j++) {
            noun_intra += cosine_sim(en + (size_t)i * n_embd, en + (size_t)j * n_embd, n_embd);
            verb_intra += cosine_sim(ev + (size_t)i * n_embd, ev + (size_t)j * n_embd, n_embd);
            adj_intra += cosine_sim(ea + (size_t)i * n_embd, ea + (size_t)j * n_embd, n_embd);
        }
    noun_intra /= 15; verb_intra /= 15; adj_intra /= 15;

    /* 算类间平均相似度 */
    float inter = 0; int n_inter = 0;
    for (int i = 0; i < 6; i++) {
        for (int j = 0; j < 6; j++) {
            inter += cosine_sim(en + (size_t)i * n_embd, ev + (size_t)j * n_embd, n_embd); n_inter++;
            inter += cosine_sim(en + (size_t)i * n_embd, ea + (size_t)j * n_embd, n_embd); n_inter++;
            inter += cosine_sim(ev + (size_t)i * n_embd, ea + (size_t)j * n_embd, n_embd); n_inter++;
        }
    }
    inter /= n_inter;

    float avg_intra = (noun_intra + verb_intra + adj_intra) / 3;
    float pos_score = avg_intra - inter;

    printf("  [L6e] 词类: 名词(%.2f) 动词(%.2f) 形容(%.2f) 类间(%.2f) 聚合=%+.3f\n",
           noun_intra, verb_intra, adj_intra, inter, pos_score);
    free(en); free(ev); free(ea);
}

/* L6 语法探针总入口 */
static void probe_syntax(Model *m) {
    probe_word_order(m);       /* L6a: 词序敏感性 */
    probe_collocation(m);      /* L6b: 搭配准确性 */
    probe_negation(m);         /* L6c: 否定理解 */
    probe_modifier(m);         /* L6d: 修饰关系 */
    probe_pos_categories(m);   /* L6e: 语法类别 */
}

/* ========================================================================
 * 高级探针总入口
 * ======================================================================== */
static void advanced_probe(Model *m) {
    printf("  --- Advanced Probe ---\n");
    probe_concept_clusters(m);     /* L2: 概念聚类 */
    /* L3/L4 暂时禁用 (多字节UTF-8 + compute_gate_input 在 Windows 栈限制下不稳定) */
    probe_syntax(m);               /* L6: 语法探针 */
}