"""Templated per-tile text for ORION (no real text is paired with ORION tiles).

Each tile gets 4 paraphrases describing cellularity and cell-type composition derived from the
marker-positive nuclei counts (fraction of all nuclei) — i.e. the text carries a coarse, discretised
version of the marker information. Kept short (< 77 CLIP tokens).
Output: cache/text.json  {tile_id: [4 strings]}
"""
import json
import os

import numpy as np
import pandas as pd

from common import CACHE, ROOT, load_all_tiles

df = load_all_tiles()
stats = pd.read_csv(os.path.join(ROOT, "results", "tile_stats.csv")).set_index("tile_id")
df["n_nuclei"] = stats.loc[df.tile_id, "n_nuclei"].values

CELLS = [  # (count column, phrase)
    ("Pan-CK_count", "Pan-CK+ epithelial cells"), ("SMA_count", "SMA+ stromal cells"),
    ("CD45_count", "CD45+ immune cells"), ("CD3e_count", "CD3e+ T cells"), ("CD8a_count", "CD8a+ cytotoxic T cells"),
    ("CD4_count", "CD4+ helper T cells"), ("FOXP3_count", "FOXP3+ regulatory T cells"), ("CD20_count", "CD20+ B cells"),
    ("CD68_count", "CD68+ macrophages"), ("CD163_count", "CD163+ macrophages"), ("CD31_count", "CD31+ endothelial cells"),
    ("Ki67_count", "Ki67+ proliferating cells"), ("PD-L1_count", "PD-L1+ cells"), ("PD-1_count", "PD-1+ cells"),
]
LEVEL_WORDS = [  # per paraphrase: (few, moderate, abundant)
    ("few", "moderate", "abundant"), ("scattered", "some", "many"),
    ("rare", "a moderate number of", "numerous"), ("sparse", "intermediate", "dense"),
]
DENSITY = [("acellular", "sparsely cellular", "moderately cellular", "highly cellular"),
           ("background", "low-density", "medium-density", "dense"),
           ("nearly empty", "sparse", "moderately dense", "crowded"),
           ("mostly non-tissue", "hypocellular", "cellular", "hypercellular")]
OPEN = ["H&E tile, {d} region with {n} nuclei.", "A {d} tissue patch ({n} nuclei).",
        "Colorectal cancer tissue, {d}, about {n} nuclei.", "{D} area containing {n} nuclei."]


def level(frac):
    return None if frac < 0.02 else 0 if frac < 0.08 else 1 if frac < 0.25 else 2


def describe(r, v, rng):
    n = int(r.n_nuclei)
    di = 0 if n < 20 else 1 if n < 100 else 2 if n < 250 else 3
    d = DENSITY[v][di]
    n_round = int(round(n, -1))
    parts = [OPEN[v].format(d=d, D=d.capitalize(), n=n_round)]
    if n >= 20:
        items = []
        for col, name in CELLS:
            lv = level(r[col] / max(n, 1))
            if lv is not None:
                items.append((lv, f"{LEVEL_WORDS[v][lv]} {name}"))
        items.sort(key=lambda x: -x[0])
        items = [s for _, s in items[:6]]  # most prominent 6 keep text < 77 tokens
        if v % 2 == 1:
            rng.shuffle(items)
        parts.append(("Shows " + ", ".join(items) + ".") if items else "No marker-positive population stands out.")
    return " ".join(parts)


rng = np.random.default_rng(0)
texts = {int(r.tile_id): [describe(r, v, rng) for v in range(4)] for _, r in df.iterrows()}
json.dump(texts, open(os.path.join(CACHE, "text.json"), "w"), indent=0)
u = len(set(t[0] for t in texts.values()))
print(f"{len(texts)} tiles, {u} unique variant-0 texts")
for k in [0, 500, 1500]:
    print(texts[k])