#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
深度数据挖掘 - 精修版
修正1: 药材名必须完整(>=2字)才计入统计
修正2: 去重短名(如同时匹配到'白术'和'术'，只取'白术')
修正3: 相似度计算排除只有一位共享的情况
修正4: 功效网络按功效关键词汇总，不再保留完整描述
"""
import json
import os
import re
import glob
from collections import Counter, defaultdict, OrderedDict
from itertools import combinations

base = "/home/songyi/Documents/ai_agent_scraper_study/data/大医网"
OUT = os.path.join(base, "03_关联融合/深度挖掘迭代/分析_精修")
os.makedirs(OUT, exist_ok=True)

# 加载药材
herbs = {}
all_herb_names = []

for f in glob.glob(os.path.join(base, "01_来源数据/中药材", "*.json")):
    with open(f, encoding='utf-8') as fh:
        d = json.load(fh)
    n = d.get('名称', '').strip()
    if n:
        herbs[n] = d
        all_herb_names.append(n)
        aliases = d.get('别名', '')
        if aliases:
            for a in re.split(r'[、,，;]', aliases):
                a = a.strip()
                if a:
                    all_herb_names.append(a)

# 去重并按长度降序排列
all_herb_names = list(dict.fromkeys(all_herb_names))
sorted_herb = sorted(all_herb_names, key=len, reverse=True)

# 构建纯药材名集合（不含别名，用于后续过滤）
herb_set = set(herbs.keys())

# 加载方剂
formulas = OrderedDict()
for f in glob.glob(os.path.join(base, "01_来源数据/方剂", "*.json")):
    with open(f, encoding='utf-8') as fh:
        d = json.load(fh)
    name = d.get('名称', os.path.basename(f).rsplit('_', 1)[0])
    formulas[name] = d

print(f"药材: {len(herbs)} 味")
print(f"方剂: {len(formulas)} 首")
print(f"药材名+别名: {len(all_herb_names)} 个")

# 药材提取 - 精修版
def extract_herbs_exact(text):
    """精确匹配：只匹配完整的药材名，且至少2个字符"""
    if not text or len(text) < 3:
        return []
    
    clean = re.sub(r'[（(][^())]*[）)]', '', text)
    clean = re.sub(r'\d+[两钱毫升克毫克]', '', clean)
    
    parts = [p.strip() for p in re.split(r'[、,\;；]', clean) if p.strip()]
    
    found = []
    for p in parts:
        p = p.strip()
        if len(p) < 2:
            continue
        for herb in sorted_herb:
            if herb == p or (len(herb) >= 2 and herb in p):
                if herb not in found:
                    found.append(herb)
                if len(found) > 30:
                    break
        if len(found) > 30:
            break
    return found

def extract_herbs_brute(text):
    """暴力匹配：检查所有药材名是否在文本中出现"""
    if not text or len(text) < 3:
        return []
    found = []
    for herb in sorted_herb:
        if len(herb) >= 2 and herb in text:
            if herb not in found:
                found.append(herb)
    return found

def deduplicate_herbs(herb_list):
    """去重短名：如果同时匹配到'白术'和'术'，只保留'白术'"""
    if not herb_list:
        return herb_list
    
    # 按长度降序去重
    result = []
    for herb in herb_list:
        is_substring = False
        for existing in result:
            if herb in existing or existing in herb:
                is_substring = True
                break
        if not is_substring:
            result.append(herb)
    return result

formula_herbs = OrderedDict()
for fname, data in formulas.items():
    all_text = ''
    for val in data.values():
        if isinstance(val, str):
            all_text += val
        elif isinstance(val, (dict, list)):
            all_text += json.dumps(val, ensure_ascii=False)
    
    # 先尝试精确匹配
    herbs_found = extract_herbs_exact(all_text)
    
    # 如果结果太少，用暴力匹配
    if len(herbs_found) < 3:
        herbs_found = extract_herbs_brute(all_text)
    
    # 去重短名
    herbs_found = deduplicate_herbs(herbs_found)
    
    # 过滤：只保留在药材库中的名
    herbs_found = [h for h in herbs_found if h in herb_set]
    
    formula_herbs[fname] = herbs_found

# 统计
total_links = sum(len(h) for h in formula_herbs.values())
avg = total_links / max(1, len(formula_herbs))
print(f"\n总关联对数: {total_links}")
print(f"平均每方药材数: {avg:.1f}")

# 检查
print("\n精修提取结果 (前10首方):")
for i, (fname, hlist) in enumerate(list(formula_herbs.items())[:10]):
    print(f"  [{i+1}] {fname}: {len(hlist)}味 -> {hlist[:8]}")

# 逆向索引
herb_formulas = defaultdict(set)
for fname, hlist in formula_herbs.items():
    for h in hlist:
        herb_formulas[h].add(fname)

# ========== 1. 高频药材 (精修) ==========
print("\n" + "="*60)
print("  1. 高频药材 Top 30 (精修版)")
print("="*60)

herb_freq = Counter()
for hlist in formula_herbs.values():
    for h in hlist:
        herb_freq[h] += 1

top30 = herb_freq.most_common(30)
for h, c in top30:
    ct = herbs.get(h, {}).get('药材分类', '未知')
    func = herbs.get(h, {}).get('功效作用', {})
    if isinstance(func, dict):
        func_str = func.get('功能', '')
    else:
        func_str = str(func)[:20]
    print(f"  {h:10s} {c:4d}首  [{ct:3s}] {func_str[:25]}")

with open(os.path.join(OUT, "精修_高频药材Top30.json"), 'w', encoding='utf-8') as f:
    json.dump([{"name": h, "count": c} for h, c in top30], f, ensure_ascii=False, indent=2)
print("\n  ✓ 已保存")

# ========== 2. 高频药对 ==========
print("\n" + "="*60)
print("  2. 高频药对 (2味) Top 20")
print("="*60)

pair_freq = Counter()
for hlist in formula_herbs.values():
    u = sorted(list(set(hlist)))
    if len(u) >= 2:
        for combo in combinations(u, 2):
            pair_freq[combo] += 1

top_pairs = [p for p, c in pair_freq.most_common(50) if c >= 5]
for (h1, h2), c in top_pairs[:20]:
    func1_data = herbs.get(h1, {}).get('功效作用', {})
    func2_data = herbs.get(h2, {}).get('功效作用', {})
    func1 = func1_data.get('功能', '')[:10] if isinstance(func1_data, dict) else ''
    func2 = func2_data.get('功能', '')[:10] if isinstance(func2_data, dict) else ''
    print(f"  {h1:>8s} + {h2:>8s}: {c}首方  [{func1[:5]}] [{func2[:5]}]")

with open(os.path.join(OUT, "精修_高频药对Top50.json"), 'w', encoding='utf-8') as f:
    json.dump([{"pair": list(p), "count": c} for p, c in pair_freq.most_common(50)], f, ensure_ascii=False, indent=2)
print("\n  ✓ 已保存")

# ========== 3. 核心药组 ==========
print("\n" + "="*60)
print("  3. 核心药组 (3-5味药) Top 10")
print("="*60)

for n in [3, 4, 5]:
    triple_freq = Counter()
    for hlist in formula_herbs.values():
        u = sorted(list(set(hlist)))
        if len(u) >= n:
            for combo in combinations(u, n):
                triple_freq[combo] += 1
    
    top_combs = [(t, c) for t, c in triple_freq.most_common(10) if c >= 5]
    print(f"\n  Top {n}味药组:")
    for combo, c in top_combs[:10]:
        names = "+".join(combo)
        print(f"    {names:>30s}: {c:4d}首方")

# 合并3-5味药组结果
with open(os.path.join(OUT, "精修_核心药组3-5味_Top10.json"), 'w', encoding='utf-8') as f:
    res = {}
    for n in [3, 4, 5]:
        triple_freq = Counter()
        for hlist in formula_herbs.values():
            u = sorted(list(set(hlist)))
            if len(u) >= n:
                for combo in combinations(u, n):
                    triple_freq[combo] += 1
        key = f"{n}味药组"
        res[key] = [{
            "group": "+".join(t),
            "count": c
        } for t, c in triple_freq.most_common(10) if c >= 5]
    json.dump(res, f, ensure_ascii=False, indent=2)
print("\n  ✓ 已保存")

# ========== 4. 功效关键词网络 ==========
print("\n" + "="*60)
print("  4. 功效关键词网络 (按功效词汇总)")
print("="*60)

# 提取核心功效词
func_keywords = Counter()
func_formulas = defaultdict(set)

for h, data in herbs.items():
    gua = data.get('功效作用', {})
    if isinstance(gua, dict) and gua.get('功能'):
        func_str = gua['功能']
        # 提取核心功效关键词 (按、和、具有等分隔)
        keywords = re.split(r'[、,，具有和的功效]', func_str)
        for kw in keywords:
            kw = kw.strip()
            if len(kw) >= 2 and len(kw) <= 10:
                func_keywords[kw] += 1
                for fname in herb_formulas.get(h, set()):
                    func_formulas[kw].add(fname)

# 按出现方剂数排序
top_funcs = [(kw, len(vx)) for kw, vx in func_formulas.items() if len(vx) >= 10]
top_funcs.sort(key=lambda x: -x[1])

print("  Top 30 功效关键词 (按关联方剂数):")
for kw, count in top_funcs[:30]:
    examples = list(func_formulas[kw])[:3]
    print(f"  {kw:12s} 关联 {count:4d}首方  例: {', '.join(examples[:2])}")

with open(os.path.join(OUT, "精修_功效关键词_Top30.json"), 'w', encoding='utf-8') as f:
    json.dump([{"keyword": kw, "formula_count": count} for kw, count in top_funcs[:30]], f, ensure_ascii=False, indent=2)
print("\n  ✓ 已保存")

# ========== 5. 方剂聚类 ==========
print("\n" + "="*60)
print("  5. 方剂聚类 (基于最高频功效词)")
print("="*60)

def get_top_func_keyword(herb_list):
    """获取方剂的最高频功效关键词"""
    kw_count = Counter()
    for herb in herb_list:
        if herb in herbs:
            gua = herbs[herb].get('功效作用', {})
            if isinstance(gua, dict) and gua.get('功能'):
                keywords = re.split(r'[、,，具有和的功效]', gua['功能'])
                for kw in keywords:
                    kw = kw.strip()
                    if len(kw) >= 2 and len(kw) <= 10:
                        kw_count[kw] += 1
    return kw_count

clusters = defaultdict(list)
for fname, hlist in formula_herbs.items():
    kw_count = get_top_func_keyword(hlist)
    if kw_count:
        top_kw = kw_count.most_common(1)[0][0]
        clusters[top_kw].append(fname)

big_clusters = [(kw, flst) for kw, flst in clusters.items() if len(flst) >= 5]
big_clusters.sort(key=lambda x: -len(x[1]))

print(f"  发现 {len(clusters)} 个功效簇")
print(f"\n  Top 20 大功效簇:")
for kw, flst in big_clusters[:20]:
    sample_herbs = set()
    for fname in flst[:5]:
        for h in formula_herbs[fname]:
            sample_herbs.add(h)
    print(f"  '{kw}' ({len(flst)}首方): 代表药材 -> {', '.join(sorted(sample_herbs)[:8])}")

with open(os.path.join(OUT, "精修_方剂聚类_Top20.json"), 'w', encoding='utf-8') as f:
    json.dump([{"cluster": kw, "count": len(flst), "formulas": flst[:20]} for kw, flst in big_clusters[:20]], f, ensure_ascii=False, indent=2)
print("\n  ✓ 已保存")

# ========== 6. 配伍禁忌验证 (支持别名匹配) ==========
print("\n" + "="*60)
print("  6. 配伍禁忌验证 (十八反/十九畏)")
print("="*60)

# 十八反
fan = {
    '甘草': ['大戟', '芫花', '甘遂', '京大戟', '红大戟'],
    '乌头': ['半夏', '瓜蒌', '天花粉', '贝母', '平贝母', '川贝母', '浙贝母', '白蔹', '白及'],
    '藜芦': ['人参', '党参', '丹参', '玄参', '沙参', '苦参', '细辛', '白芍', '赤芍'],
}
# 十九畏
wei = {
    '硫黄': ['朴硝', '芒硝', '牙硝'],
    '水银': ['铅丹', '砒霜'],
    '巴豆': ['牵牛', '牵牛子'],
    '丁香': ['郁金'],
    '人参': ['五灵脂'],
    '肉桂': ['石脂', '赤石脂', '黑锡丹'],
    '半夏': ['羊脂', '羊角'],
    '厚朴': ['硝石', '滑石'],
}

# 构建别名映射
alias_to_canon = {}
for h, data in herbs.items():
    aliases = data.get('别名', '')
    if aliases:
        for a in re.split(r'[、,，;]', aliases):
            a = a.strip()
            if a:
                alias_to_canon[a] = h

# 扩展禁忌药材：加入别名
def expand_contraindicated(names_list):
    """扩展：将别名也加入匹配列表"""
    result = list(names_list)
    for name in names_list:
        # 检查是否在别名映射中
        if name in alias_to_canon:
            result.append(alias_to_canon[name])
        # 也检查是否在药材名中
        if name in herb_set:
            pass  # 已经在药材名中
    # 反过来：查找所有别名
    for alt, canon in alias_to_canon.items():
        if canon in names_list:
            result.append(alt)
    return result

contra = defaultdict(int)
contra_details = defaultdict(list)

for fname, hlist in formula_herbs.items():
    # 扩展禁忌药材
    all_fan = {}
    for k, v in fan.items():
        all_fan[k] = expand_contraindicated(v)
    all_wei = {}
    for k, v in wei.items():
        all_wei[k] = expand_contraindicated(v)
    
    total_all = {**all_fan, **all_wei}
    
    for herb, contra_set in total_all.items():
        if herb in hlist:
            for ch in contra_set:
                if ch in hlist:
                    contra[(herb, ch)] += 1
                    contra_details[(herb, ch)].append(fname)
                    
            # 也用别名查找
            for alt, canon in alias_to_canon.items():
                if alt in hlist and canon == herb:
                    for ch in contra_set:
                        if ch in hlist:
                            contra[(alt, ch)] += 1
                            contra_details[(alt, ch)].append(fname)

top_contra = sorted(contra.items(), key=lambda x: -x[1])[:15]
if top_contra:
    print("\n  发现配伍禁忌:")
    for (h1, h2), count in top_contra:
        print(f"  ⚠️ {h1} + {h2}: {count}首方")
        print(f"     方剂: {', '.join(contra_details[(h1, h2)][:3])}")
else:
    print("\n  未发现有直接配伍禁忌的方剂 (数据中不存在十八反/十九畏配对)")
    
    # 检查我们找到的药材是否在方剂中出现
    print("\n  验证：检索禁忌相关药材在方剂中的出现情况:")
    search_names = list(fan.keys()) + list(wei.keys())
    search_names.extend(['大戟', '芫花', '半夏', '瓜蒌', '天花粉', '贝母', '人参', '党参', '丹参'])
    for sn in search_names:
        if sn in herb_formulas:
            print(f"    {sn}: 出现在 {len(herb_formulas[sn])} 首方中")
            print(f"       方剂: {', '.join(list(herb_formulas[sn])[:3])}")

with open(os.path.join(OUT, "精修_配伍禁忌验证.json"), 'w', encoding='utf-8') as f:
    json.dump({
        "十八反": fan,
        "十九畏": wei,
        "发现": {f"{h1}+{h2}": c for (h1, h2), c in top_contra},
        "详情": {f"{h1}+{h2}": contra_details[(h1, h2)] for (h1, h2) in top_contra}
    }, f, ensure_ascii=False, indent=2)
print("\n  ✓ 已保存")

# ========== 7. 方剂相似度 (精修) ==========
print("\n" + "="*60)
print("  7. 方剂相似度 (修正版)")
print("="*60)

def jaccard_fixed(s1, s2):
    """Jaccard相似度，空集=0"""
    if not s1 or not s2:
        return 0.0
    intersection = len(s1 & s2)
    if intersection == 0:
        return 0.0  # 没有共享药材，相似度为0
    union = len(s1 | s2)
    return intersection / union

similar_pairs = []
formula_list = list(formulas.keys())
total_formulas = len(formula_list)

print(f"  正在计算 {total_formulas} 首方剂的相似度...")
count = 0
for i in range(len(formula_list)):
    f1 = formula_list[i]
    h1 = formula_herbs.get(f1, [])
    if not h1:
        continue
    for j in range(i + 1, len(formula_list)):
        f2 = formula_list[j]
        h2 = formula_herbs.get(f2, [])
        if not h2:
            continue
        s1 = set(h1)
        s2 = set(h2)
        sim = jaccard_fixed(s1, s2)
        if sim >= 0.3:
            shared = s1 & s2
            similar_pairs.append((f1, f2, sim, len(shared)))
        count += 1

print(f"  已检查 {count} 对")
similar_pairs.sort(key=lambda x: (-x[2], -x[3]))
top_similar = similar_pairs[:20]

print(f"\n  找到 {len(similar_pairs)} 对相似方剂 (Jaccard >= 0.3)")
print(f"\n  Top 15 相似方剂对:")
for f1, f2, sim, shared_count in top_similar[:15]:
    shared_herbs = set(formula_herbs.get(f1, [])) & set(formula_herbs.get(f2, []))
    print(f"  {f1:30s} + {f2:30s} -> {sim:.3f}  (共享{shared_count}味: {', '.join(list(shared_herbs)[:3])})")

with open(os.path.join(OUT, "精修_方剂相似度_Top20.json"), 'w', encoding='utf-8') as f:
    json.dump([{"pair1": f1, "pair2": f2, "similarity": round(sim, 4), "shared_count": sc, "shared_herbs": list(formula_herbs.get(f1, []) & formula_herbs.get(f2, []))[:5]} for f1, f2, sim, sc in top_similar], f, ensure_ascii=False, indent=2)
print("\n  ✓ 已保存")

# ========== 8. 生成汇总报告 ==========
print("\n" + "="*60)
print("  8. 汇总报告")
print("="*60)

summary = OrderedDict()
summary["标题"] = "大医网 方剂-中药材 深度挖掘报告 (精修版)"
summary["数据规模"] = {
    "药材": len(herbs),
    "方剂": len(formulas),
    "总关联对数": total_links,
    "平均每方药材数": round(total_links / max(1, len(formula_herbs)), 2),
    "成功提取方剂": sum(1 for h in formula_herbs.values() if h),
}

# 高频药材前20
summary["高频药材_Top20"] = [{"name": h, "count": c} for h, c in top30[:20]]

# 高频药对前15
summary["高频药对_Top15"] = [{"pair": list(p), "count": c} for p, c in pair_freq.most_common(15)]

# 功效关键词前15
summary["功效关键词_Top15"] = [{"keyword": kw, "formula_count": count} for kw, count in top_funcs[:15]]

# 方剂聚类
summary["方剂聚类"] = [{"cluster": kw, "count": len(flst), "representative": list(flst[:3])} for kw, flst in big_clusters[:10]]

# 配伍禁忌
summary["配伍禁忌"] = {
    "十八反": fan,
    "十九畏": wei,
    "发现": {f"{h1}+{h2}": c for (h1, h2), c in top_contra},
}

# 方剂相似度
summary["方剂相似度_Top10"] = [{"similarity": round(sim, 4), "shared": sc, "pair": f"{f1} + {f2}"} for f1, f2, sim, sc in top_similar[:10]]

summary["输出目录"] = os.path.abspath(OUT)

with open(os.path.join(OUT, "精修_总汇总报告.json"), 'w', encoding='utf-8') as f:
    json.dump(summary, f, ensure_ascii=False, indent=2)

print("\n  精修深度挖掘完成！")
print(f"\n  输出目录: {os.path.abspath(OUT)}")
print("\n  文件列表:")
for fn in sorted(os.listdir(OUT)):
    fp = os.path.join(OUT, fn)
    size = os.path.getsize(fp)
    print(f"    {fn:40s} {size:>10,} B")