#!/usr/bin/env python3
# RAG 离线检索质量评测：四模式(向量/全文/图谱/混合) recall@k + MRR，量化"混合 vs 单路"。
import json, sys, time, urllib.request

BASE = "http://127.0.0.1:8080"
KB = "evalset"
TOKEN = open(sys.argv[1]).read().strip()

# 语料：每条一句话事实，唯一标记=载具名（用于判定召回）。
CORPUS = [
    ("启明星一号", "启明星一号是一颗对地观测卫星，由总设计师周文渊研制，搭载高光谱相机。"),
    ("玄鸟二号", "玄鸟二号是一型运载火箭，由总工程师林书豪设计，使用液氧甲烷发动机。"),
    ("蛟龙三号", "蛟龙三号是一台载人潜水器，由首席科学家陈海平主持，最大下潜深度7000米。"),
    ("烛龙四号", "烛龙四号是一颗深空探测器，由项目主管吴宇航负责，用于探测木星磁场。"),
    ("麒麟五号", "麒麟五号是一座空间站，由总指挥赵长青建造，配备机械臂。"),
    ("白泽六号", "白泽六号是一颗通信卫星，由总设计师钱明轩研制，实现量子通信。"),
    ("饕餮七号", "饕餮七号是一艘破冰船，由总工程师孙铁柱主持，采用核动力驱动。"),
    ("穷奇八号", "穷奇八号是一艘行星际飞船，由首席科学家周伯通设计，依靠太阳帆推进。"),
    ("梼杌九号", "梼杌九号是一架预警机，由总设计师黄药师研制，配红外制导系统。"),
    ("应龙十号", "应龙十号是一款高超声速飞行器，由总工程师欧阳锋设计，采用折叠翼。"),
    ("重明十一号", "重明十一号是一架电动垂直起降飞行器，由总设计师洪七公研制，用全固态电池。"),
    ("毕方十二号", "毕方十二号是一座防御平台，由总指挥段誉建造，装备激光武器。"),
    # —— 干扰项：与各目标同族近义，逼向量路在相似项中区分 ——
    ("天罡卫星", "天罡卫星是一颗对地观测卫星，由总设计师周文博研制，搭载多光谱相机。"),
    ("地煞火箭", "地煞火箭是一型运载火箭，由总工程师林书文设计，使用液氧煤油发动机。"),
    ("潜龙号", "潜龙号是一台载人潜水器，由首席科学家陈海峰主持，最大下潜深度5000米。"),
    ("探火号", "探火号是一颗深空探测器，由项目主管吴宇辉负责，用于探测火星大气。"),
    ("天宫号", "天宫号是一座空间站，由总指挥赵长河建造，配备太阳能帆板。"),
    ("鹊桥号", "鹊桥号是一颗通信卫星，由总设计师钱明远研制，做激光通信。"),
    ("破冰者号", "破冰者号是一艘破冰船，由总工程师孙铁山主持，采用柴油动力。"),
    ("远航号", "远航号是一艘行星际飞船，由首席科学家周伯文设计，依靠核电推进。"),
    ("天眼号", "天眼号是一架预警机，由总设计师黄药剂研制，配相控阵雷达。"),
    ("风刃号", "风刃号是一款高超声速飞行器，由总工程师欧阳峰设计，采用乘波体。"),
    ("飞羽号", "飞羽号是一架电动垂直起降飞行器，由总设计师洪七叔研制，用液流电池。"),
    ("盾卫号", "盾卫号是一座防御平台，由总指挥段正淳建造，装备动能拦截弹。"),
]
# 查询：q=问句, m=期望命中的标记, t=类型(便于看各路强弱)
QUERIES = [
    ("哪台设备可以下潜到七千米的深海", "蛟龙三号", "语义"),
    ("用来在冰封海域开路的核动力船", "饕餮七号", "语义"),
    ("能垂直起降的电动飞行器用了什么电池", "重明十一号", "语义"),
    ("靠激光拦截目标的防御系统是哪个", "毕方十二号", "语义"),
    ("液氧甲烷发动机用在哪", "玄鸟二号", "关键词"),
    ("木星磁场探测", "烛龙四号", "关键词"),
    ("量子通信卫星", "白泽六号", "关键词"),
    ("太阳帆推进", "穷奇八号", "关键词"),
    ("周文渊负责研制的是什么", "启明星一号", "实体关系"),
    ("孙铁柱主持的项目", "饕餮七号", "实体关系"),
    ("黄药师设计了哪款预警机", "梼杌九号", "实体关系"),
    ("赵长青建造的空间站叫什么", "麒麟五号", "实体关系"),
    # —— 纯语义改写：尽量避开与目标文档的字面重叠，逼全文路失手、看混合是否兜住 ——
    ("可搭乘人员探入数千米幽暗海沟的勘探艇", "蛟龙三号", "纯语义"),
    ("环绕巨行星测量其磁场强度的太空仪器", "烛龙四号", "纯语义"),
    ("能从陆地竖直升空、不烧油的载客飞行器", "重明十一号", "纯语义"),
    ("以高能光束摧毁来袭目标的拦截阵地", "毕方十二号", "纯语义"),
]
MODES = ["vector", "fulltext", "graph", "hybrid"]

def post(path, body):
    req = urllib.request.Request(BASE + path, data=json.dumps(body).encode(),
        headers={"Authorization": "Bearer " + TOKEN, "Content-Type": "application/json"}, method="POST")
    return json.load(urllib.request.urlopen(req, timeout=20))

def ingest():
    for name, text in CORPUS:
        post("/api/v1/kb/note", {"kb": KB, "name": name, "content": text})
    print(f"已提交 {len(CORPUS)} 篇，等待入库(切块/向量/图谱)…")

def search(q, mode, topk=5):
    r = post("/api/v1/kb/search", {"kb": KB, "q": q, "topK": topk, "mode": mode})
    return [h.get("Text") or h.get("text") or "" for h in (r.get("hits") or [])]

def rank_of(hits, marker):
    for i, t in enumerate(hits):
        if marker in t:
            return i + 1
    return 0

def main():
    if "--noingest" not in sys.argv:
        ingest()
        time.sleep(40)  # 等队列把 12 篇全部 embed+图谱抽取完
    print(f"\n{'模式':<10}{'recall@3':>10}{'recall@5':>10}{'MRR':>8}")
    print("-" * 40)
    per_type = {m: {} for m in MODES}
    for mode in MODES:
        r3 = r5 = mrr = 0.0
        bytype = {}
        for q, marker, typ in QUERIES:
            hits = search(q, mode)
            rk = rank_of(hits, marker)
            hit3, hit5 = (1 if 0 < rk <= 3 else 0), (1 if 0 < rk <= 5 else 0)
            r3 += hit3; r5 += hit5; mrr += (1.0 / rk if rk else 0.0)
            bytype.setdefault(typ, [0, 0]); bytype[typ][0] += hit5; bytype[typ][1] += 1
        n = len(QUERIES)
        print(f"{mode:<10}{r3/n:>10.2f}{r5/n:>10.2f}{mrr/n:>8.2f}")
        per_type[mode] = {t: f"{v[0]}/{v[1]}" for t, v in bytype.items()}
    types = list(dict.fromkeys(t for _, _, t in QUERIES))
    print("\n各模式 recall@5 分查询类型（看各路强弱）:")
    print("模式".ljust(10) + "".join(t.rjust(10) for t in types))
    for mode in MODES:
        d = per_type[mode]
        print(mode.ljust(10) + "".join(d.get(t, "").rjust(10) for t in types))

if __name__ == "__main__":
    main()