"""Unified pool of ALL H&E image-text pairs in PathGen-1.6M/ (every text-image source of PAIR_SUMMARY.md; all TCGA
patches). Images are de-duplicated by patch key <wsi>_<x>_<y> (the same patch recurs across sources); one image can
carry several texts. Text types:
  caption  : PathGen-1.6M captions, figure_caption_summary / summary_testing / subsample_summary, slide_report desc,
             sample_data/output txt, output_subsample txt, path_reason_data caption
  summary  : path_summary_train(.json, _5w.json) one-paragraph trainee summaries
  qa_open  : open-ended Q/A -> assistant answers joined (PathGen-Instruct OPEN, output_subsample json,
             qa_pathinstruct_withreason*.csv answers, path_reason_data question+answer)
  qa_close : multiple-choice Q/A -> question stem + text of the correct option (PathGen-Instruct CLOSE,
             qa_instruct_merged_onlyqaforrl, output_reason json, multiple-choice csv rows)
Outputs
  cache/allimg_meta.csv : unique images  (idx, key, wsi_id, patient, image_path)            -> image features
  cache/alltxt_meta.csv : unique (image, text) pairs (idx, img_idx, caption, ttype, source)   -> text features
(the `caption` column holds the text so the extractors' --set option works unchanged)
Not used: pathor1_output*.csv (hashed non-TCGA image names, images absent).
"""
import ast
import glob
import json
import os
import re

import pandas as pd

from common import CACHE, ROOT

PG = os.path.join(ROOT, "..", "PathGen-1.6M")
img, txt = {}, []          # key -> image path ; list of (key, text, ttype, source)
KEY = re.compile(r"(TCGA-[A-Z0-9-]+\.[A-Za-z0-9-]+)_(\d+)_(\d+)")


def key_of(s):
    m = KEY.search(os.path.basename(str(s)))
    return f"{m.group(1)}_{m.group(2)}_{m.group(3)}" if m else None


def add_img(key, path):
    if key and key not in img and os.path.exists(path):
        img[key] = path


def add_txt(key, text, ttype, source):
    text = re.sub(r"\s+", " ", str(text).replace("<image>", " ")).strip()
    if key and text and text.lower() != "nan":
        txt.append((key, text, ttype, source))


def close_text(question, answer):
    """multiple-choice -> 'question stem + correct option text'"""
    q = question.replace("<image>", " ")
    opts = dict(re.findall(r"(?:^|\n)\s*\(?([A-F])[\).:]\s*([^\n]+)", q))
    stem = re.split(r"\n\s*\(?A[\).:]", q)[0].strip()
    a = answer.strip()
    letter = re.match(r"^\(?([A-F])\)?(?:[\).:\s]|$)", a)
    if letter and letter.group(1) in opts:
        a = opts[letter.group(1)]
    elif letter and len(a) <= 3:
        return None
    else:
        a = re.sub(r"^\(?[A-F]\)\s*", "", a)
    return f"{stem} {a.strip()}"


def conv_text(conv):
    """returns (ttype, text) for a LLaVA-style conversation list"""
    gpt = [c["value"].strip() for c in conv if c["from"] == "gpt"]
    if len(gpt) == 1 and re.fullmatch(r"\(?[A-F]\)?\.?", gpt[0]):
        human = [c["value"] for c in conv if c["from"] == "human"][0]
        t = close_text(human, gpt[0])
        return ("qa_close", t) if t else (None, None)
    return "qa_open", " ".join(gpt)


# 1. PathGen-1.6M captions (Data/images, 672 px)
for x in json.load(open(os.path.join(PG, "Data", "PathGen-1.6M.json"))):
    k = f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}"
    add_img(k, os.path.join(PG, "Data", "images", k + ".png"))
    add_txt(k, x["caption"], "caption", "pathgen1.6m")
# 2. PathGen-Instruct (OPEN + CLOSE)
idx = pd.read_csv(os.path.join(PG, "PathGen-Instruct.json.image_index.csv")).drop_duplicates("key").set_index("key")
for x in json.load(open(os.path.join(PG, "PathGen-Instruct.json"))):
    k = f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}"
    if k in idx.index:
        add_img(k, os.path.join(PG, idx.loc[k, "local_image"]))
    tt, t = conv_text(x["conversations"])
    if tt:
        add_txt(k, t, tt, "pathgen-instruct")
# 3. qa_instruct_merged_onlyqaforrl
for x in json.load(open(os.path.join(PG, "qa_instruct_merged_onlyqaforrl.json"))):
    k = x.get("path_list") or f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}"
    for c in glob.glob(os.path.join(PG, "qa_instruct_merged_images", k + ".*")):
        if c.endswith((".png", ".jpg")):
            add_img(k, c)
    tt, t = conv_text(x["conversations"])
    if tt:
        add_txt(k, t, tt, "qa_instruct_merged")
# 4. summaries
for f in ["path_summary_train.json", "path_summary_train_5w.json"]:
    for x in json.load(open(os.path.join(PG, f))):
        k = key_of(x["images"][0])
        add_img(k, os.path.join(PG, x["images"][0]))
        add_txt(k, [m["content"] for m in x["messages"] if m["role"] == "assistant"][0], "summary", f.replace(".json", ""))
# 5. output_reason (png + json conversation), output_subsample (png + txt caption / json conversation), sample_data
for d, src in [("output_reason", "output_reason"), ("output_subsample", "output_subsample"), ("sample_data/output", "sample_data")]:
    for pth in glob.glob(os.path.join(PG, d, "*.png")):
        k = key_of(pth)
        add_img(k, pth)
        base = pth[:-4]
        if os.path.exists(base + ".txt"):
            add_txt(k, open(base + ".txt").read(), "caption", src)
        if os.path.exists(base + ".json"):
            tt, t = conv_text(json.load(open(base + ".json")))
            if tt:
                add_txt(k, t, tt, src)
# 6. qa_pathinstruct_withreason*.csv (images in output_reason/)
for f in sorted(glob.glob(os.path.join(PG, "qa_pathinstruct_withreason*.csv"))):
    for _, r in pd.read_csv(f).iterrows():
        k = key_of(r.get("path_list", r.get("imagep")))
        add_img(k, os.path.join(PG, str(r["imagep"]).lstrip("./")))
        t = close_text(r["question"], str(r["answer"])) if re.search(r"\n\s*\(?A[\).:]", str(r["question"])) else r["answer"]
        if t:
            add_txt(k, t, "qa_close" if t is not r["answer"] else "qa_open", "qa_pathinstruct_withreason")
# 7. path_reason_data (caption + diagnostic Q/A)
for f in sorted(glob.glob(os.path.join(PG, "path_reason_data_3600*.csv"))):
    for _, r in pd.read_csv(f).iterrows():
        k = key_of(r["patch_id"])
        for c in [os.path.join(PG, str(r["patch_id"]).lstrip("./")), os.path.join(PG, "output_subsample", os.path.basename(str(r["patch_id"])))]:
            add_img(k, c)
        add_txt(k, r["caption"], "caption", "path_reason_data")
        add_txt(k, f"{r['question']} {r['answer']}", "qa_open", "path_reason_data")
# 8. figure_caption_summary (summary_testing_data / subsample_summary_sentence), slide_report
for f in ["summary_testing_data.csv", "subsample_summary_sentence.csv"]:
    for _, r in pd.read_csv(os.path.join(PG, f)).iterrows():
        x, y = ast.literal_eval(r["position"])
        k = f"{r['wsi_id']}_{x}_{y}"
        add_img(k, os.path.join(PG, "figure_caption_summary", k + ".png"))
        add_txt(k, r["caption"], "caption", "figure_caption_summary")
for f in glob.glob(os.path.join(PG, "slide_report", "*", "*_desc.csv")):
    k = key_of(f.replace("_desc.csv", ""))
    add_img(k, f.replace("_desc.csv", ".png"))
    for _, r in pd.read_csv(f).iterrows():
        add_txt(k, r["caption"], "caption", "slide_report")

T = pd.DataFrame(txt, columns=["key", "caption", "ttype", "source"])
n_raw = len(T)
T = T[T.key.isin(img)]
n_noimg = n_raw - len(T)
T = T.groupby(["key", "caption"], as_index=False).agg(ttype=("ttype", "first"), source=("source", lambda s: "|".join(sorted(set(s)))))
I = pd.DataFrame({"key": sorted(T.key.unique())})
I["image_path"] = I.key.map(img)
I["wsi_id"] = I.key.str.rsplit("_", n=2).str[0]
I["patient"] = I.wsi_id.str[:12]
I["caption"] = ""
T["img_idx"] = T.key.map(dict(zip(I.key, I.index)))
I.to_csv(os.path.join(CACHE, "allimg_meta.csv"), index_label="idx")
T[["img_idx", "key", "caption", "ttype", "source"]].to_csv(os.path.join(CACHE, "alltxt_meta.csv"), index_label="idx")
print(f"texts: {n_raw} raw, {n_noimg} without an image on disk, {len(T)} unique (image, text) pairs")
print(f"images: {len(I)} unique patches, {I.wsi_id.nunique()} slides, {I.patient.nunique()} patients")
print(T.ttype.value_counts().to_string())
print(T.source.str.split("|").explode().value_counts().to_string())
print("texts per image:", T.groupby("img_idx").size().describe().round(2).to_dict())
print("words per text by type:", T.assign(w=T.caption.str.split().str.len()).groupby("ttype").w.median().to_dict())