"""CT support: Hounsfield-unit window presets and tissue segmentation.

On a synthetic head CT, shows the clinical window/level presets (brain / blood /
bone) and HU-threshold tissue masks (skull, brain, CSF/ventricle, acute clot).

RESEARCH SOFTWARE — NOT FOR CLINICAL USE.
"""
import os, sys
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, HERE)
from sulcora.volume import Volume
from sulcora import ct as C
from sulcora import hematoma as H


def main():
    docs = os.path.join(HERE, "docs", "images"); os.makedirs(docs, exist_ok=True)
    shape = (120, 120, 120); ni, nj, nk = shape
    ii, jj, kk = np.meshgrid(np.arange(ni), np.arange(nj), np.arange(nk), indexing="ij")
    c = np.array(shape) / 2
    r = np.sqrt((ii - c[0]) ** 2 + (jj - c[1]) ** 2 + (kk - c[2]) ** 2)
    d = np.full(shape, -1000.0, np.float32)                 # air
    d[r <= 52] = 35.0                                       # brain
    d[(r > 48) & (r <= 52)] = 1100.0                        # skull
    d[r <= 10] = 6.0                                        # ventricles (CSF)
    d[((ii - 70) ** 2 / 14 ** 2 + (jj - 60) ** 2 / 8 ** 2 + (kk - 60) ** 2 / 8 ** 2) <= 1] = 70.0  # clot
    d += np.random.default_rng(0).normal(0, 3, shape).astype(np.float32)
    vol = Volume(d, np.eye(4), meta={"modality": "CT"})
    print("is_ct:", C.is_ct(vol))

    z = nk // 2
    skull = C.skull_mask(vol).data[:, :, z]
    brain = C.brain_mask(vol).data[:, :, z]
    csf = C.csf_mask(vol).data[:, :, z]
    clot = H.segment_hematoma(vol, 45, 100).data[:, :, z]

    fig, ax = plt.subplots(1, 4, figsize=(15, 4.2), facecolor="#0b0f18")
    for a, name in zip(ax[:3], ("brain", "blood", "bone")):
        a.imshow(C.window_preset(vol, name).data[:, :, z].T, cmap="gray", origin="lower")
        a.set_title(f"{name.capitalize()} window", color="#9fd0ff", fontsize=11); a.axis("off")
    # segmentation overlay on the brain window
    ax[3].imshow(C.window_preset(vol, "brain").data[:, :, z].T, cmap="gray", origin="lower")
    for m, cmap in ((skull, "bone_r"), (brain, "winter"), (csf, "cool"), (clot, "spring")):
        ax[3].imshow(np.ma.masked_where(m.T == 0, m.T), cmap=cmap, alpha=0.45, origin="lower")
    ax[3].set_title("HU segmentation\n(skull·brain·CSF·clot)", color="#9fd0ff", fontsize=11); ax[3].axis("off")
    fig.suptitle("CT support — Hounsfield-unit windowing and tissue segmentation", color="#dce6f5", fontsize=13)
    fig.text(0.5, 0.02, "Same CT, three clinical windows; right: HU-threshold masks for skull, brain, CSF, and acute clot. "
             "RESEARCH SOFTWARE — NOT FOR CLINICAL USE.", color="#c0392b", ha="center", fontsize=9)
    fig.tight_layout(rect=[0, 0.05, 1, 0.93])
    out = os.path.join(docs, "ct_support.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("wrote", out, "\nCT PROOF OK")


if __name__ == "__main__":
    main()
