"""Few-shot MedCLIPSeg-style training: frozen CLIP + text-conditioned decoder.

Train on 8 animals (1A,1B,2A,3A,3B,4Aa,4Bb,5A); test on C1 and C2b (unseen).
Faithful to the MedCLIPSeg approach that this machine can run: FROZEN CLIP
encoders provide image patch features + a text-prompt embedding; a lightweight
decoder (FiLM text conditioning) is trained data-efficiently with Dice+BCE to
segment the fibers. Reports Dice (DSC) on the held-out animals and saves
predictions. Labels = IMARIS reconstruction masks (approximate).
"""
import os, sys, glob, time
import numpy as np
import torch
import torch.nn as nn
from skimage.transform import resize
from skimage.io import imsave

MED = "/Users/slf20757/Desktop/Dr. Fernandes/.claude/worktrees/intelligent-kalam-5d8731/Desktop/Dr. Fernandes/hf_neuron_tracer_medclip"
sys.path.insert(0, MED)
import processing as P
import medclipseg as MC
from prepare_medclip_data import seg_xyz, skel_mask   # reconstruction label

ROOT = "/Users/slf20757/Desktop/Dr. Fernandes"
TRAIN = {"1A", "1B", "2A", "3A", "3B", "4Aa", "4Bb", "5A"}
TEST = {"C1", "C2b"}
SIZE = 224
dev = ("mps" if torch.backends.mps.is_available()
       else "cuda" if torch.cuda.is_available() else "cpu")


def clip_features(mip_u8):
    """Frozen CLIP vision patch grid (768,14,14) for a 224 grayscale MIP."""
    clip, proc = MC._load()
    rgb = np.stack([resize(mip_u8, (SIZE, SIZE), preserve_range=True)] * 3, -1).astype(np.uint8)
    px = proc(images=rgb, return_tensors="pt")["pixel_values"]
    with torch.no_grad():
        toks = clip.vision_model(pixel_values=px).last_hidden_state[0, 1:, :]  # (196,768)
    g = int(round(toks.shape[0] ** 0.5))
    return toks.reshape(g, g, -1).permute(2, 0, 1).contiguous()   # (768,g,g)


class Decoder(nn.Module):
    def __init__(self, cin=768, tdim=512, C=128):
        super().__init__()
        self.reduce = nn.Conv2d(cin, C, 1)
        self.film = nn.Linear(tdim, 2 * C)

        def up(ci, co):
            return nn.Sequential(nn.ConvTranspose2d(ci, co, 2, 2),
                                 nn.Conv2d(co, co, 3, padding=1),
                                 nn.BatchNorm2d(co), nn.ReLU(True))
        self.up = nn.Sequential(up(C, C), up(C, C), up(C, C // 2), up(C // 2, C // 4))
        self.out = nn.Conv2d(C // 4, 1, 1)

    def forward(self, feat, tvec):
        x = self.reduce(feat)
        g, b = self.film(tvec).chunk(2, -1)
        x = x * (1 + g[:, :, None, None]) + b[:, :, None, None]
        x = self.up(x)                                # 14 -> 224
        return torch.sigmoid(self.out(x))


def dice_metric(p, t, thr=0.5):
    p = (p > thr).float()
    return float((2 * (p * t).sum() + 1) / (p.sum() + t.sum() + 1))


# ---- build cached features + masks ----
print(f"device={dev}; extracting frozen CLIP features ...", flush=True)
tvec = MC._text_embeds(MC.FG_PROMPTS).mean(0, keepdim=True).to(dev)   # (1,512) prompt
feats, masks, animals, stems = [], [], [], []
t0 = time.time()
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
    if animal not in TRAIN and animal not in TEST:
        continue
    img = P.load_image(czi); nf, _ = P.guess_channels(img); vol = img.data[nf]
    mip = P.channel_preview(vol)
    feats.append(clip_features(mip))
    m = skel_mask(seg_xyz(sd, stem), img.voxel, (vol.shape[1], vol.shape[2]))
    masks.append(torch.tensor(m.astype(np.float32)))
    animals.append(animal); stems.append(stem)
    del img, vol
F = torch.stack(feats).to(dev)                      # (N,768,14,14)
M = torch.stack(masks)[:, None].to(dev)             # (N,1,224,224)
animals = np.array(animals)
tr = np.isin(animals, list(TRAIN)); te = np.isin(animals, list(TEST))
print(f"extracted {len(F)} imgs ({tr.sum()} train / {te.sum()} test) in {time.time()-t0:.0f}s", flush=True)

# ---- train decoder ----
net = Decoder().to(dev)
opt = torch.optim.Adam(net.parameters(), 1e-3, weight_decay=1e-5)
idx = np.where(tr)[0]
t1 = time.time(); net.train()
for ep in range(300):
    perm = np.random.permutation(idx)
    for k in range(0, len(perm), 8):
        b = perm[k:k + 8]
        opt.zero_grad()
        p = net(F[b], tvec.expand(len(b), -1))
        bce = nn.functional.binary_cross_entropy(p, M[b])
        dsc = 1 - (2 * (p * M[b]).sum() + 1) / (p.sum() + M[b].sum() + 1)
        (bce + dsc).backward(); opt.step()
    if (ep + 1) % 100 == 0:
        print(f"  epoch {ep+1}/300 ({time.time()-t1:.0f}s)", flush=True)

# ---- evaluate on held-out C1 + C2b ----
net.eval()
O = os.path.dirname(__file__)
print("\n=== held-out test (C1, C2b — never trained on) ===", flush=True)
train_dsc = []
with torch.no_grad():
    for i in np.where(tr)[0]:
        train_dsc.append(dice_metric(net(F[i:i+1], tvec), M[i:i+1]))
    dscs = []
    for i in np.where(te)[0]:
        p = net(F[i:i+1], tvec)
        d = dice_metric(p, M[i:i+1]); dscs.append(d)
        pred = (p[0, 0].cpu().numpy() > 0.5).astype(np.uint8) * 255
        gt = (M[i, 0].cpu().numpy() * 255).astype(np.uint8)
        imsave(os.path.join(O, f"pred_{stems[i].replace(' ','_')}.png"),
               np.concatenate([gt, np.full((224, 6), 128, np.uint8), pred], 1))
        print(f"  {stems[i]:16s} Dice={d:.3f}", flush=True)
print(f"\nTRAIN mean Dice={np.mean(train_dsc):.3f}  |  "
      f"HELD-OUT (C1+C2b) mean Dice={np.mean(dscs):.3f}", flush=True)
torch.save({"model": net.state_dict()}, os.path.join(MED, "weights_medclip_decoder.pt"))
print("saved decoder weights + prediction images (pred_*.png: GT | prediction)", flush=True)