models/scope_pooler.py

"""
Scope-Aware Pooler

Extracts semantic regions from palette using scope markers (0=START, 1=END).
Implements exact scope matching via stack-based algorithm.
"""

import logging
import torch
import torch.nn as nn
from typing import List, Tuple, NamedTuple
from dataclasses import dataclass


class RegionMetadata(NamedTuple):
    """
    Metadata about detected semantic regions

    Fields:
    - masks: BoolTensor[R, H, W] - spatial masks for each region
    - starts: List[int] - flattened start indices
    - ends: List[int] - flattened end indices
    - depths: List[int] - nesting depth of each region
    - types: List[str] - region type hints
    """
    masks: torch.Tensor
    starts: List[int]
    ends: List[int]
    depths: List[int]
    types: List[str]


class ScopeImbalanceError(Exception):
    """Raised when scope markers are critically unbalanced"""
    pass


class ScopePooler(nn.Module):
    """
    Extract semantic regions from palette using scope markers

    This module identifies code scopes (functions, loops, classes, etc.)
    by matching START_OF_SCOPE (0) and END_OF_SCOPE (1) tokens.

    Algorithm:
    1. Flatten palette to 1D sequence
    2. Stack-based matching of scope markers
    3. Extract features for each matched region
    4. Pool features via mean+max aggregation

    Edge Cases Handled:
    - Unbalanced scopes (warning + best-effort matching)
    - Nested scopes (via stack depth tracking)
    - No scopes found (fallback to uniform grid)
    - Empty regions (skip + warning)
    """

    def __init__(
        self,
        hidden_dim: int = 768,
        min_region_size: int = 2,
        fallback_grid_size: int = 4
    ):
        """
        Args:
            hidden_dim: Feature dimension
            min_region_size: Minimum tokens per region
            fallback_grid_size: Grid size when no scopes found
        """
        super().__init__()

        self.hidden_dim = hidden_dim
        self.min_region_size = min_region_size
        self.fallback_grid_size = fallback_grid_size

        # Learned pooling projection
        # Concat [mean, max] then project back to hidden_dim
        self.pool_proj = nn.Linear(hidden_dim * 2, hidden_dim)

    def forward(
        self,
        features: torch.Tensor,  # (B, H, W, D)
        palette: torch.Tensor    # (B, H, W)
    ) -> Tuple[torch.Tensor, List[RegionMetadata]]:
        """
        Extract semantic regions and pool features

        Args:
            features: (B, H, W, D) - ViT output features
            palette: (B, H, W) - palette indices

        Returns:
            regions: (B, R, D) - per-region pooled features
            metadata: List[RegionMetadata] - one per batch item

        Guarantees:
            - R >= 1 always (at least one region)
            - All regions non-empty
            - Features normalized (unit norm)
        """
        B, H, W, D = features.shape
        assert palette.shape == (B, H, W), f"Shape mismatch: features{features.shape} vs palette{palette.shape}"
        assert D == self.hidden_dim, f"Hidden dim mismatch: {D} != {self.hidden_dim}"

        all_regions = []
        all_metadata = []

        for b in range(B):
            feat_b = features[b]  # (H, W, D)
            pal_b = palette[b]    # (H, W)

            # Extract regions for this sample
            regions_b, meta_b = self._extract_regions_single(feat_b, pal_b, H, W)

            all_regions.append(regions_b)  # (R_b, D)
            all_metadata.append(meta_b)

        # Pad to max number of regions in batch
        max_regions = max(r.shape[0] for r in all_regions)
        padded_regions = []

        for regions_b in all_regions:
            R_b = regions_b.shape[0]
            if R_b < max_regions:
                # Pad with zeros
                padding = torch.zeros(
                    max_regions - R_b, D,
                    device=regions_b.device,
                    dtype=regions_b.dtype
                )
                regions_b = torch.cat([regions_b, padding], dim=0)
            padded_regions.append(regions_b)

        batched_regions = torch.stack(padded_regions, dim=0)  # (B, R_max, D)

        return batched_regions, all_metadata

    def _extract_regions_single(
        self,
        features: torch.Tensor,  # (H, W, D)
        palette: torch.Tensor,   # (H, W)
        H: int,
        W: int
    ) -> Tuple[torch.Tensor, RegionMetadata]:
        """
        Extract regions from a single sample

        Returns:
            regions: (R, D) - pooled features
            metadata: RegionMetadata
        """
        # 1. Flatten to sequence
        seq = palette.flatten()  # (H*W,)
        features_flat = features.view(-1, self.hidden_dim)  # (H*W, D)

        # 2. Match scopes
        try:
            scope_pairs, depths = self._match_scopes(seq)
        except ScopeImbalanceError as e:
            # Critical error - scopes too broken to recover
            logging.warning(f"{e}. Using fallback uniform grid.")
            scope_pairs, depths = self._fallback_uniform_grid(H, W)

        # 3. Filter invalid regions
        valid_pairs = []
        valid_depths = []
        for (start, end), depth in zip(scope_pairs, depths):
            if (end - start + 1) >= self.min_region_size:
                valid_pairs.append((start, end))
                valid_depths.append(depth)

        if not valid_pairs:
            # No valid regions - use full sequence
            valid_pairs = [(0, H*W - 1)]
            valid_depths = [0]

        # 4. Extract features for each region
        region_features = []
        region_masks = []
        starts = []
        ends = []

        for (start, end) in valid_pairs:
            # Extract features in range
            region_feat = features_flat[start:end+1]  # (L, D)

            # Pool: mean + max
            mean_pool = region_feat.mean(dim=0)  # (D,)
            max_pool = region_feat.max(dim=0)[0]  # (D,)

            # Concatenate and project
            combined = torch.cat([mean_pool, max_pool], dim=0)  # (2D,)
            pooled = self.pool_proj(combined)  # (D,)

            # Normalize
            pooled = torch.nn.functional.normalize(pooled, dim=0)

            region_features.append(pooled)

            # Create mask
            mask = torch.zeros(H * W, dtype=torch.bool, device=palette.device)
            mask[start:end+1] = True
            mask_2d = mask.view(H, W)
            region_masks.append(mask_2d)

            starts.append(start)
            ends.append(end)

        # Stack regions
        regions = torch.stack(region_features, dim=0)  # (R, D)
        masks = torch.stack(region_masks, dim=0)  # (R, H, W)

        # Create metadata
        types = ['scope'] * len(valid_pairs)  # Generic type for now
        metadata = RegionMetadata(
            masks=masks,
            starts=starts,
            ends=ends,
            depths=valid_depths,
            types=types
        )

        return regions, metadata

    def _match_scopes(
        self,
        seq: torch.Tensor  # (N,)
    ) -> Tuple[List[Tuple[int, int]], List[int]]:
        """
        Stack-based scope matching

        Returns:
            pairs: List of (start_idx, end_idx) tuples
            depths: List of nesting depths

        Algorithm:
        - Maintain stack of open scope indices
        - When seeing START (0), push index
        - When seeing END (1), pop and create pair
        - Track depth = current stack size

        Edge Cases:
        - Unmatched START: close at sequence end
        - Unmatched END: skip with warning
        - No scopes: return empty list (caller handles)
        """
        START_OF_SCOPE = 0
        END_OF_SCOPE = 1

        stack = []  # Stack of (index, depth)
        pairs = []
        depths = []

        seq_np = seq.cpu().numpy()  # Faster iteration

        for i, token in enumerate(seq_np):
            if token == START_OF_SCOPE:
                # Open new scope
                depth = len(stack)
                stack.append((i, depth))

            elif token == END_OF_SCOPE:
                # Close scope
                if stack:
                    start_idx, depth = stack.pop()
                    pairs.append((start_idx, i))
                    depths.append(depth)
                else:
                    # Unmatched END - skip
                    logging.warning(f"Unmatched END_OF_SCOPE at position {i}")

        # Handle unmatched STARTs
        if stack:
            logging.warning(f"{len(stack)} unmatched START_OF_SCOPE tokens")
            # Close them at sequence end
            seq_len = len(seq_np)
            for start_idx, depth in stack:
                pairs.append((start_idx, seq_len - 1))
                depths.append(depth)

        # Validate: check for severe imbalance
        num_starts = (seq == START_OF_SCOPE).sum().item()
        num_ends = (seq == END_OF_SCOPE).sum().item()

        if abs(num_starts - num_ends) > max(num_starts, num_ends) * 0.5:
            # More than 50% imbalance - critical error
            raise ScopeImbalanceError(
                f"Severe scope imbalance: {num_starts} starts vs {num_ends} ends"
            )

        return pairs, depths

    def _fallback_uniform_grid(
        self,
        H: int,
        W: int
    ) -> Tuple[List[Tuple[int, int]], List[int]]:
        """
        Fallback when scope matching fails

        Returns uniform grid of regions

        Args:
            H, W: palette dimensions

        Returns:
            pairs: List of (start, end) for grid cells
            depths: All depth=0 (flat)
        """
        total = H * W
        grid_size = self.fallback_grid_size
        region_size = total // grid_size

        pairs = []
        for i in range(grid_size):
            start = i * region_size
            end = (i + 1) * region_size - 1 if i < grid_size - 1 else total - 1
            pairs.append((start, end))

        depths = [0] * grid_size

        return pairs, depths

    def visualize_regions(
        self,
        palette: torch.Tensor,  # (H, W)
        metadata: RegionMetadata
    ) -> str:
        """
        Generate human-readable visualization of regions

        Returns: String representation
        """
        H, W = palette.shape
        output = []
        output.append(f"Detected {len(metadata.starts)} regions:")

        for i, (start, end, depth) in enumerate(zip(
            metadata.starts,
            metadata.ends,
            metadata.depths
        )):
            region_size = end - start + 1
            indent = "  " * depth
            output.append(
                f"{indent}Region {i}: [{start:4d}, {end:4d}] "
                f"(size={region_size:3d}, depth={depth})"
            )

        return "\n".join(output)