# -*- coding: utf-8 -*-
"""
训练数据准备: train.xlsx(Sheet1, 1800条) -> 微调用 Alpaca 格式
输出: train.jsonl(1700) / val.jsonl(100) / train_alpaca.json
"""
import pandas as pd, json, random, os, numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
df = pd.read_excel(os.path.join(ROOT, "o_data/train.xlsx"), sheet_name="Sheet1")
cand = json.load(open(os.path.join(ROOT, "icd_knowledge_base/task_candidates.json"), encoding="utf-8"))

FIELDS = ["主诉","现病史","既往史","入院诊断","诊疗经过","出院情况","出院医嘱","术前诊断","术中诊断","手术经过"]
DXM, DXO, OPM, OPO = "主要疾病（诊断）编码", "其他疾病（诊断）编码", "主要手术编码", "其他手术编码"

CAND_DX = "\n".join(f"{c['code']} {c['name']}" for c in cand['main_diagnosis_candidates'])
CAND_OP = "\n".join(f"{c['code']} {c['name']}" for c in cand['main_procedure_candidates'])
INSTR = (
    "你是三甲医院资深病案首页ICD编码员。根据电子病历，输出主要诊断编码、其他诊断编码、主要手术编码、其他手术编码，"
    "四者用|分隔，其他诊断/手术内部用;分隔，格式：主诊断|其他诊断1;其他诊断2|主手术|其他手术1;其他手术2。\n"
    f"主要诊断必须从以下18个中选1个：\n{CAND_DX}\n主要手术必须从以下16个中选1个：\n{CAND_OP}"
)

def clean(s):
    if pd.isna(s): return ""
    return ";".join(x.strip() for x in str(s).replace("；", ";").split(";") if x.strip())
def record(row):
    return "\n".join(f"【{f}】{str(row[f]).strip()}" for f in FIELDS if pd.notna(row[f]) and str(row[f]).strip())
def output(row):
    return f"{str(row[DXM]).strip()}|{clean(row[DXO])}|{str(row[OPM]).strip()}|{clean(row[OPO])}"

samples = [{"instruction": INSTR, "input": record(df.iloc[i]), "output": output(df.iloc[i])} for i in range(len(df))]
random.seed(42); random.shuffle(samples)
val, train = samples[:100], samples[100:]

out = os.path.join(ROOT, "train_data")
for name, data in [("train.jsonl", train), ("val.jsonl", val)]:
    with open(os.path.join(out, name), "w", encoding="utf-8") as f:
        for s in data: f.write(json.dumps(s, ensure_ascii=False) + "\n")
json.dump(train, open(os.path.join(out, "train_alpaca.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=1)
print(f"train {len(train)} / val {len(val)} 已生成")