"""Generate H&E descriptions with GPT-6.1 Sol on Azure OpenAI for images that have no real text (plan A:
one FIXED 3-shot prefix of real PathGen-1.6M captions -> prompt-cached; short caption; low reasoning effort).

Targets (images on disk without paired text):
  ge_patches : PathGen-1.6M gene-expression patches (brain*, gliomas*, idc*, saromas*, skin*)   1640
  slide_rep  : PathGen-1.6M/slide_report/image_set/*.png                                           1
  orion      : ORION H&E tiles (he/*.jpeg)                                                         6335
  hemit      : HEMIT H&E inputs (he/*/input/*)                                                     55
Config (env): AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_KEY, AZURE_OPENAI_DEPLOYMENT,
              optional AZURE_OPENAI_API_VERSION (if unset the v1 endpoint <endpoint>/openai/v1/ is used)
Usage:  python 27_azure_caption.py --dry                # list targets + 3-shot prompt, no API call
        python 27_azure_caption.py --pilot 20           # 20 images, prints measured cost / image
        python 27_azure_caption.py --budget 40          # everything; stops when the USD budget is reached
Output: cache/gen_text/captions.jsonl (resumable; one line per image: set, key, path, caption, usage)
Generated text is SYNTHETIC (seen only the H&E) -> never use it as test text; tag ttype='synthetic'.
"""
import argparse
import base64
import concurrent.futures as cf
import glob
import io
import json
import os
import random
import threading
import time

import numpy as np
import pandas as pd
from PIL import Image, PngImagePlugin

from common import CACHE, ROOT

PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024
P = os.path.join(ROOT, "..")
OUT = os.path.join(CACHE, "gen_text")
os.makedirs(OUT, exist_ok=True)
PRICE = {"in": 2.0, "cached": 0.10, "out": 10.0}  # USD / 1M tokens, GPT-6.1 Sol standard
MAXPX = 512  # images are downscaled to <=512 px (one vision tile) to keep the token count low

ap = argparse.ArgumentParser()
ap.add_argument("--sets", nargs="+", default=["ge_patches", "slide_rep", "orion", "hemit"],
                help="tcga_notext = TCGA images left without real text after 40_clean_text.py; dlpfc / idc = external test sets (all spots); tcga_calib = 1000 TCGA patches that HAVE real text (held-out patients of 31_), for calibrating synthetic vs real text")
ap.add_argument("--pilot", type=int, default=0)
ap.add_argument("--budget", type=float, default=40.0)
ap.add_argument("--workers", type=int, default=8)
ap.add_argument("--effort", default="low")
ap.add_argument("--dry", action="store_true")
args = ap.parse_args()


def targets():
    rows = []
    ge = ["brain_image", "brain_train_image", "brain_trainnew_image", "gliomas_test_image_new", "gliomas_train_image",
          "idc_image", "idc_train_image", "saromas_test_image_new", "saromas_train_image_new", "skin_image", "skin_train_image"]
    src = {"ge_patches": [f for d in ge for f in glob.glob(os.path.join(P, "PathGen-1.6M", d, "*"))],
           "slide_rep": glob.glob(os.path.join(P, "PathGen-1.6M", "slide_report", "image_set", "*")),
           "orion": glob.glob(os.path.join(P, "ORION", "ORIONCRC_dataset_tile_20x", "he", "*.jpeg")),
           "hemit": glob.glob(os.path.join(P, "HEMIT", "he", "*", "input", "*"))}
    if "tcga_calib" in args.sets:  # same held-out image set as 31_virtual_validation.py (TCGA), random 1000
        I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
        T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
        seen = set(pd.read_csv(os.path.join(CACHE, "pathgen_meta.csv")).query("split in ['train', 'val']").patient)
        ev = sorted(set(T.img_idx) & set(I.index[~I.patient.isin(seen)]))
        src["tcga_calib"] = list(I.image_path.values[np.random.default_rng(0).choice(ev, 1000, replace=False)])
    if "tcga_notext" in args.sets:  # TCGA images whose real texts are all junk (40_clean_text.py)
        src["tcga_notext"] = list(pd.read_csv(os.path.join(CACHE, "notext_meta.csv")).image_path)
    for s in ("dlpfc", "idc"):  # external test sets (36_external_prep.py): the same 224-px patches as the anchor
        if s in args.sets:      # (0.5 um/px; = UNI2 / MIPHEI input) -> path = "npy:<set>:<row>"
            n = len(np.load(os.path.join(CACHE, "external", EXT[s]), mmap_mode="r"))
            rows += [{"set": s, "key": f"external/{s}/{i}", "path": f"npy:{s}:{i}"} for i in range(n)]
    for s in args.sets:
        if s in EXT:
            continue
        for f in sorted(src[s]):
            if f.lower().endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff")):
                rows.append({"set": s, "key": os.path.relpath(f, P), "path": f})
    return pd.DataFrame(rows)


EXT = {"dlpfc": "dlpfc_he_05um.npy", "idc": "idc_he.npy"}


def b64(path):
    if path.startswith("npy:"):
        _, st, i = path.split(":")
        im = Image.fromarray(np.asarray(np.load(os.path.join(CACHE, "external", EXT[st]), mmap_mode="r")[int(i)]))
    else:
        im = Image.open(path).convert("RGB")
    im.thumbnail((MAXPX, MAXPX))
    buf = io.BytesIO()
    im.save(buf, format="JPEG", quality=90)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


def few_shot():
    """3 fixed real PathGen-1.6M captions (40-60 words), different slides; same for every request -> cacheable."""
    T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
    I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
    T = T[(T.ttype == "caption") & T.source.str.contains("pathgen1.6m")]
    w = T.caption.str.split().str.len()
    T = T[(w >= 40) & (w <= 60)]
    rng = random.Random(0)
    picks, seen = [], set()
    for i in rng.sample(list(T.index), len(T)):
        wsi = I.wsi_id[T.img_idx[i]]
        if wsi[:12] in seen:
            continue
        try:
            img = b64(I.image_path[T.img_idx[i]])
        except Exception:
            continue
        seen.add(wsi[:12]); picks.append((img, T.caption[i], T.key[i]))
        if len(picks) == 3:
            return picks


SYSTEM = ("You are a pathologist writing a short description of one H&E-stained histopathology image patch. "
          "Describe only what is visible: tissue type and architecture, cell types and their density, nuclear "
          "features, stroma, inflammation, vessels, necrosis or mitoses when present. Follow the style and length "
          "of the examples (one paragraph, about 50 words). Do not mention the examples, staining quality, "
          "magnification or patient information, and do not invent a diagnosis that the image does not support.")


def messages(shots, target_img):
    m = [{"role": "system", "content": SYSTEM}]
    for img, cap, _ in shots:
        m.append({"role": "user", "content": [{"type": "text", "text": "Describe this H&E patch."},
                                              {"type": "image_url", "image_url": {"url": img, "detail": "auto"}}]})
        m.append({"role": "assistant", "content": cap})
    m.append({"role": "user", "content": [{"type": "text", "text": "Describe this H&E patch."},
                                          {"type": "image_url", "image_url": {"url": target_img, "detail": "auto"}}]})
    return m


def client():
    from openai import AzureOpenAI, OpenAI
    ep, key = os.environ["AZURE_OPENAI_ENDPOINT"].rstrip("/"), os.environ["AZURE_OPENAI_API_KEY"]
    if os.environ.get("AZURE_OPENAI_API_VERSION"):
        return AzureOpenAI(azure_endpoint=ep, api_key=key, api_version=os.environ["AZURE_OPENAI_API_VERSION"])
    return OpenAI(base_url=ep + "/openai/v1/", api_key=key)


def cost(u):
    cached = (u.get("prompt_tokens_details") or {}).get("cached_tokens") or 0
    return ((u["prompt_tokens"] - cached) * PRICE["in"] + cached * PRICE["cached"] + u["completion_tokens"] * PRICE["out"]) / 1e6


D = targets()
shots = few_shot()
print(f"targets: {len(D)}  " + json.dumps(D.set.value_counts().to_dict()))
print("3-shot examples (PathGen-1.6M keys):", [k for _, _, k in shots])
fout = os.path.join(OUT, "captions.jsonl")
done = set()
spent = 0.0
if os.path.exists(fout):
    for line in open(fout):
        r = json.loads(line)
        done.add(r["key"]); spent += r.get("usd", 0)
todo = D[~D.key.isin(done)]
if args.pilot:
    todo = todo.groupby("set", group_keys=False).apply(lambda g: g.sample(min(len(g), max(1, args.pilot // D.set.nunique())), random_state=0))
print(f"already done {len(done)} (USD {spent:.2f}); to do now {len(todo)}")
json.dump({"system": SYSTEM, "shots": [{"key": k, "caption": c} for _, c, k in shots], "maxpx": MAXPX,
           "effort": args.effort, "deployment": os.environ.get("AZURE_OPENAI_DEPLOYMENT")},
          open(os.path.join(OUT, "prompt.json"), "w"), indent=1)
if args.dry:
    for _, c, k in shots:
        print(f"- {k}: {c}")
    raise SystemExit

cl, dep = client(), os.environ["AZURE_OPENAI_DEPLOYMENT"]
lock = threading.Lock()
stop = threading.Event()


def run(row):
    if stop.is_set():
        return
    for attempt in range(6):
        try:
            r = cl.chat.completions.create(model=dep, messages=messages(shots, b64(row.path)),
                                           reasoning_effort=args.effort, max_completion_tokens=2000,
                                           prompt_cache_key="he-caption-3shot-v1")
            u = r.usage.model_dump()
            rec = {"set": row.set, "key": row.key, "path": row.path, "caption": r.choices[0].message.content.strip(),
                   "finish": r.choices[0].finish_reason, "usage": u, "usd": cost(u), "ttype": "synthetic"}
            break
        except Exception as e:  # rate limits / transient errors -> exponential backoff
            err = repr(e)[:300]
            if "content_filter" in err or attempt == 5:
                rec = {"set": row.set, "key": row.key, "path": row.path, "caption": None, "error": err, "usd": 0}
                break
            time.sleep(2 ** attempt * 5)
    global spent
    with lock:
        with open(fout, "a") as f:
            f.write(json.dumps(rec) + "\n")
        spent += rec["usd"]
        if spent >= args.budget:
            stop.set()


t0 = time.time()
with cf.ThreadPoolExecutor(args.workers) as ex:
    for i, _ in enumerate(ex.map(run, [r for r in todo.itertuples()])):
        if i % 200 == 0:
            print(f"{i}/{len(todo)}  USD {spent:.2f}  {time.time() - t0:.0f}s", flush=True)
R = pd.read_json(fout, lines=True)
ok = R[R.caption.notna()]
print(f"done {len(ok)} ok / {len(R)} total; failures {R.caption.isna().sum()}; USD spent {R.usd.sum():.2f}; "
      f"USD/image {ok.usd.mean():.4f} -> projected for all {len(D)}: {ok.usd.mean() * len(D):.0f}")
if len(ok):
    u = pd.json_normalize(ok.usage)
    print(u.filter(regex="prompt_tokens$|completion_tokens$|cached_tokens|reasoning_tokens").mean().round(0).to_string())
if stop.is_set():
    print(f"STOPPED: budget {args.budget} USD reached")