AgentSSL: Can MLE Agents Leverage Unlabeled Data?

from __future__ import annotations
import json
import os
from typing import List, Tuple

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T
from PIL import Image


def _read_json(path: str):
    with open(path, "r") as f:
        return json.load(f)


def _coco_label_mapping(train_ann_file: str):
    coco = _read_json(train_ann_file)
    cat_ids = sorted(c["id"] for c in coco["categories"])
    cat_id_to_idx = {cid: i for i, cid in enumerate(cat_ids)}
    idx_to_cat_id = {i: cid for cid, i in cat_id_to_idx.items()}
    return cat_id_to_idx, idx_to_cat_id, len(cat_ids)


def _labeled_records(img_dir: str, train_ann_file: str, cat_id_to_idx: dict):
    coco = _read_json(train_ann_file)
    id_to_name = {img["id"]: img["file_name"] for img in coco["images"]}
    out = []
    for ann in coco["annotations"]:
        fn = id_to_name.get(ann["image_id"])
        if fn is None:
            continue
        path = os.path.join(img_dir, fn)
        if not os.path.isfile(path):
            continue
        cid = ann["category_id"]
        if cid not in cat_id_to_idx:
            continue
        out.append((path, cat_id_to_idx[cid]))
    return out


def _unlabeled_records(img_dir: str, unlabel_ann_file: str):
    coco = _read_json(unlabel_ann_file)
    out = []
    for img in coco["images"]:
        fn = img["file_name"]
        path = os.path.join(img_dir, fn)
        if os.path.isfile(path):
            out.append(path)
    return out


class _ImageOnlyDataset(Dataset):
    def __init__(self, paths, transform):
        self.paths = paths
        self.transform = transform

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert("RGB")
        return self.transform(img)


class _GraphConfig:
    def __init__(
        self,
        k_ul: int = 24,
        k_uu: int = 24,
        temperature: float = 0.07,
        alpha: float = 0.85,
        iters: int = 25,
        chunk_q: int = 1024,
    ):
        self.k_ul = k_ul
        self.k_uu = k_uu
        self.temperature = temperature
        self.alpha = alpha
        self.iters = iters
        self.chunk_q = chunk_q


class ClassificationAgent:
    NET_NAME = "timm/vit_base_patch14_reg4_dinov2.lvd142m"

    def __init__(self, net_builder_fn, get_peft_config_fn, num_classes):
        self.net_builder_fn = net_builder_fn
        self.get_peft_config_fn = get_peft_config_fn
        self.num_classes = num_classes
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.img_size = 224
        self.transform = self._eval_transform(self.img_size)
        self.idx_to_cat_id = {i: i for i in range(num_classes)}

        peft_config = self.get_peft_config_fn({"freeze_backbone": True})
        vit_config = {"drop_path_rate": 0.0}
        net_builder = self.net_builder_fn(self.NET_NAME, peft_config, vit_config)
        self.model = net_builder(
            num_classes=num_classes, pretrained=True, pretrained_path=""
        ).to(self.device)
        self.model.eval()

        self.graph_cfg = _GraphConfig()

        self.train_feat = None
        self.train_labels = None
        self.train_prop = None
        self.class_prototypes = None
        self.residual_quantiles = None

        self.k_query = 32
        self.k_proto = 5
        self.tta_shift = 8

        self.fuse_w_knn_hard = 1.15
        self.fuse_w_knn_prop = 0.85
        self.fuse_w_proto = 0.55
        self.fuse_w_ord = 0.30
        self.score_temp = 0.07
        self.proto_temp = 0.09
        self.conformal_q = 0.80

    def _eval_transform(self, img_size=224):
        resize = int(img_size * 256 / 224)
        return T.Compose(
            [
                T.Resize(resize, interpolation=T.InterpolationMode.BICUBIC),
                T.CenterCrop(img_size),
                T.ToTensor(),
                T.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                ),
            ]
        )

    @torch.no_grad()
    def _extract_features(self, paths, batch_size=64):
        if len(paths) == 0:
            return torch.empty((0, 768), dtype=torch.float32)

        ds = _ImageOnlyDataset(paths, self.transform)
        loader = DataLoader(
            ds,
            batch_size=min(batch_size, max(1, len(ds))),
            shuffle=False,
            num_workers=min(8, os.cpu_count() or 0),
            pin_memory=self.device.type == "cuda",
            drop_last=False,
        )

        feats = []
        for x in loader:
            x = x.to(self.device, non_blocking=True)
            out = self.model(x)
            feats.append(out["feat"].detach().cpu())
        if len(feats) == 0:
            return torch.empty((0, 768), dtype=torch.float32)
        return torch.cat(feats, dim=0)

    def _build_sparse_neighbors(self, xu: torch.Tensor, xl: torch.Tensor):
        cfg = self.graph_cfg
        device = self.device
        xu = xu.to(device)
        xl = xl.to(device)

        Nu = xu.shape[0]
        Nl = xl.shape[0]
        if Nu == 0 or Nl == 0:
            raise RuntimeError(
                "Neighbor graph requires both labeled and unlabeled features."
            )

        k_ul = min(cfg.k_ul, Nl)
        k_uu = min(cfg.k_uu + 1, Nu)

        ul_idx_parts, ul_w_parts = [], []
        uu_idx_parts, uu_w_parts = [], []

        for start in range(0, Nu, cfg.chunk_q):
            end = min(start + cfg.chunk_q, Nu)
            q = xu[start:end]

            sim_ul = q @ xl.T
            vals_ul, idx_ul = torch.topk(sim_ul, k=max(1, k_ul), dim=1)
            w_ul = torch.softmax(vals_ul / cfg.temperature, dim=1)
            ul_idx_parts.append(idx_ul.cpu())
            ul_w_parts.append(w_ul.cpu())

            if Nu == 1:
                idx_uu = torch.zeros((end - start, 1), dtype=torch.long, device=device)
                vals_uu = torch.zeros(
                    (end - start, 1), dtype=torch.float32, device=device
                )
            else:
                sim_uu = q @ xu.T
                vals_uu, idx_uu = torch.topk(sim_uu, k=max(2, k_uu), dim=1)
                rows = torch.arange(start, end, device=device).unsqueeze(1)
                keep = idx_uu != rows
                idx_uu = idx_uu[keep].view(end - start, -1)
                vals_uu = vals_uu[keep].view(end - start, -1)
                if idx_uu.shape[1] == 0:
                    idx_uu = torch.zeros(
                        (end - start, 1), dtype=torch.long, device=device
                    )
                    vals_uu = torch.zeros(
                        (end - start, 1), dtype=torch.float32, device=device
                    )
                else:
                    idx_uu = idx_uu[:, : min(cfg.k_uu, idx_uu.shape[1])]
                    vals_uu = vals_uu[:, : min(cfg.k_uu, vals_uu.shape[1])]

            w_uu = torch.softmax(vals_uu / cfg.temperature, dim=1)
            uu_idx_parts.append(idx_uu.cpu())
            uu_w_parts.append(w_uu.cpu())

        ul_idx = torch.cat(ul_idx_parts, dim=0)
        ul_w = torch.cat(ul_w_parts, dim=0)
        uu_idx = torch.cat(uu_idx_parts, dim=0)
        uu_w = torch.cat(uu_w_parts, dim=0)
        return ul_idx, ul_w, uu_idx, uu_w

    def _label_propagation(self, xl: torch.Tensor, yl: torch.Tensor, xu: torch.Tensor):
        num_classes = self.num_classes
        cfg = self.graph_cfg

        xl = F.normalize(xl.float(), dim=1)
        xu = F.normalize(xu.float(), dim=1)

        y_onehot = F.one_hot(yl, num_classes=num_classes).float()
        ul_idx, ul_w, uu_idx, uu_w = self._build_sparse_neighbors(xu, xl)

        yu = torch.full(
            (xu.shape[0], num_classes), 1.0 / num_classes, dtype=torch.float32
        )

        for _ in range(cfg.iters):
            msg_l = torch.sum(ul_w.unsqueeze(-1) * y_onehot[ul_idx], dim=1)
            msg_u = torch.sum(uu_w.unsqueeze(-1) * yu[uu_idx], dim=1)
            yu = (1.0 - cfg.alpha) * msg_l + cfg.alpha * msg_u
            yu = yu / yu.sum(dim=1, keepdim=True).clamp_min(1e-8)

        return y_onehot, yu

    def _compute_train_propagated_probs(
        self, xl: torch.Tensor, yl: torch.Tensor, xu: torch.Tensor, yu: torch.Tensor
    ):
        xl_n = F.normalize(xl.float(), dim=1)
        xu_n = F.normalize(xu.float(), dim=1)
        y_onehot = F.one_hot(yl, num_classes=self.num_classes).float()

        sims_ll = xl_n @ xl_n.T
        sims_lu = xl_n @ xu_n.T

        if xl.shape[0] <= 1:
            msg_ll = y_onehot.clone()
        else:
            k_ll = min(16, xl.shape[0] - 1)
            vals_ll, idx_ll = torch.topk(
                sims_ll, k=min(k_ll + 1, sims_ll.shape[1]), dim=1
            )
            rows = torch.arange(xl.shape[0]).unsqueeze(1)
            keep = idx_ll != rows
            idx_ll = idx_ll[keep].view(xl.shape[0], -1)[:, :k_ll]
            vals_ll = vals_ll[keep].view(xl.shape[0], -1)[:, :k_ll]
            w_ll = torch.softmax(vals_ll / 0.07, dim=1)
            msg_ll = torch.sum(w_ll.unsqueeze(-1) * y_onehot[idx_ll], dim=1)

        if xu.shape[0] == 0:
            msg_lu = torch.full_like(msg_ll, 1.0 / self.num_classes)
        else:
            k_lu = min(32, xu.shape[0])
            vals_lu, idx_lu = torch.topk(sims_lu, k=k_lu, dim=1)
            w_lu = torch.softmax(vals_lu / 0.07, dim=1)
            msg_lu = torch.sum(w_lu.unsqueeze(-1) * yu[idx_lu], dim=1)

        prop = 0.6 * msg_ll + 0.4 * msg_lu
        prop = prop / prop.sum(dim=1, keepdim=True).clamp_min(1e-8)
        return prop

    def _build_prototypes(self, xl: torch.Tensor, yl: torch.Tensor):
        protos = []
        for c in range(self.num_classes):
            mask = yl == c
            if mask.any():
                proto = xl[mask].mean(dim=0)
            else:
                proto = xl.mean(dim=0)
            protos.append(F.normalize(proto, dim=0))
        return torch.stack(protos, dim=0)

    def _fused_local_logits(
        self,
        query_feat: torch.Tensor,
        leave_one_out_index: int | None = None,
    ):
        feat = F.normalize(query_feat.float(), dim=1).cpu()
        train_feat = self.train_feat
        train_labels = self.train_labels
        train_prop = self.train_prop
        prototypes = self.class_prototypes
        num_classes = self.num_classes

        sim = feat @ train_feat.T

        if leave_one_out_index is not None and feat.shape[0] == 1:
            sim[:, leave_one_out_index] = -1e4

        k = min(
            self.k_query,
            train_feat.shape[0] - (1 if leave_one_out_index is not None else 0),
        )
        k = max(1, k)

        vals, idx = torch.topk(sim, k=k, dim=1)
        w = torch.softmax(vals / self.score_temp, dim=1)

        neigh_labels = train_labels[idx]
        hard = F.one_hot(neigh_labels, num_classes=num_classes).float()
        hard_prob = torch.sum(w.unsqueeze(-1) * hard, dim=1)

        prop_prob = torch.sum(w.unsqueeze(-1) * train_prop[idx], dim=1)
        prop_prob = prop_prob / prop_prob.sum(dim=1, keepdim=True).clamp_min(1e-8)

        sim_proto = feat @ prototypes.T
        proto_prob = torch.softmax(sim_proto / self.proto_temp, dim=1)

        class_vals = torch.arange(num_classes, dtype=feat.dtype).unsqueeze(0)
        exp_severity = torch.sum(hard_prob * class_vals, dim=1, keepdim=True)
        ord_bias = -torch.abs(class_vals - exp_severity)

        fused = (
            self.fuse_w_knn_hard * torch.log(hard_prob.clamp_min(1e-8))
            + self.fuse_w_knn_prop * torch.log(prop_prob.clamp_min(1e-8))
            + self.fuse_w_proto * torch.log(proto_prob.clamp_min(1e-8))
            + self.fuse_w_ord * ord_bias
        )
        return fused

    def _build_conformal_residuals(self):
        n = self.train_feat.shape[0]
        residuals = []
        correct = 0
        for i in range(n):
            logits_i = self._fused_local_logits(
                self.train_feat[i : i + 1], leave_one_out_index=i
            )
            probs_i = torch.softmax(logits_i, dim=-1).squeeze(0)
            true_y = int(self.train_labels[i].item())
            true_prob = probs_i[true_y].item()
            residuals.append(1.0 - true_prob)
            if int(torch.argmax(probs_i).item()) == true_y:
                correct += 1

        residuals = torch.tensor(residuals, dtype=torch.float32)
        q = torch.quantile(residuals, self.conformal_q).item()
        self.residual_quantiles = torch.full(
            (self.num_classes,), float(q), dtype=torch.float32
        )
        print(
            f"local_loo_acc={correct / max(n, 1):.4f} "
            f"conformal_q={self.conformal_q:.2f} "
            f"resid_quantile={q:.4f}",
            flush=True,
        )

    def fit(self, img_dir, train_ann_file, unlabel_ann_file):
        cat_id_to_idx, idx_to_cat_id, num_classes = _coco_label_mapping(train_ann_file)
        self.idx_to_cat_id = idx_to_cat_id
        self.num_classes = num_classes

        labeled = _labeled_records(img_dir, train_ann_file, cat_id_to_idx)
        unlabeled = _unlabeled_records(img_dir, unlabel_ann_file)

        if len(labeled) == 0:
            raise RuntimeError("No labeled samples found.")
        if len(unlabeled) == 0:
            raise RuntimeError("No unlabeled samples found.")

        labeled_paths = [p for p, _ in labeled]
        yl = torch.tensor([y for _, y in labeled], dtype=torch.long)

        print(
            f"extract_features labeled={len(labeled_paths)} unlabeled={len(unlabeled)}",
            flush=True,
        )
        xl = self._extract_features(labeled_paths, batch_size=64)
        xu = self._extract_features(unlabeled, batch_size=64)

        xl = F.normalize(xl.float(), dim=1)
        xu = F.normalize(xu.float(), dim=1)

        print("running_label_propagation", flush=True)
        _, yu = self._label_propagation(xl, yl, xu)
        prop_l = self._compute_train_propagated_probs(xl, yl, xu, yu)

        self.train_feat = xl.cpu()
        self.train_labels = yl.cpu()
        self.train_prop = prop_l.cpu()
        self.class_prototypes = self._build_prototypes(
            self.train_feat, self.train_labels
        ).cpu()

        print("building_local_conformal_calibration", flush=True)
        self._build_conformal_residuals()

        with torch.no_grad():
            train_logits = self._fused_local_logits(self.train_feat)
            train_probs = torch.softmax(train_logits, dim=-1)
            train_pred = train_probs.argmax(dim=-1)
            acc = (train_pred == self.train_labels).float().mean().item()
            nll = F.nll_loss(
                torch.log(train_probs.clamp_min(1e-8)), self.train_labels
            ).item()
            print(f"fit_train_acc={acc:.4f} fit_train_nll={nll:.4f}", flush=True)

        self.model.eval()
        print("fit_done", flush=True)

    def _tta_views(self, x: torch.Tensor):
        views = [x]
        views.append(torch.flip(x, dims=[3]))
        s = min(self.tta_shift, x.shape[-1] // 16)
        if s > 0:
            views.append(torch.roll(x, shifts=-s, dims=3))
            views.append(torch.roll(x, shifts=s, dims=3))
        return views

    @torch.no_grad()
    def _predict_from_feat(self, feat: torch.Tensor):
        raw_logits = self._fused_local_logits(feat).to(self.device)
        probs = torch.softmax(raw_logits, dim=-1)

        penalty = self.residual_quantiles.to(self.device).unsqueeze(0)
        calibrated_probs = (probs - penalty).clamp_min(1e-8)
        calibrated_probs = calibrated_probs / calibrated_probs.sum(
            dim=1, keepdim=True
        ).clamp_min(1e-8)
        calibrated_logits = torch.log(calibrated_probs.clamp_min(1e-8))
        return calibrated_logits

    @torch.no_grad()
    def predict(self, image_batch):
        self.model.eval()

        if (
            self.train_feat is None
            or self.train_prop is None
            or self.residual_quantiles is None
        ):
            raise RuntimeError("Call fit() before predict().")

        if isinstance(image_batch, torch.Tensor):
            x = image_batch
            if x.ndim == 3:
                x = x.unsqueeze(0)
            x = x.float().to(self.device, non_blocking=True)
        elif isinstance(image_batch, list):
            x = torch.stack(
                [self.transform(im.convert("RGB")) for im in image_batch], dim=0
            ).to(self.device, non_blocking=True)
        else:
            raise TypeError("image_batch must be a Tensor or list of PIL images")

        views = self._tta_views(x)
        logits_sum = None
        feat_main = None

        for i, xv in enumerate(views):
            out = self.model(xv)
            feat = F.normalize(out["feat"].float(), dim=1)
            if i == 0:
                feat_main = feat
            logits_v = self._predict_from_feat(feat)
            logits_sum = logits_v if logits_sum is None else (logits_sum + logits_v)

        logits = logits_sum / float(len(views))
        pred_idx = logits.argmax(dim=-1).detach().cpu().tolist()
        predictions = [{"category_id": int(self.idx_to_cat_id[i])} for i in pred_idx]
        return logits, feat_main, predictions