"""Build a MedCLIPSeg dataset for the neurofilament task (no manual masks needed).

For every CZI with an IMARIS _Statistics folder, writes:
  data/Neurofilament/<split>_Folder/img/<name>.png     (neurofilament MIP, RGB)
  data/Neurofilament/<split>_Folder/label/<name>.png   (IMARIS-reconstruction mask)
  data/Neurofilament/Prompts_Folder/<split>_text.xlsx  (Image, Ground Truth, Description)

The mask is the k-NN reconstruction of IMARIS's traced skeleton from its segment
coordinates (dilated) — the same approximate label used elsewhere. Splits are by
ANIMAL so validation/test are unseen animals. Then fine-tune the bundled model:

  cd reference_medclipseg
  python train.py --config-file configs/Neurofilament.yaml   # needs a CUDA GPU

Honest note: this label is approximate (reconstructed, not IMARIS's exact voxel
mask), and the target is IMARIS's over-traced tangle — see README.
"""
import argparse, glob, os, sys
import numpy as np
import pandas as pd
from scipy import ndimage as ndi
from scipy.spatial import cKDTree
from skimage.transform import resize
from skimage.io import imsave

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import processing as P

ROOT = "/Users/slf20757/Desktop/Dr. Fernandes"
SIZE = 224
PROMPT = "a fluorescence microscopy image of neurofilament nerve fibers"
# animal -> split
VAL = {"C1"}
TEST = {"C2b"}


def seg_xyz(sd, stem):
    def cols(name, nc):
        r = []
        for line in open(os.path.join(sd, f"{stem}_{name}.csv"), errors="ignore"):
            p = line.split(",")
            try:
                r.append([float(p[i]) for i in range(nc)])
            except (ValueError, IndexError):
                pass
        return np.array(r)
    return cols("Segment_Position", 3)


def skel_mask(pos, voxel, shape_yx, k=3):
    dz, dy, dx = voxel; ny, nx = shape_yx
    col = pos[:, 0] / dx * (SIZE / nx); row = pos[:, 1] / dy * (SIZE / ny)
    pts = np.stack([row, col], 1)
    tree = cKDTree(pts); dist, idx = tree.query(pts, k=min(k + 1, len(pts)))
    radius = 2.0 / dx * (SIZE / nx)
    m = np.zeros((SIZE, SIZE), bool)
    for i in range(len(pts)):
        for j, d in zip(idx[i, 1:], dist[i, 1:]):
            if d <= radius:
                p0, p1 = pts[i], pts[j]
                n = max(int(np.abs(p1 - p0).max()) + 1, 2)
                t = np.linspace(0, 1, n)[:, None]; line = p0 + t * (p1 - p0)
                m[np.clip(line[:, 0].astype(int), 0, SIZE - 1),
                  np.clip(line[:, 1].astype(int), 0, SIZE - 1)] = True
    return ndi.binary_dilation(m, iterations=1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/Neurofilament")
    ap.add_argument("--limit", type=int, default=0, help="cap images (for testing)")
    args = ap.parse_args()
    rows = {"Train": [], "Val": [], "Test": []}
    n = 0
    for csv in sorted(glob.glob(os.path.join(ROOT, "*", "*_Statistics",
                                             "*_Filament_Length_(sum).csv"))):
        sd = os.path.dirname(csv)
        stem = os.path.basename(csv).replace("_Filament_Length_(sum).csv", "")
        animal = os.path.relpath(csv, ROOT).split(os.sep)[0]
        czi = os.path.join(ROOT, animal, stem.replace(" filament", "") + ".czi")
        if not (os.path.exists(czi) and
                os.path.exists(os.path.join(sd, f"{stem}_Segment_Position_X.csv"))):
            continue
        split = "Val" if animal in VAL else "Test" if animal in TEST else "Train"
        img = P.load_image(czi); nf, _ = P.guess_channels(img)
        vol = img.data[nf]; ny, nx = vol.shape[1], vol.shape[2]
        mip = resize(P.channel_preview(vol), (SIZE, SIZE), preserve_range=True).astype(np.uint8)
        mask = (skel_mask(seg_xyz(sd, stem), img.voxel, (ny, nx)) * 255).astype(np.uint8)
        name = stem.replace(" ", "_") + ".png"
        for sub in ("img", "label"):
            os.makedirs(os.path.join(args.out, f"{split}_Folder", sub), exist_ok=True)
        imsave(os.path.join(args.out, f"{split}_Folder", "img", name),
               np.stack([mip] * 3, -1))
        imsave(os.path.join(args.out, f"{split}_Folder", "label", name), mask)
        rows[split].append({"Image": name, "Ground Truth": name, "Description": PROMPT})
        n += 1
        print(f"[{split}] {stem}: mask px={int((mask>0).sum())}", flush=True)
        del img, vol
        if args.limit and n >= args.limit:
            break

    os.makedirs(os.path.join(args.out, "Prompts_Folder"), exist_ok=True)
    for split, r in rows.items():
        if r:
            pd.DataFrame(r).to_excel(
                os.path.join(args.out, "Prompts_Folder", f"{split}_text.xlsx"),
                index=False)
    print(f"\nDone. Train={len(rows['Train'])} Val={len(rows['Val'])} "
          f"Test={len(rows['Test'])} -> {args.out}")
    print("Now fine-tune (GPU): cd reference_medclipseg && "
          "python train.py --config-file configs/Neurofilament.yaml")


if __name__ == "__main__":
    main()