"""Descriptive analysis of the ORION tri-modal (+markers) data. Writes results/data_summary.json + figures."""
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import tifffile
from PIL import Image

from common import ORION, ROOT, load_all_tiles, marker_columns

out = os.path.join(ROOT, "results")
os.makedirs(out, exist_ok=True)
df = load_all_tiles()
prop, cnt = marker_columns(df)

summary = {
    "n_tiles": len(df),
    "n_slides": int(df.in_slide_name.nunique()),
    "n_patients": int(df.patient.nunique()),
    "official_split_tiles": df.orion_split.value_counts().to_dict(),
    "official_split_slides": df.groupby("orion_split").in_slide_name.nunique().to_dict(),
    "if_batches_tiles": df.if_batch.value_counts().to_dict(),
    "if_batches_patients": df.groupby("if_batch").patient.nunique().to_dict(),
    "tiles_per_slide": df.groupby("in_slide_name").size().describe().round(1).to_dict(),
}

# per-tile stats from images (all tiles; small set)
he_mean, if_ch_mean, n_nuc, nuc_cov = [], [], [], []
for _, r in df.iterrows():
    he = np.asarray(Image.open(os.path.join(ORION, r.image_path)).convert("RGB"))
    iff = tifffile.imread(os.path.join(ORION, r.target_path))
    nuc = tifffile.imread(os.path.join(ORION, r.nuclei_path))
    he_mean.append(he.mean())
    if_ch_mean.append(iff.reshape(-1, iff.shape[-1]).mean(0))
    n_nuc.append(len(np.unique(nuc)) - 1)
    nuc_cov.append((nuc > 0).mean())
df["he_mean"], df["n_nuclei"], df["nuc_cov"] = he_mean, n_nuc, nuc_cov
if_ch_mean = np.stack(if_ch_mean)
summary["shapes"] = {"he": list(he.shape), "if": list(iff.shape), "nuclei": list(nuc.shape),
                     "markers": len(prop) + len(cnt)}
summary["if_channel_names_from_prop_cols"] = [c[:-5] for c in prop]
summary["n_nuclei_per_tile"] = df.n_nuclei.describe().round(1).to_dict()
summary["frac_tiles_lt20_nuclei"] = float((df.n_nuclei < 20).mean())
summary["he_mean_vs_n_nuclei_spearman"] = float(df[["he_mean", "n_nuclei"]].corr("spearman").iloc[0, 1])
summary["marker_prop_mean"] = df[prop].mean().round(4).to_dict()
summary["marker_count_mean"] = df[cnt].mean().round(2).to_dict()
summary["if_channel_mean_intensity"] = dict(zip([c[:-5] for c in prop], if_ch_mean.mean(0).round(2).tolist()))
df[["tile_id", "in_slide_name", "patient", "if_batch", "orion_split", "n_nuclei", "nuc_cov", "he_mean"]].to_csv(
    os.path.join(out, "tile_stats.csv"), index=False)
json.dump(summary, open(os.path.join(out, "data_summary.json"), "w"), indent=1)
print(json.dumps(summary, indent=1))

# figure: example quadruplets
rng = np.random.default_rng(0)
ex = df[df.n_nuclei > 100].sample(4, random_state=0)
fig, ax = plt.subplots(4, 5, figsize=(15, 12))
for i, (_, r) in enumerate(ex.iterrows()):
    he = Image.open(os.path.join(ORION, r.image_path)).convert("RGB")
    iff = tifffile.imread(os.path.join(ORION, r.target_path))
    nuc = tifffile.imread(os.path.join(ORION, r.nuclei_path))
    ax[i, 0].imshow(he); ax[i, 0].set_title(f"H&E  {r.orion_slide_id}")
    # Hoechst / Pan-CK / CD45 as RGB-ish composite
    names = [c[:-5] for c in prop]
    comp = np.stack([iff[..., names.index(n)] for n in ["Pan-CK", "CD45", "Hoechst"]], -1)
    ax[i, 1].imshow(comp); ax[i, 1].set_title("IF: Pan-CK/CD45/Hoechst")
    ax[i, 2].imshow(iff[..., names.index("SMA")], cmap="magma"); ax[i, 2].set_title("IF: SMA")
    ax[i, 3].imshow(nuc > 0, cmap="gray"); ax[i, 3].set_title(f"nuclei ({r.n_nuclei})")
    ax[i, 4].barh(range(len(cnt)), np.log1p(r[cnt].values.astype(float)))
    ax[i, 4].set_yticks(range(len(cnt))); ax[i, 4].set_yticklabels([c[:-6] for c in cnt], fontsize=6)
    ax[i, 4].set_title("markers log1p(count)")
    for a in ax[i, :4]:
        a.axis("off")
plt.tight_layout(); plt.savefig(os.path.join(out, "example_quadruplets.png"), dpi=80)

fig, ax = plt.subplots(1, 3, figsize=(15, 4))
ax[0].hist(df.n_nuclei, bins=50); ax[0].set_title("nuclei per tile"); ax[0].set_yscale("log")
df.groupby("if_batch").size().plot.bar(ax=ax[1], title="tiles per IF batch")
ax[2].imshow(np.corrcoef(np.concatenate([df[prop].values, np.log1p(df[cnt].values)], 1).T), cmap="RdBu_r", vmin=-1, vmax=1)
ax[2].set_title("marker vector correlation (17 prop + 16 log count)")
plt.tight_layout(); plt.savefig(os.path.join(out, "data_overview.png"), dpi=80)