classifier.py

import torch
import torch.nn as nn
from pathlib import Path
import torch.nn.functional as F
import models
from model_utils import generate_rollout


def check_embeddings(embeddings):
    if embeddings.dim() == 1:
        return embeddings.unsqueeze(0).unsqueeze(0)  # [D] -> [1, 1, D]
    elif embeddings.dim() == 2:
        return embeddings.unsqueeze(0)  # [N, D] -> [1, N, D]
    elif embeddings.dim() != 3:
        raise ValueError(f"Unexpected embedding shape: {embeddings.shape}")
    return embeddings


class cAItomorph(nn.Module):
    def __init__(self, class_count, arch='ViT', embedding_dim=768, device='cpu'):
        super().__init__()
        if arch not in models.__dict__:
            raise ValueError(f"Unknown model architecture '{arch}'")
        self.model = models.__dict__[arch](input_dim=embedding_dim, num_classes=class_count)
        self.device = device

    def forward(self, embeddings, return_latent=False, register_hook=False):
        embeddings = check_embeddings(embeddings)
        latent, logits = self.model(embeddings, register_hook=register_hook)
        return (latent, logits) if return_latent else logits

    def load_checkpoint(self, checkpoint):
        """
        Load model weights from either:
        - a filepath (.pth or .pt)
        - a dictionary containing a 'weights' key
        - a state_dict directly
        """
        # --- Case 1: Path string or Path object ---
        if isinstance(checkpoint, (str, Path)):
            checkpoint = Path(checkpoint)
            if not checkpoint.exists():
                raise FileNotFoundError(f"Checkpoint not found: {checkpoint}")
            state_dict = torch.load(checkpoint, map_location='cpu')
            if "weights" in state_dict:
                state_dict = state_dict["weights"]

        # --- Case 2: Dictionary with 'weights' key ---
        elif isinstance(checkpoint, dict) and "weights" in checkpoint:
            state_dict = checkpoint["weights"]

        # --- Case 3: Direct state_dict ---
        elif isinstance(checkpoint, dict):
            state_dict = checkpoint

        else:
            raise ValueError("Invalid checkpoint type. Provide a path, state_dict, or dict with 'weights' key.")

        # --- Load and prepare model ---
        self.model.load_state_dict(state_dict, strict=True)
        self.model.to(self.device)
        self.eval()

        print("✅ Checkpoint successfully loaded.")

    def __repr__(self):
        return f"cAItomorph(model={self.model})"


class EnsembleModel(nn.Module):
    def __init__(self, class_count, arch='ViT', embedding_dim=768, weights_path=None,
                 return_attention=False, device="cpu"):
        super().__init__()
        self.device = device
        self.arch = arch
        self.class_count = class_count
        self.embedding_dim = embedding_dim
        self.return_attention = return_attention

        if weights_path is None:
            raise ValueError("`weights_path` must be provided.")
        self.weights_path = Path(weights_path)

        # Find all .pth files in directory
        self.weight_files = sorted(self.weights_path.glob("*.pth"))
        if len(self.weight_files) == 0:
            raise FileNotFoundError(f"No .pth files found in {self.weights_path}")
        self.num_models = len(self.weight_files)

        # Load all models
        self.models = nn.ModuleList()
        for wf in self.weight_files:
            model = cAItomorph(class_count, arch, embedding_dim, device=self.device)
            model.load_checkpoint(wf)
            self.models.append(model)

        print(f"✅ Loaded {self.num_models} models from {self.weights_path}")

    @torch.no_grad()
    def forward(self, embeddings):
        embeddings = check_embeddings(embeddings).to(self.device)
        b, n, d = embeddings.shape

        results = {}
        probs_list = []
        attention_list = []

        for i, model in enumerate(self.models):
            model_name = f"model_{i}"
            latent, logits = model(embeddings, return_latent=True)

            if self.return_attention:
                attention_raw = generate_rollout(model.model, embeddings, start_layer=0)
                attention_raw = attention_raw.squeeze(0)  # [N]
                attention_soft = F.softmax(attention_raw, dim=-1)  # [N]
            else:
                attention_raw = torch.zeros(n, device=self.device)
                attention_soft = torch.zeros(n, device=self.device)

            probs = F.softmax(logits, dim=-1)

            results[model_name] = {
                "latent": latent.detach().cpu(),
                "logits": logits.detach().cpu(),
                "attention_raw": attention_raw.detach().cpu(),
                "attention_soft": attention_soft.detach().cpu(),
            }
            probs_list.append(probs)
            attention_list.append(attention_soft)

        # === Ensemble averaging ===
        probs_stack = torch.stack(probs_list, dim=0)  # [num_models, batch, num_classes]
        ensemble_prob = probs_stack.mean(dim=0)       # [batch, num_classes]

        attention_stack = torch.stack(attention_list, dim=0)  # [num_models, N]
        attentions = attention_stack.mean(dim=0)              # [N]

        results["ensemble"] = {
            "probability": ensemble_prob.detach().cpu(),
            "attentions": attentions.detach().cpu(),
        }

        return results

    def __repr__(self):
        return f"EnsembleModel(num_models={self.num_models}, arch={self.arch})"