model.py

"""Model definitions for AS-TIME.

The proposed model is ``ASTIME`` (time-conditioned multimodal fusion:
MIL -> delta_t-conditioned gated fusion -> delta_t-guided MoE). The
remaining classes are the Table 1 comparison baselines. Use the
``ASTimeModel`` factory at the bottom of this file to instantiate a model
from a set of modality / fusion flags.
"""

from typing import Dict, Any, List, Optional, Tuple
import math

import torch
from torch import nn
from torch.nn import functional as F

from time_embedding import TimePosConfig, PosEmbedding, build_time_features


def _to_device_if_needed(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
    if tensor.device == device:
        return tensor
    return tensor.to(device, non_blocking=True)


# ===========================================================================
# Baselines (Table 1 comparison methods)
# ===========================================================================
class TransformerAggregator(nn.Module):
    """
    Aggregator with attention export.
    - returns CLS embedding
    - optionally returns attention from CLS to input tokens (avg over heads)
    """
    def __init__(self, dim: int, n_heads: int = 4, dropout: float = 0.3):
        super().__init__()
        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
        self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=n_heads,
                                          dropout=dropout, batch_first=True)
        self.ln1 = nn.LayerNorm(dim)
        self.ff = nn.Sequential(
            nn.Linear(dim, 4 * dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(4 * dim, dim),
        )
        self.ln2 = nn.LayerNorm(dim)

    def forward(self, x: torch.Tensor, return_attn: bool = False):
        """
        x: [N, D] or [B, N, D]
        return_attn:
          - if True: returns (cls_emb, attn_cls) where attn_cls is [B, N]
                    (attention weights from CLS to each input token, averaged over heads)
        """
        squeeze_back = False
        if x.ndim == 2:
            x = x.unsqueeze(0)  # [1, N, D]
            squeeze_back = True

        B, N, D = x.shape
        cls = self.cls_token.expand(B, 1, D)   # [B,1,D]
        x = torch.cat([cls, x], dim=1)         # [B,1+N,D]

        # Self-attn
        attn_out, attn_w = self.attn(x, x, x, need_weights=True, average_attn_weights=False)
        # attn_w: [B, num_heads, (1+N), (1+N)]

        x = self.ln1(x + attn_out)
        x = self.ln2(x + self.ff(x))

        cls_out = x[:, 0, :]  # [B,D]

        if squeeze_back:
            cls_out = cls_out.squeeze(0)  # [D]

        if not return_attn:
            return cls_out

        # CLS attending to input tokens only (exclude CLS itself)
        # attn_w[:, :, 0, 1:] -> [B, heads, N]
        attn_cls = attn_w[:, :, 0, 1:].mean(dim=1)  # [B, N]

        if squeeze_back:
            attn_cls = attn_cls.squeeze(0)  # [N]

        return cls_out, attn_cls

class ASVideoOnlyModel(nn.Module):
    """
    Video-only AS-time model.

    - Input: per-video embeddings (each video = 1 token)
    - No reports, no time
    - Pipeline: video_proj -> TransformerAggregator -> classifier

    Forward is kept compatible:
      forward(video_0, view_0, report0_cls, delta_t) -> logits [B, num_classes]
    Only video_0 is used.
    """

    def __init__(
        self,
        video_dim: int = 512,
        d_model: int = 256,
        num_classes: int = 4,
        dropout: float = 0.3,
        n_heads: int = 4,
    ):
        super().__init__()

        self.d_model = d_model
        self.num_classes = num_classes

        # Project raw EchoPrime video embeddings -> latent space
        self.video_proj = nn.Linear(video_dim, d_model)

        # Aggregate across all videos for the study
        self.video_aggregator = TransformerAggregator(dim=d_model, n_heads=n_heads, dropout=dropout)

        # Simple MLP classifier on the study-level video embedding
        hidden_dim = d_model * 2
        self.classifier = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes),
        )

    # def encode_study_videos(self, videos: torch.Tensor, return_attn: bool = False) -> torch.Tensor:
    #     """
    #     videos: [N_videos, video_dim]
    #     returns: [d_model] study-level video embedding
    #     """
    #     if videos is None or videos.numel() == 0:
    #         raise ValueError("encode_video_tokens: videos is None")

    #     if torch.isnan(videos).any() or torch.isinf(videos).any():
    #         raise ValueError(f"videos contain NaN or Inf values: {videos}")

    #     v = self.video_proj(videos)              # [N_v, d_model]
    #     v_emb = self.video_aggregator(v, return_attn=return_attn)        # [d_model] - get the cls tokens summarizing all video embeddings
        
    #     return v_emb

    def encode_study_videos(self, videos: torch.Tensor, return_attn: bool = False):
        v = self.video_proj(videos)  # [N_v, d_model]
        if return_attn:
            cls_emb, attn = self.video_aggregator(v, return_attn=True)  # attn: [N_v]
            return cls_emb, attn
        else:
            cls_emb = self.video_aggregator(v, return_attn=False)
            return cls_emb

    def forward(
        self,
        video_0: List[torch.Tensor],     # list of [N0_i, video_dim]
        view_0: List[torch.Tensor],      # unused
        report0_cls: List[torch.Tensor], # unused
        delta_t: torch.Tensor,           # unused
    ) -> torch.Tensor:
        """
        Returns logits: [B, num_classes]
        """
        device = next(self.parameters()).device
        B = len(video_0)

        vid_emb_list = []
        for i in range(B):
            v0 = video_0[i].to(device)           # [N0_i, video_dim]
            vid_emb = self.encode_study_videos(v0, return_attn=False)  # [d_model] 
            vid_emb_list.append(vid_emb)

        vid_batch = torch.stack(vid_emb_list, dim=0)  # [B, d_model]
        logits = self.classifier(vid_batch)           # [B, num_classes]
        
        return logits
    
class ASReportOnlyModel(nn.Module):
    """
    Report-only AS-time model.

    - Input: per-sentence CLS embeddings for study 0 (report0_cls)
    - No videos, no time
    - Pipeline: report_proj -> TransformerAggregator -> classifier

    Forward signature kept compatible with your training loop.
    Only report0_cls is used.
    """

    def __init__(
        self,
        report_dim: int = 768,
        d_model: int = 256,
        num_classes: int = 4,
        dropout: float = 0.3,
        n_heads: int = 4,
    ):
        super().__init__()

        self.d_model = d_model
        self.num_classes = num_classes

        self.report_proj = nn.Linear(report_dim, d_model)

        self.report_aggregator = TransformerAggregator(dim=d_model, n_heads=n_heads, dropout=dropout)

        hidden_dim = d_model * 2
        self.classifier = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes),
        )

    # def encode_study_report(self, report_cls: torch.Tensor) -> torch.Tensor:
    #     """
    #     report_cls: [N_sentences, report_dim]
    #     returns: [d_model] study-level report embedding
    #     """
        
    #     if report_cls is None or report_cls.numel() == 0:
    #         raise ValueError("encode_study_tokens: report is None")

    #     if torch.isnan(report_cls).any() or torch.isinf(report_cls).any():
    #         raise ValueError(f"report_cls contains NaN or Inf values: {report_cls}")

    #     rep = self.report_proj(report_cls)               # [N_s, d_model]
    #     rep_emb = self.report_aggregator(rep)            # [d_model]
        
    #     return rep_emb

    def encode_study_report(self, report_cls: torch.Tensor, return_attn: bool = False):
        rep = self.report_proj(report_cls)  # [N_s, d_model]
        if return_attn:
            cls_emb, attn = self.report_aggregator(rep, return_attn=True)  # attn: [N_s]
            return cls_emb, attn
        else:
            cls_emb = self.report_aggregator(rep, return_attn=False)
            return cls_emb

    def forward(
        self,
        video_0: List[torch.Tensor],     # unused
        view_0: List[torch.Tensor],      # unused
        report0_cls: List[torch.Tensor], # list of [M0_i, report_dim]
        delta_t: torch.Tensor,           # unused
    ) -> torch.Tensor:
        """
        Returns logits: [B, num_classes]
        """
        device = next(self.parameters()).device
        B = len(report0_cls)

        rep_emb_list = []
        for i in range(B):
            r0 = report0_cls[i].to(device)             # [M0_i, report_dim]
            rep_emb = self.encode_study_report(r0)     # [d_model]
            rep_emb_list.append(rep_emb)

        rep_batch = torch.stack(rep_emb_list, dim=0)   # [B, d_model]
        logits = self.classifier(rep_batch)            # [B, num_classes]
        
        return logits
    
class ASTimeOnlyModel(nn.Module):
    """
    Time-only AS-time model.

    - Input: scalar delta_t for each pair (days)
    - No videos, no reports
    - Pipeline: build_time_features(delta_t) -> PosEmbedding -> classifier

    Forward signature kept compatible; only delta_t is used.
    """

    def __init__(
        self,
        d_model: int = 256,
        num_classes: int = 4,
        dropout: float = 0.3,
    ):
        super().__init__()

        self.d_model = d_model
        self.num_classes = num_classes

        # Time encoder config
        self.time_cfg = TimePosConfig(
            emb_dim=d_model,
            seq_length=1,
            pos_emb="temporal",
            n_time_features=1,
            temperature=10000.0,
        )
        self.time_pos = PosEmbedding(self.time_cfg, emb_dim=d_model)

        self.use_log_standardize = True
        self.register_buffer("mean_log", torch.zeros(1))
        self.register_buffer("std_log", torch.ones(1))

        hidden_dim = d_model * 2
        self.classifier = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes),
        )

        self.scalar_proj = nn.Linear(1, d_model)

    # def encode_time_batch(self, delta_t: torch.Tensor) -> torch.Tensor:
    #     """
    #     delta_t: [B] tensor of scalars
    #     returns: [B, d_model] time embeddings
    #     """
    #     if delta_t is None or delta_t.numel() == 0:
    #         raise ValueError("No delta_t values provided.")

    #     if torch.isnan(delta_t).any() or torch.isinf(delta_t).any():
    #         raise ValueError(f"delta_t contains NaN or Inf values: {delta_t}")

    #     if delta_t.ndim > 1:
    #         delta_t = delta_t.view(-1)  # [B]

    #     time_feats = build_time_features(delta_t)   # [B, 1, n_time_features]
    #     t_emb = self.time_pos(time_feats)          # [B, 1, d_model]
    #     t_emb = t_emb.squeeze(1)                   # [B, d_model]
        
    #     return t_emb
    
    @torch.no_grad()
    def set_time_stats(self, mean_log: float, std_log: float):
        self.mean_log.fill_(float(mean_log))
        self.std_log.fill_(float(std_log))
        # return
    
    def encode_time_batch(self, delta_t: torch.Tensor) -> torch.Tensor:
        if delta_t is None or delta_t.numel() == 0:
            raise ValueError("No delta_t values provided.")
        if torch.isnan(delta_t).any() or torch.isinf(delta_t).any():
            raise ValueError(f"delta_t contains NaN or Inf values: {delta_t}")
        if delta_t.ndim > 1:
            delta_t = delta_t.view(-1)  # [B]

        if self.use_log_standardize:
            delta_t = torch.log1p(delta_t)
            delta_t = (delta_t - self.mean_log) / (self.std_log + 1e-8)

        time_feats = build_time_features(delta_t)   # [B, 1, 1]
        t_emb = self.time_pos(time_feats)          # [B, 1, d_model]
        return t_emb.squeeze(1)                    # [B, d_model]

    # def encode_time_batch(self, delta_t: torch.Tensor) -> torch.Tensor:
    #     if delta_t is None or delta_t.numel() == 0:
    #         raise ValueError("No delta_t values provided.")
    #     if torch.isnan(delta_t).any() or torch.isinf(delta_t).any():
    #         raise ValueError(f"delta_t contains NaN or Inf values: {delta_t}")
    #     if delta_t.ndim > 1:
    #         delta_t = delta_t.view(-1)  # [B]

    #     # OPTIONAL: keep log+standardization for fairness
    #     if self.use_log_standardize:
    #         delta_t = torch.log1p(delta_t)
    #         delta_t = (delta_t - self.mean_log) / (self.std_log + 1e-8)

        # # ---- NEW PART ----
        # delta_t = delta_t.unsqueeze(-1)          # [B, 1]
        # return self.scalar_proj(delta_t)

    def forward(
        self,
        video_0: List[torch.Tensor],     # unused
        view_0: List[torch.Tensor],      # unused
        report0_cls: List[torch.Tensor], # unused
        delta_t: torch.Tensor,           # [B]
    ) -> torch.Tensor:
        """
        Returns logits: [B, num_classes]
        """
        device = next(self.parameters()).device
        delta_t = delta_t.to(device)                # [B]

        t_emb = self.encode_time_batch(delta_t)
        logits = self.classifier(t_emb)

        return logits

class ASMLP(nn.Module):
    """
    Multi-modal AS-time model (MLP fusion).

    - Supports any 2 or 3 of: video, report, time.
    - For each used modality, we get a [B, d_model] study-level embedding,
      then concatenate -> [B, M * d_model] (M = #modalities) and feed to an MLP classifier.

    Forward signature is kept compatible with your training loop:
      forward(video_0, view_0, report0_cls, delta_t) -> logits [B, num_classes]
    """

    def __init__(
        self,
        video_dim: int = 512,
        report_dim: int = 768,
        d_model: int = 256,
        num_classes: int = 4,
        dropout: float = 0.3,
        n_heads: int = 4,
        use_video: bool = False,
        use_report: bool = False,
        use_time: bool = False,
    ):
        super().__init__()

        # ---- which modalities are used ----
        self.use_video  = use_video
        self.use_report = use_report
        self.use_time   = use_time

        active_count = int(use_video) + int(use_report) + int(use_time)
        self.active_modalities = active_count

        # if (use_video + use_report + use_time) != 2:
        #     raise ValueError(
        #         "ASBiModalMLP expects EXACTLY two of "
        #         "use_video/use_report/use_time to be True."
        #     )

        self.d_model = d_model
        self.num_classes = num_classes

        # ---- video branch (if used) ----
        if self.use_video:
            self.video_proj = nn.Linear(video_dim, d_model)
            self.video_aggregator = TransformerAggregator(
                dim=d_model,
                n_heads=n_heads,
                dropout=dropout,
            )

        # ---- report branch (if used) ----
        if self.use_report:
            self.report_proj = nn.Linear(report_dim, d_model)
            self.report_aggregator = TransformerAggregator(
                dim=d_model,
                n_heads=n_heads,
                dropout=dropout,
            )

        # ---- time branch (if used) ----
        if self.use_time:
            # # Simple scalar -> d_model projection, same flavour as ASTimeOnlyModel
            # self.scalar_proj = nn.Linear(1, d_model)
            # Time encoder config
            self.time_cfg = TimePosConfig(
                emb_dim=d_model,
                seq_length=1,
                pos_emb="temporal",
                n_time_features=1,
                temperature=10000.0,
            )
            self.time_pos = PosEmbedding(self.time_cfg, emb_dim=d_model)

            self.use_log_standardize = True
            self.register_buffer("mean_log", torch.zeros(1))
            self.register_buffer("std_log", torch.ones(1))

        # ---- classifier: (M * d_model) -> num_classes ----
        fused_dim = self.active_modalities * d_model
        hidden_dim = fused_dim * 2

        self.classifier = nn.Sequential(
            nn.LayerNorm(fused_dim),
            nn.Linear(fused_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes),
        )

    # --------- helpers: encoders per modality ---------
    def encode_study_videos(self, videos: torch.Tensor) -> torch.Tensor:
        """
        videos: [N_v, video_dim]
        returns: [d_model] study-level video embedding
        """
        v = self.video_proj(videos)                    # [N_v, d_model]
        cls_emb = self.video_aggregator(v, return_attn=False)  # [d_model]
        return cls_emb

    def encode_study_report(self, report_cls: torch.Tensor) -> torch.Tensor:
        """
        report_cls: [N_s, report_dim]
        returns: [d_model] study-level report embedding
        """
        rep = self.report_proj(report_cls)             # [N_s, d_model]
        cls_emb = self.report_aggregator(rep, return_attn=False)  # [d_model]
        return cls_emb

    # def encode_time_batch(self, delta_t: torch.Tensor) -> torch.Tensor:
    #     """
    #     delta_t: [B] raw (or preprocessed) time gaps in days.
    #     returns: [B, d_model]
    #     """
    #     if delta_t is None or delta_t.numel() == 0:
    #         raise ValueError("No delta_t values provided.")
    #     if torch.isnan(delta_t).any() or torch.isinf(delta_t).any():
    #         raise ValueError(f"delta_t contains NaN or Inf values: {delta_t}")
    #     if delta_t.ndim > 1:
    #         delta_t = delta_t.view(-1)  # [B]

    #     delta_t = delta_t.unsqueeze(-1)       # [B, 1]
    #     return self.scalar_proj(delta_t)      # [B, d_model]

    # @torch.no_grad()
    # def set_time_stats(self, *args, **kwargs):
    #     """
    #     Kept for API compatibility with time-standardized models.
    #     No-op here (we're using raw/scalar-projected time).
    #     """
    #     return

    @torch.no_grad()
    def set_time_stats(self, mean_log: float, std_log: float):
        self.mean_log.fill_(float(mean_log))
        self.std_log.fill_(float(std_log))
        # return
    
    def encode_time_batch(self, delta_t: torch.Tensor) -> torch.Tensor:
        if delta_t is None or delta_t.numel() == 0:
            raise ValueError("No delta_t values provided.")
        if torch.isnan(delta_t).any() or torch.isinf(delta_t).any():
            raise ValueError(f"delta_t contains NaN or Inf values: {delta_t}")
        if delta_t.ndim > 1:
            delta_t = delta_t.view(-1)  # [B]

        if self.use_log_standardize:
            delta_t = torch.log1p(delta_t)
            delta_t = (delta_t - self.mean_log) / (self.std_log + 1e-8)

        time_feats = build_time_features(delta_t)   # [B, 1, 1]
        t_emb = self.time_pos(time_feats)          # [B, 1, d_model]
        return t_emb.squeeze(1)                    # [B, d_model]

    # --------- forward ---------
    def forward(
        self,
        video_0: List[torch.Tensor],     # list of [N_v_i, video_dim]
        view_0: List[torch.Tensor],      # unused, for API compatibility
        report0_cls: List[torch.Tensor], # list of [N_s_i, report_dim]
        delta_t: torch.Tensor,           # [B]
    ) -> torch.Tensor:
        """
        Returns logits: [B, num_classes]
        """
        device = next(self.parameters()).device

        # Infer batch size safely from available inputs
        B_candidates = []
        if self.use_video:
            B_candidates.append(len(video_0))
        if self.use_report:
            B_candidates.append(len(report0_cls))
        if self.use_time:
            B_candidates.append(delta_t.size(0))
        B = B_candidates[0] if B_candidates else 0

        if B == 0:
            raise ValueError("Empty batch passed to ASMLP")

        # ---- encode each used modality to [B, d_model] ----
        modal_embs = []

        if self.use_video:
            vid_emb_list = []
            for i in range(B):
                v0 = video_0[i].to(device)           # [N_v_i, video_dim]
                vid_emb = self.encode_study_videos(v0)  # [d_model]
                vid_emb_list.append(vid_emb)
            vid_emb = torch.stack(vid_emb_list, dim=0)  # [B, d_model]
            modal_embs.append(vid_emb)

        if self.use_report:
            rep_emb_list = []
            for i in range(B):
                r0 = report0_cls[i].to(device)          # [N_s_i, report_dim]
                rep_emb = self.encode_study_report(r0)  # [d_model]
                rep_emb_list.append(rep_emb)
            rep_emb = torch.stack(rep_emb_list, dim=0)  # [B, d_model]
            modal_embs.append(rep_emb)

        if self.use_time:
            t_emb = self.encode_time_batch(delta_t.to(device))  # [B, d_model]
            modal_embs.append(t_emb)

        # # We asserted exactly two modalities are True → modal_embs has length 2
        # if len(modal_embs) != 2:
        #     raise RuntimeError("ASBiModalMLP internal error: expected 2 modality embeddings.")

        fused = torch.cat(modal_embs, dim=-1)  # [B, 2 * d_model]
        logits = self.classifier(fused)        # [B, num_classes]
        return logits

class ASTimeModel_CrossAttn(nn.Module):
    """
    AS-time model with report/time-guided video interpretation via cross-attention.

    - Video: per-video embeddings (each video = 1 token; all videos kept)
    - Report: per-sentence CLS embeddings -> TransformerAggregator -> 1 report token
    - Time: scalar delta_t -> 1 time token via temporal positional embedding
    - Cross-attn: (report + time) tokens query the video tokens (K/V)
    - Summaries: video / report / time -> concat -> classifier

    Forward signature:
      forward(video_0, view_0, report0_cls, delta_t) -> logits [B, num_classes]
    """

    def __init__(
        self,
        video_dim: int = 512,
        report_dim: int = 768,
        d_model: int = 256,
        num_classes: int = 2,   # early/significant
        dropout: float = 0.3,
        n_heads: int = 4,
        use_video: bool = True,
        use_report: bool = True,
        use_time: bool = True,
    ):
        super().__init__()

        self.d_model = d_model
        self.num_classes = num_classes

        self.use_video = use_video
        self.use_report = use_report
        self.use_time = use_time

        # Project raw embeddings into shared latent space
        self.video_proj = nn.Linear(video_dim, d_model)
        self.report_proj = nn.Linear(report_dim, d_model)

        # Aggregate all report CLS tokens into a single token (like original script style)
        self.report_aggregator = TransformerAggregator(dim=d_model, n_heads=n_heads, dropout=dropout)

        # Time encoding (delta_t -> temporal position embedding)
        self.time_cfg = TimePosConfig(
            emb_dim=d_model,
            seq_length=1,
            pos_emb="temporal",
            n_time_features=1,
            temperature=10000.0,
        )
        self.time_pos = PosEmbedding(self.time_cfg, emb_dim=d_model)
        
        self.use_log_standardize = True
        self.register_buffer("mean_log", torch.zeros(1))
        self.register_buffer("std_log", torch.ones(1))

        # Cross-attention: Q = [report_token, time_token], K/V = video_tokens
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=n_heads,
            dropout=dropout,
            batch_first=True,
        )

        # Lightweight FFN on the cross-attn output
        self.cross_ffn = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

        # Final fusion: video + report + time
        fusion_dim = d_model * 3
        hidden_dim = d_model * 2

        self.classifier = nn.Sequential(
            nn.Linear(fusion_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes),
        )

    @torch.no_grad()
    def set_time_stats(self, mean_log: float, std_log: float):
        self.mean_log.fill_(float(mean_log))
        self.std_log.fill_(float(std_log))

    def encode_video_tokens(self, videos: torch.Tensor) -> torch.Tensor:
        """
        videos: [N_videos, video_dim]
        returns: [N_videos, d_model]
        """
        if not self.use_video:
            raise ValueError("ASTimeModel_CrossAttn requires use_video=True for cross-attention.")

        if videos is None or videos.numel() == 0:
            raise ValueError("No video embeddings available for this study.")

        if torch.isnan(videos).any() or torch.isinf(videos).any():
            raise ValueError(f"videos contain NaN or Inf values: {videos}")

        return self.video_proj(videos)  # [N_v, d_model]

    def encode_report_tokens(self, report_cls: torch.Tensor) -> torch.Tensor:
        """
        Turn report CLS embeddings into a single d_model token via TransformerAggregator.
        report_cls: [N_sentences, report_dim]
        returns:
            [1, d_model] if we have a report and use_report=True
            [0, d_model] (empty) if reports are disabled or missing
        """
        if not self.use_report:
            device = report_cls.device if isinstance(report_cls, torch.Tensor) else "cpu"
            return torch.empty(0, self.d_model, device=device)

        if report_cls is None or report_cls.numel() == 0:
            device = report_cls.device if isinstance(report_cls, torch.Tensor) else "cpu"
            return torch.empty(0, self.d_model, device=device)

        if torch.isnan(report_cls).any() or torch.isinf(report_cls).any():
            raise ValueError(f"report_cls contains NaN or Inf values: {report_cls}")

        rep = self.report_proj(report_cls)           # [M, d_model]
        rep_agg = self.report_aggregator(rep)        # [d_model]
        
        return rep_agg.unsqueeze(0)                  # [1, d_model]  (single token)

    def encode_time_token(self, delta_t: torch.Tensor) -> torch.Tensor:
        """
        delta_t: scalar tensor (single value for this pair)
        returns: [1, d_model] time token
        """
        if not self.use_time:
            raise ValueError("encode_time_token called but use_time=False.")

        if delta_t is None or delta_t.numel() == 0:
            raise ValueError("No time value available for this study.")

        if torch.isnan(delta_t).any() or torch.isinf(delta_t).any():
            raise ValueError(f"delta_t contains NaN or Inf values: {delta_t}")

        if delta_t.ndim == 0:
            delta_t = delta_t.unsqueeze(0)
        elif delta_t.ndim > 1:
            delta_t = delta_t.view(-1)

        if delta_t.numel() != 1:
            raise ValueError(
                f"encode_time_token expects scalar delta_t, got shape {delta_t.shape}"
            )
        
        if self.use_log_standardize:
            delta_t = torch.log1p(delta_t)
            delta_t = (delta_t - self.mean_log) / (self.std_log + 1e-8)

        time_feats = build_time_features(delta_t)   # [1, 1, n_time_features]
        t_emb = self.time_pos(time_feats)          # [1, 1, d_model]
        
        return t_emb.squeeze(0)                    # [1, d_model]

    def forward(
        self,
        video_0: List[torch.Tensor],     # list of [N0_i, video_dim]
        view_0: List[torch.Tensor],      # list of [N0_i], unused but kept for API compatibility
        report0_cls: List[torch.Tensor], # list of [M0_i, report_dim]
        delta_t: torch.Tensor,           # [B]
    ) -> torch.Tensor:
        """
        All inputs come directly from your current dataloader + collate_fn.
        Returns:
            logits: [B, num_classes]
        """
        B = len(video_0)
        device = next(self.parameters()).device

        fusion_list = []

        for i in range(B):
            v0 = video_0[i].to(device)                       # [N_v_i, video_dim]
            video_tokens = self.encode_video_tokens(v0)      # [N_v_i, d_model]

            r0 = report0_cls[i].to(device)
            report_token = self.encode_report_tokens(r0)     # [1, d_model] or [0, d_model]
            has_report = (report_token.size(0) > 0 and self.use_report)

            if self.use_time:
                t_i = delta_t[i].to(device)                  # scalar
                time_token = self.encode_time_token(t_i)     # [1, d_model]
                has_time = True
            else:
                time_token = None
                has_time = False

            # building Q
            q_pieces = []
            if has_report:
                q_pieces.append(report_token)                # [1, d_model]
            if has_time:
                q_pieces.append(time_token)                  # [1, d_model]

            if len(q_pieces) == 0:
                raise ValueError("No Q available.")
            else:
                q = torch.cat(q_pieces, dim=0).unsqueeze(0)  # [1, Lq<=2, d_model]
                k = video_tokens.unsqueeze(0)                # [1, N_v_i, d_model]
                v = video_tokens.unsqueeze(0)                # [1, N_v_i, d_model]

                attn_out, _ = self.cross_attn(q, k, v)       # [1, Lq<=2, d_model]
                attn_out = self.cross_ffn(attn_out)          # [1, Lq<=2, d_model]

                video_summary = attn_out.mean(dim=1).squeeze(0)   # [d_model]

                # Split back into report/time summaries
                if has_report and has_time:
                    # order: [report_token, time_token]
                    rep_out  = attn_out[:, 0:1, :]           # [1, 1, d_model]
                    time_out = attn_out[:, 1:2, :]           # [1, 1, d_model]

                    report_summary = rep_out.squeeze(0).squeeze(0)  # [d_model]
                    time_summary   = time_out.squeeze(0).squeeze(0) # [d_model]

                elif has_report and not has_time:
                    rep_out        = attn_out                # [1, 1, d_model]
                    report_summary = rep_out.squeeze(0).squeeze(0)
                    time_summary   = torch.zeros(self.d_model, device=device)

                elif not has_report and has_time:
                    time_out       = attn_out                # [1, 1, d_model]
                    time_summary   = time_out.squeeze(0).squeeze(0)
                    report_summary = torch.zeros(self.d_model, device=device)

                else:
                    # logically unreachable due to len(q_pieces) check
                    report_summary = torch.zeros(self.d_model, device=device)
                    time_summary   = torch.zeros(self.d_model, device=device)

            fused = torch.cat([video_summary, report_summary, time_summary], dim=-1,)  # [3 * d_model]

            fusion_list.append(fused)

        fusion_batch = torch.stack(fusion_list, dim=0)   # [B, 3 * d_model]
        logits = self.classifier(fusion_batch)           # [B, num_classes]
        
        return logits
    
class CrossAttnFusionBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float):
        super().__init__()
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=n_heads,
            dropout=dropout,
            batch_first=True,
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(4 * d_model, d_model),
            nn.Dropout(dropout),
        )
        self.ln2 = nn.LayerNorm(d_model)

    def forward(
        self,
        q_tokens: torch.Tensor,
        kv_tokens: torch.Tensor,
        key_padding_mask: torch.Tensor,
        return_attn: bool = False,
    ):
        attn_out, attn_w = self.cross_attn(
            q_tokens,
            kv_tokens,
            kv_tokens,
            key_padding_mask=key_padding_mask,
            need_weights=return_attn,
            average_attn_weights=False,
        )
        x = self.ln1(q_tokens + attn_out)
        x = self.ln2(x + self.ffn(x))
        if return_attn:
            return x, attn_w
        return x


class ASTimeModel_CrossAttnV2(nn.Module):
    """
    Latest batched cross-attention architecture.
    Keeps legacy ASTimeModel_CrossAttn intact for backward compatibility.
    """

    def __init__(
        self,
        video_dim: int = 512,
        report_dim: int = 768,
        d_model: int = 256,
        num_classes: int = 2,
        dropout: float = 0.3,
        n_heads: int = 4,
        n_cross_layers: int = 2,
        use_video: bool = True,
        use_report: bool = True,
        use_time: bool = True,
    ):
        super().__init__()

        self.d_model = d_model
        self.num_classes = num_classes
        self.use_video = use_video
        self.use_report = use_report
        self.use_time = use_time

        if not self.use_video:
            raise ValueError("ASTimeModel_CrossAttnV2 requires use_video=True")

        self.video_proj = nn.Linear(video_dim, d_model)
        self.report_proj = nn.Linear(report_dim, d_model)
        self.report_aggregator = TransformerAggregator(dim=d_model, n_heads=n_heads, dropout=dropout)

        self.time_cfg = TimePosConfig(
            emb_dim=d_model,
            seq_length=1,
            pos_emb="temporal",
            n_time_features=1,
            temperature=10000.0,
        )
        self.time_pos = PosEmbedding(self.time_cfg, emb_dim=d_model)

        self.use_log_standardize = True
        self.register_buffer("mean_log", torch.zeros(1))
        self.register_buffer("std_log", torch.ones(1))

        self.null_report_token = nn.Parameter(torch.zeros(1, 1, d_model))
        self.null_time_token = nn.Parameter(torch.zeros(1, 1, d_model))

        self.cross_blocks = nn.ModuleList(
            [
                CrossAttnFusionBlock(d_model=d_model, n_heads=n_heads, dropout=dropout)
                for _ in range(n_cross_layers)
            ]
        )

         # classifier dims depend on whether time is used
        if self.use_time:
            fusion_dim = d_model * 3  # [video, report, time]
        else:
            fusion_dim = d_model * 2  # [video, report] only

        hidden_dim = d_model * 2    
        self.classifier = nn.Sequential(
            nn.LayerNorm(fusion_dim),
            nn.Linear(fusion_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes),
        )

    @torch.no_grad()
    def set_time_stats(self, mean_log: float, std_log: float):
        self.mean_log.fill_(float(mean_log))
        self.std_log.fill_(float(std_log))

    def encode_time_token(self, delta_t: torch.Tensor) -> torch.Tensor:
        if delta_t.ndim == 0:
            delta_t = delta_t.unsqueeze(0)
        elif delta_t.ndim > 1:
            delta_t = delta_t.view(-1)

        if delta_t.numel() != 1:
            raise ValueError(f"encode_time_token expects scalar delta_t, got shape {delta_t.shape}")

        if self.use_log_standardize:
            delta_t = torch.log1p(delta_t)
            delta_t = (delta_t - self.mean_log) / (self.std_log + 1e-8)

        time_feats = build_time_features(delta_t)
        t_emb = self.time_pos(time_feats)
        return t_emb.squeeze(0)

    def _masked_mean(self, x: torch.Tensor, valid_mask: torch.Tensor) -> torch.Tensor:
        valid = valid_mask.unsqueeze(-1).to(dtype=x.dtype)
        summed = (x * valid).sum(dim=1)
        denom = valid.sum(dim=1).clamp(min=1.0)
        return summed / denom

    def forward(
        self,
        video_0: List[torch.Tensor],
        view_0: List[torch.Tensor],
        report0_cls: List[torch.Tensor],
        delta_t: torch.Tensor,
    ) -> torch.Tensor:
        B = len(video_0)
        device = next(self.parameters()).device

        if B == 0:
            raise ValueError("Empty batch passed to ASTimeModel_CrossAttnV2")

        video_lengths = [v.size(0) for v in video_0]
        if min(video_lengths) == 0:
            raise ValueError("No video embeddings available for at least one study.")

        max_n = max(video_lengths)
        video_tokens = torch.zeros(B, max_n, self.d_model, device=device)
        video_valid_mask = torch.zeros(B, max_n, dtype=torch.bool, device=device)

        for i in range(B):
            v_i = _to_device_if_needed(video_0[i], device)
            v_proj = self.video_proj(v_i)
            n_i = v_proj.size(0)
            video_tokens[i, :n_i] = v_proj
            video_valid_mask[i, :n_i] = True

        key_padding_mask = ~video_valid_mask

        report_queries = []
        time_queries = []
        delta_t = _to_device_if_needed(delta_t, device)

        for i in range(B):
            if self.use_report and report0_cls[i].numel() > 0:
                r_i = _to_device_if_needed(report0_cls[i], device)
                rep_proj = self.report_proj(r_i)
                rep_token = self.report_aggregator(rep_proj).view(1, self.d_model)
            else:
                rep_token = self.null_report_token.view(1, self.d_model)
            report_queries.append(rep_token)

             # time token (only if use_time=True)
            if self.use_time:
                t_i = self.encode_time_token(delta_t[i])           # [D]
                time_token = t_i.view(1, self.d_model)            # [1, D]
                time_queries.append(time_token)

        report_q = torch.stack(report_queries, dim=0)
        
        if self.use_time:
            time_q = torch.stack(time_queries, dim=0)  # [B, 1, D]
            q_tokens = torch.cat([report_q, time_q], dim=1)  # [B, 2, D]
        else:
            # no time token at all → single query (report only)
            q_tokens = report_q  # [B, 1, D]


        for block in self.cross_blocks:
            q_tokens = block(
                q_tokens=q_tokens,
                kv_tokens=video_tokens,
                key_padding_mask=key_padding_mask,
            )

         # ---- summaries ----
        video_summary = self._masked_mean(video_tokens, video_valid_mask)  # [B, D]
        report_summary = q_tokens[:, 0, :]                                 # [B, D]

        if self.use_time:
            # q_tokens shape: [B, 2, D] → index 1 exists
            time_summary = q_tokens[:, 1, :]                               # [B, D]
            fused = torch.cat(
                [video_summary, report_summary, time_summary], dim=-1
            )  # [B, 3D]
        else:
            # q_tokens shape: [B, 1, D] → no time token
            fused = torch.cat(
                [video_summary, report_summary], dim=-1
            )  # [B, 2D]

        return self.classifier(fused)

    @torch.no_grad()
    def explain_one(
        self,
        batch: Dict[str, Any],
        sample_idx: int = 0,
        topk: int = 3,
        device: torch.device = None,
    ):
        if device is None:
            device = next(self.parameters()).device

        v0 = _to_device_if_needed(batch["video_0"][sample_idx], device)
        if v0.ndim != 2 or v0.size(0) == 0:
            raise ValueError("explain_one requires non-empty video embeddings for the selected sample")

        r0 = _to_device_if_needed(batch["report0_cls"][sample_idx], device)
        dt = batch["delta_t"][sample_idx]
        dt = _to_device_if_needed(dt, device)

        kept_video_ids = batch.get("video_0_ids", None)
        kept_video_ids = kept_video_ids[sample_idx] if kept_video_ids is not None else None

        video_tokens = self.video_proj(v0).unsqueeze(0)
        video_valid_mask = torch.ones(1, video_tokens.size(1), dtype=torch.bool, device=device)
        key_padding_mask = ~video_valid_mask

        if self.use_report and r0.numel() > 0:
            rep_proj = self.report_proj(r0)
            report_token = self.report_aggregator(rep_proj).view(1, 1, self.d_model)
            has_report = True
        else:
            report_token = self.null_report_token
            has_report = False

        # time token (only if use_time=True)
        if self.use_time:
            time_token = self.encode_time_token(dt).view(1, 1, self.d_model)
            has_time = True
            q_tokens = torch.cat([report_token, time_token], dim=1)  # [1, 2, D]
        else:
            has_time = False
            q_tokens = report_token  # [1, 1, D]

        last_attn = None
        for block in self.cross_blocks:
            q_tokens, attn_w = block(
                q_tokens=q_tokens,
                kv_tokens=video_tokens,
                key_padding_mask=key_padding_mask,
                return_attn=True,
            )
            last_attn = attn_w

        if last_attn is None:
            raise RuntimeError("No attention weights were produced in explain_one")

        # last_attn: [1, heads, 2, N_v]
        attn_mean = last_attn.mean(dim=1).squeeze(0)
        report_attn = attn_mean[0]
        time_attn = attn_mean[1] if (self.use_time and attn_mean.size(0) > 1) else None

        active_attn = []
        if has_report:
            active_attn.append(report_attn)
        if has_time and (time_attn is not None):
            active_attn.append(time_attn)

        if len(active_attn) > 0:
            combined_attn = torch.stack(active_attn, dim=0).mean(dim=0)
        else:
            combined_attn = attn_mean.mean(dim=0)  # fallback

        k = min(topk, combined_attn.numel())
        top_idx = torch.topk(combined_attn, k=k).indices.cpu().tolist()

        out = {
            "video_topk_idx": top_idx,
            "video_topk_attn": combined_attn[top_idx].detach().cpu().tolist(),
            "video_topk_stream_ids": [kept_video_ids[i] for i in top_idx] if kept_video_ids is not None else None,
            "video_stream_ids": kept_video_ids,
            "cross_report_attn": report_attn.detach().cpu().tolist(),
            "cross_time_attn": time_attn.detach().cpu().tolist() if time_attn is not None else None,
        }

        if has_report:
            report_top_idx = torch.topk(report_attn, k=k).indices.cpu().tolist()
            out["cross_report_topk_idx"] = report_top_idx
            out["cross_report_topk_attn"] = report_attn[report_top_idx].detach().cpu().tolist()

        if has_time and (time_attn is not None):
            time_top_idx = torch.topk(time_attn, k=k).indices.cpu().tolist()
            out["cross_time_topk_idx"] = time_top_idx
            out["cross_time_topk_attn"] = time_attn[time_top_idx].detach().cpu().tolist()

        return out

# ===========================================================================
# AS-TIME (proposed model) -- our contribution
# ===========================================================================
class ASTIME(nn.Module):
    """AS-TIME: time-conditioned multimodal model for AS progression.

    Hierarchical fusion conditioned on the inter-visit interval delta_t:
      Stage 1 (token aggregation): attention-based multiple-instance
        learning (MIL) over video and report tokens within each modality.
      Stage 2 (modality gating): a delta_t-conditioned gate weights the
        aggregated video and report embeddings using their cosine
        similarity and the time embedding (video weight w in [0, 1],
        report weight 1 - w).
      Stage 3 (expert routing): a delta_t-guided mixture-of-experts (MoE)
        head produces the final prediction.
    """

    def __init__(
        self,
        video_dim: int = 512,
        report_dim: int = 768,
        d_model: int = 256,
        num_classes: int = 2,
        dropout: float = 0.3,
        n_experts: int = 12,
        condition_moe_on_time: bool = True,
        zero_time_signal: bool = False,
        use_video: bool = True,
        use_report: bool = True,
        use_time: bool = True,
    ):
        super().__init__()

        if not (use_video and use_report and use_time):
            raise ValueError(
                "ASTIME is implemented for use_video=True, "
                "use_report=True, use_time=True only."
            )

        self.d_model = d_model
        self.num_classes = num_classes
        self.use_video = use_video
        self.use_report = use_report
        self.use_time = use_time
        self.condition_moe_on_time = bool(condition_moe_on_time)
        # Temporal ablation (Table 2): zero out delta_t everywhere.
        self.zero_time_signal = bool(zero_time_signal)


        # ---- projections ----
        self.video_proj = nn.Linear(video_dim, d_model)
        self.report_proj = nn.Linear(report_dim, d_model)

        # ---- time encoder ----
        self.time_cfg = TimePosConfig(
            emb_dim=d_model,
            seq_length=1,
            pos_emb="temporal",
            n_time_features=1,
            temperature=10000.0,
        )
        self.time_pos = PosEmbedding(self.time_cfg, emb_dim=d_model)
        self.use_log_standardize = True
        self.register_buffer("mean_log", torch.zeros(1))
        self.register_buffer("std_log", torch.ones(1))

        # ---- MIL attention scorers ----
        hidden = max(64, d_model // 2)
        self.video_attn = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.Tanh(),
            nn.Linear(hidden, 1),
        )
        self.report_attn = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.Tanh(),
            nn.Linear(hidden, 1),
        )

        gate_in_dim = d_model + 1  # [time_summary, cos_sim]
        self.modality_gate = nn.Sequential(
            nn.LayerNorm(gate_in_dim),
            nn.Linear(gate_in_dim, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, 1),
        )

        # ---- MoE head ----
        fused_dim = 2 * d_model  # [fused_modality (D) + time_summary (D)]
        moe_gate_in_dim = fused_dim + (1 if self.condition_moe_on_time else 0)
        self.moe_gate = nn.Sequential(
            nn.LayerNorm(moe_gate_in_dim),
            nn.Linear(moe_gate_in_dim, n_experts),
        )
        self.moe_experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.LayerNorm(fused_dim),
                    nn.Linear(fused_dim, 2 * d_model),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(2 * d_model, num_classes),
                )
                for _ in range(n_experts)
            ]
        )

        # for debugging / analysis
        self._last_explain = {}

    # ---------------- time stats ----------------
    @torch.no_grad()
    def set_time_stats(self, mean_log: float, std_log: float):
        self.mean_log.fill_(float(mean_log))
        self.std_log.fill_(float(std_log))

    # ---------------- time encoder ----------------
    def _encode_time_batch(self, delta_t: torch.Tensor) -> torch.Tensor:
        """
        delta_t: [B] -> [B, d_model]
        """
        if delta_t.ndim > 1:
            delta_t = delta_t.view(-1)

        if self.use_log_standardize:
            delta_t = torch.log1p(torch.clamp(delta_t, min=0.0))
            delta_t = (delta_t - self.mean_log) / (self.std_log + 1e-8)

        time_feats = build_time_features(delta_t)   # [B, 1, n_time_features]
        t_emb = self.time_pos(time_feats)          # [B, 1, d_model]
        return t_emb.squeeze(1)                    # [B, d_model]

    # ---------------- padding + projection ----------------
    def _pad_and_project(
        self,
        seq_list: List[torch.Tensor],
        proj: nn.Linear,
        device: torch.device,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Pads a list of [Ni, F] tensors to [B, maxN, d_model] with mask.
        Returns:
          tokens: [B, maxN, d_model]
          valid_mask: [B, maxN] (bool)
        """
        lengths = [s.size(0) for s in seq_list]
        max_len = max(lengths)
        B = len(seq_list)

        tokens = torch.zeros(B, max_len, self.d_model, device=device)
        valid_mask = torch.zeros(B, max_len, dtype=torch.bool, device=device)

        for i, s in enumerate(seq_list):
            s = _to_device_if_needed(s, device)
            if s.numel() == 0:
                continue
            proj_s = proj(s)              # [Ni, d_model]
            n_i = proj_s.size(0)
            tokens[i, :n_i] = proj_s
            valid_mask[i, :n_i] = True

        return tokens, valid_mask

    # ---------------- MIL pooling ----------------
    def _mil_pool_with_attn(
        self,
        tokens: torch.Tensor,
        valid_mask: torch.Tensor,
        scorer: nn.Module,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        tokens: [B, L, D]
        valid_mask: [B, L]
        scorer: small net producing [B, L, 1] scores
        Returns:
          pooled: [B, D]
          attn:   [B, L]
        """
        scores = scorer(tokens).squeeze(-1)                  # [B, L]
        scores = scores.masked_fill(~valid_mask, float("-inf"))
        attn = torch.softmax(scores, dim=1)                  # [B, L]
        attn = attn * valid_mask.to(attn.dtype)              # zero out padded
        denom = attn.sum(dim=1, keepdim=True).clamp(min=1e-8)
        attn = attn / denom
        pooled = torch.bmm(attn.unsqueeze(1), tokens).squeeze(1)  # [B, D]
        return pooled, attn

    # ---------------- explainability getter ----------------
    def get_last_explainability(self):
        return self._last_explain

    # ---------------- forward ----------------
    def forward(
        self,
        video_0: List[torch.Tensor],     # list of [N_v_i, video_dim]
        view_0: List[torch.Tensor],      # unused, kept for API compatibility
        report0_cls: List[torch.Tensor], # list of [N_s_i, report_dim]
        delta_t: torch.Tensor,           # [B]
    ) -> torch.Tensor:
        B = len(video_0)
        device = next(self.parameters()).device

        if B == 0:
            raise ValueError("Empty batch passed to ASTIME")

        # ---- video & report tokens ----
        video_tokens, video_valid_mask = self._pad_and_project(
            video_0, self.video_proj, device
        )
        report_tokens, report_valid_mask = self._pad_and_project(
            report0_cls, self.report_proj, device
        )

        if not video_valid_mask.any():
            raise ValueError("No video embeddings available for any study.")
        if not report_valid_mask.any():
            raise ValueError("No report embeddings available for any study.")
        
        # ---- time embeddings ----
        delta_t = _to_device_if_needed(delta_t, device)   # [B]
        if delta_t.ndim > 1:
            delta_t = delta_t.view(-1)

        time_summary = self._encode_time_batch(delta_t)

        # Temporal ablation: remove all delta_t signal (Table 2, "no temporal signaling").
        if self.zero_time_signal:
            time_summary = torch.zeros_like(time_summary)
            delta_t = torch.zeros_like(delta_t)

        # ---- Stage 1: attention-based MIL within each modality ----
        video_mil, video_attn = self._mil_pool_with_attn(
            video_tokens, video_valid_mask, self.video_attn
        )
        report_mil, report_attn = self._mil_pool_with_attn(
            report_tokens, report_valid_mask, self.report_attn
        )

        # ---- Stage 2: delta_t-conditioned gated modality fusion ----
        # gate on the time embedding + video/report cosine similarity
        cos_sim = F.cosine_similarity(video_mil, report_mil, dim=-1, eps=1e-8).unsqueeze(-1)  # [B, 1]
        gate_in = torch.cat([time_summary, cos_sim], dim=-1)  # [B, D + 1]

        video_weight = torch.sigmoid(self.modality_gate(gate_in)).squeeze(-1).clamp(0.0, 1.0)  # [B]
        report_weight = 1.0 - video_weight

        fused_modality = (
            video_weight.unsqueeze(-1) * video_mil
            + report_weight.unsqueeze(-1) * report_mil
        )  # [B, D]
        fused_repr = torch.cat([fused_modality, time_summary], dim=-1)  # [B, 2D]

        # ---- Stage 3: delta_t-guided mixture-of-experts ----
        if self.condition_moe_on_time:
            moe_gate_in = torch.cat([fused_repr, delta_t.view(B, 1)], dim=-1)  # [B, 2D+1]
        else:
            moe_gate_in = fused_repr                                      # [B, 2D]

        moe_weight = torch.softmax(self.moe_gate(moe_gate_in), dim=-1)    # [B, E]
        expert_logits = torch.stack(
            [expert(fused_repr) for expert in self.moe_experts], dim=1
        )  # [B, E, C]
        logits = (moe_weight.unsqueeze(-1) * expert_logits).sum(dim=1)    # [B, C]

        # store multi-level explainability (detached): instance attention,
        # modality gating weights, and expert routing probabilities.
        self._last_explain = {
            "video_attn": video_attn.detach(),
            "report_attn": report_attn.detach(),
            "video_weight": video_weight.detach(),
            "report_weight": report_weight.detach(),
            "moe_weight": moe_weight.detach(),
        }

        return logits


# ===========================================================================
# Factory
# ===========================================================================
def ASTimeModel(
    video_dim: int = 512,
    report_dim: int = 768,
    d_model: int = 256,
    num_classes: int = 2,
    dropout: float = 0.3,
    n_heads: int = 4,
    use_video: bool = True,
    use_report: bool = True,
    use_time: bool = True,
    n_cross_layers: int = 2,
    use_tri_top: bool = False,
    use_bimodal_mlp: bool = False,
    n_experts: int = 12,
) -> nn.Module:
    """Factory returning the model selected by the modality / fusion flags.

    Proposed model
    --------------
    use_video + use_report + use_time + use_tri_top -> ASTIME
        The proposed time-conditioned multimodal model (MIL token
        aggregation -> delta_t-conditioned gated fusion -> delta_t-guided MoE).

    Baselines (Table 1)
    -------------------
    use_video + use_report + use_time (default)      -> ASTimeModel_CrossAttnV2 (Cross-Attention)
    ... + use_bimodal_mlp                            -> ASMLP (Concatenation)
    use_video + use_time                             -> ASTimeModel_CrossAttn  (EchoPrime + PosEmb)
    use_video only                                   -> ASVideoOnlyModel
    use_report only                                  -> ASReportOnlyModel
    use_time only                                    -> ASTimeOnlyModel        (Positional Embedding)
    """

    # ----- Proposed model and tri-modal baselines: video + report + time -----
    if use_video and use_report and use_time:
        if use_tri_top:
            print("Using AS-TIME (MIL -> gated fusion -> MoE).")
            return ASTIME(
                video_dim=video_dim,
                report_dim=report_dim,
                d_model=d_model,
                num_classes=num_classes,
                dropout=dropout,
                n_experts=n_experts,
                use_video=True,
                use_report=True,
                use_time=True,
            )
        elif use_bimodal_mlp:
            print("Using Concatenation baseline (ASMLP feature fusion).")
            return ASMLP(
                video_dim=video_dim,
                report_dim=report_dim,
                d_model=d_model,
                num_classes=num_classes,
                dropout=dropout,
                n_heads=n_heads,
                use_video=use_video,
                use_report=use_report,
                use_time=use_time,
            )
        else:
            print("Using Cross-Attention baseline.")
            return ASTimeModel_CrossAttnV2(
                video_dim=video_dim,
                report_dim=report_dim,
                d_model=d_model,
                num_classes=num_classes,
                dropout=dropout,
                n_heads=n_heads,
                n_cross_layers=n_cross_layers,
                use_video=True,
                use_report=True,
                use_time=True,
            )

    # Cross-Attention bimodal: video + report, NO time
    if use_video and use_report and (not use_time) and (not use_tri_top):
        print("Using Cross-Attention baseline (video+report, no time).")
        return ASTimeModel_CrossAttnV2(
            video_dim=video_dim,
            report_dim=report_dim,
            d_model=d_model,
            num_classes=num_classes,
            dropout=dropout,
            n_heads=n_heads,
            n_cross_layers=n_cross_layers,
            use_video=True,
            use_report=True,
            use_time=False,
        )

    # Concatenation baseline on any two modalities
    if use_bimodal_mlp and (use_video + use_report + use_time == 2):
        print("Using Concatenation baseline (ASMLP feature fusion).")
        return ASMLP(
            video_dim=video_dim,
            report_dim=report_dim,
            d_model=d_model,
            num_classes=num_classes,
            dropout=dropout,
            n_heads=n_heads,
            use_video=use_video,
            use_report=use_report,
            use_time=use_time,
        )

    # Video-only baseline
    if use_video and not use_report and not use_time:
        print("Using Video-only baseline.")
        return ASVideoOnlyModel(
            video_dim=video_dim,
            d_model=d_model,
            num_classes=num_classes,
            dropout=dropout,
            n_heads=n_heads,
        )

    # Report-only baseline
    if (not use_video) and use_report and (not use_time):
        print("Using Report-only baseline.")
        return ASReportOnlyModel(
            report_dim=report_dim,
            d_model=d_model,
            num_classes=num_classes,
            dropout=dropout,
            n_heads=n_heads,
        )

    # Time-only baseline (Positional Embedding)
    if (not use_video) and (not use_report) and use_time:
        print("Using Time-only baseline (Positional Embedding).")
        return ASTimeOnlyModel(
            d_model=d_model,
            num_classes=num_classes,
            dropout=dropout,
        )

    # Video + time baseline (EchoPrime + PosEmb)
    if use_video and (not use_report) and use_time:
        print("Using EchoPrime + PosEmb baseline (video + time).")
        return ASTimeModel_CrossAttn(
            video_dim=video_dim,
            report_dim=report_dim,  # unused if use_report=False
            d_model=d_model,
            num_classes=num_classes,
            dropout=dropout,
            n_heads=n_heads,
            use_video=True,
            use_report=False,
            use_time=True,
        )


    # Any other combo (e.g., video+time only, video+report only) not implemented yet
    raise ValueError(
        f"Model combination not implemented: use_video={use_video}, "
        f"use_report={use_report}, use_time={use_time}"
    )