import json, re, importlib.util, sys, os

BASE = os.path.dirname(os.path.abspath(__file__))

def load_batch(name):
    spec = importlib.util.spec_from_file_location(name, os.path.join(BASE, name + ".py"))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod.DATA

ALL = {}
for b in ["batch1", "batch2", "batch3", "batch4", "batch5", "batch6"]:
    d = load_batch(b)
    for k, v in d.items():
        if k in ALL:
            print("DUPLICATE ARTICLE ID", k)
        ALL[k] = v

with open(os.path.join(BASE, "mdm-06.json"), encoding="utf-8") as f:
    original = json.load(f)

print("Total articles in input:", len(original))
print("Total articles drafted:", len(ALL))

missing = [a["article_id"] for a in original if a["article_id"] not in ALL]
if missing:
    print("MISSING article ids:", missing)

output = []
ALLOWED_TITLES = ["Clinical paths", "Diagnosis", "Management", "Prognosis and outcome"]

fail_count = 0
prognosis_count = 0

for art in original:
    aid = art["article_id"]
    if aid not in ALL:
        continue
    new_sections = ALL[aid]
    out_art = {"article_id": aid, "article_title": art["article_title"], "sections": new_sections}
    output.append(out_art)

    # verification
    before = " ".join(s["content"] for s in art["sections"])
    after = " ".join(s["content"] for s in new_sections)

    before_refs = sorted(re.findall(r"\[\[(\d+)\|", before))
    after_refs = sorted(re.findall(r"\[\[(\d+)\|", after))
    before_nums = sorted(re.findall(r"\d+(?:\.\d+)?", before))
    after_nums = sorted(re.findall(r"\d+(?:\.\d+)?", after))
    ratio = len(after) / len(before) if before else 1.0

    titles = [s["title"] for s in new_sections]
    titles_ok = (
        all(t in ALLOWED_TITLES for t in titles)
        and len(set(titles)) == len(titles)
        and titles[0] == "Clinical paths"
        and "Diagnosis" in titles
        and "Management" in titles
        and titles.index("Clinical paths") < titles.index("Diagnosis") < titles.index("Management")
        and (len(titles) == 3 or (len(titles) == 4 and titles[3] == "Prognosis and outcome"))
    )
    if "Prognosis and outcome" in titles:
        prognosis_count += 1

    problems = []
    if before_refs != after_refs:
        problems.append(f"REF MISMATCH before={before_refs} after={after_refs}")
    if before_nums != after_nums:
        # show the diff
        from collections import Counter
        cb, ca = Counter(before_nums), Counter(after_nums)
        diff = {}
        for k in set(cb) | set(ca):
            if cb[k] != ca[k]:
                diff[k] = (cb[k], ca[k])
        problems.append(f"NUM MISMATCH diff(before,after)={diff}")
    if not (0.8 <= ratio <= 1.3):
        problems.append(f"LENGTH RATIO {ratio:.3f} (before={len(before)} after={len(after)})")
    if not titles_ok:
        problems.append(f"TITLES BAD {titles}")

    if problems:
        fail_count += 1
        print(f"=== Article {aid} ({art['article_title']}) FAILED ===")
        for p in problems:
            print("  -", p)

print()
print("Articles restructured:", len(output))
print("Articles with Prognosis and outcome section:", prognosis_count)
print("Articles failing verification:", fail_count)

with open(os.path.join(BASE, "mdm-06.done.json"), "w", encoding="utf-8") as f:
    json.dump(output, f, ensure_ascii=False, indent=2)

print("Wrote output to mdm-06.done.json")