"""
GAP-008: Learned Type Projections for Schema Matching
═══════════════════════════════════════════════════════════════════

Research Reference: README_RESEARCH_3.md, Lines 17-20
Target: MRR 0.866 vs 0.45 baseline (92% improvement)

This module implements learned type projections that encode data types
into semantic embeddings, enabling type-aware matching that goes beyond
random initialization.

Key Features:
- Learned type embeddings via contrastive learning
- Type-aware semantic matching
- Support for LoRA-style efficient fine-tuning
- Integration with existing embedding pipeline

Dependencies:
    pip install torch numpy
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional, Union
import random

import numpy as np

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.utils.data import Dataset, DataLoader
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    torch = None
    nn = None
    F = None


# =============================================================================
# DATA STRUCTURES
# =============================================================================

@dataclass
class TypeProjectionConfig:
    """Configuration for type projection model."""
    
    # Embedding dimensions
    base_embedding_dim: int = 384  # Match MiniLM output
    type_embedding_dim: int = 64   # Learnable type embedding
    projection_dim: int = 384      # Output dimension
    
    # Type vocabulary
    type_vocab_size: int = 50      # Number of distinct types
    
    # Training config
    learning_rate: float = 1e-4
    batch_size: int = 32
    num_epochs: int = 10
    temperature: float = 0.07     # Contrastive loss temperature
    
    # Model architecture
    use_layer_norm: bool = True
    dropout_rate: float = 0.1


@dataclass
class TypeMatchPair:
    """A pair of fields for contrastive learning."""
    
    source_text: str
    source_type: str
    target_text: str
    target_type: str
    is_match: bool  # True if they should match
    
    def to_dict(self) -> dict:
        return {
            "source_text": self.source_text,
            "source_type": self.source_type,
            "target_text": self.target_text,
            "target_type": self.target_type,
            "is_match": self.is_match,
        }


# =============================================================================
# TYPE VOCABULARY
# =============================================================================

class TypeVocabulary:
    """
    Vocabulary for data types.
    
    Maps string type names to integer IDs for embedding lookup.
    """
    
    # Standard type mappings
    STANDARD_TYPES = {
        # String types
        "string": 1, "text": 1, "varchar": 1, "char": 1, "str": 1,
        "nvarchar": 1, "clob": 1, "longtext": 1,
        
        # Integer types
        "integer": 2, "int": 2, "bigint": 3, "smallint": 4, "tinyint": 5,
        "long": 3, "short": 4,
        
        # Decimal types
        "decimal": 6, "float": 7, "double": 8, "real": 7, 
        "number": 6, "numeric": 6, "money": 9,
        
        # Boolean
        "boolean": 10, "bool": 10, "bit": 10,
        
        # Date/Time
        "date": 11, "datetime": 12, "timestamp": 12, "time": 13,
        "datetime2": 12, "datetimeoffset": 14,
        
        # Binary
        "binary": 15, "blob": 16, "bytes": 15, "bytea": 15,
        "varbinary": 15, "image": 16,
        
        # Complex types
        "array": 17, "list": 17, "set": 18,
        "object": 19, "struct": 19, "record": 19, "map": 20,
        "json": 21, "jsonb": 21, "xml": 22,
        
        # Special
        "uuid": 23, "guid": 23,
        "enum": 24,
        "null": 0,
        "unknown": 25,
    }
    
    def __init__(self):
        self._type_to_id = dict(self.STANDARD_TYPES)
        self._id_to_type = {}
        for type_name, type_id in self.STANDARD_TYPES.items():
            if type_id not in self._id_to_type:
                self._id_to_type[type_id] = type_name
        self._next_id = max(self.STANDARD_TYPES.values()) + 1
    
    def get_type_id(self, type_name: str) -> int:
        """Get the ID for a type name."""
        normalized = type_name.lower().strip()
        
        if normalized in self._type_to_id:
            return self._type_to_id[normalized]
        
        # Check for partial matches (e.g., "varchar(255)" -> "varchar")
        for known_type in self._type_to_id:
            if normalized.startswith(known_type):
                return self._type_to_id[known_type]
        
        # Unknown type - assign new ID
        self._type_to_id[normalized] = self._next_id
        self._id_to_type[self._next_id] = normalized
        self._next_id += 1
        
        return self._type_to_id[normalized]
    
    def get_type_name(self, type_id: int) -> str:
        """Get the type name for an ID."""
        return self._id_to_type.get(type_id, "unknown")
    
    @property
    def vocab_size(self) -> int:
        """Get the vocabulary size."""
        return self._next_id


# =============================================================================
# TYPE PROJECTION MODEL
# =============================================================================

if TORCH_AVAILABLE:
    
    class TypeProjectionLayer(nn.Module):
        """
        Learnable type projection that combines base embeddings with type info.
        
        Architecture:
        1. Type embedding lookup
        2. Concatenation with base embedding
        3. Projection to output dimension
        4. Optional layer normalization
        """
        
        def __init__(self, config: TypeProjectionConfig):
            super().__init__()
            self.config = config
            
            # Type embedding
            self.type_embedding = nn.Embedding(
                config.type_vocab_size,
                config.type_embedding_dim,
                padding_idx=0,
            )
            
            # Projection layers
            input_dim = config.base_embedding_dim + config.type_embedding_dim
            
            self.projection = nn.Sequential(
                nn.Linear(input_dim, config.projection_dim * 2),
                nn.GELU(),
                nn.Dropout(config.dropout_rate),
                nn.Linear(config.projection_dim * 2, config.projection_dim),
            )
            
            if config.use_layer_norm:
                self.layer_norm = nn.LayerNorm(config.projection_dim)
            else:
                self.layer_norm = None
            
            # Initialize weights
            self._init_weights()
        
        def _init_weights(self):
            """Initialize weights for better training."""
            nn.init.normal_(self.type_embedding.weight, mean=0, std=0.02)
            
            for module in self.projection:
                if isinstance(module, nn.Linear):
                    nn.init.xavier_uniform_(module.weight)
                    if module.bias is not None:
                        nn.init.zeros_(module.bias)
        
        def forward(
            self,
            base_embeddings: torch.Tensor,  # (batch, base_dim)
            type_ids: torch.Tensor,          # (batch,)
        ) -> torch.Tensor:
            """
            Project base embeddings with type information.
            
            Args:
                base_embeddings: Base semantic embeddings (batch, base_dim)
                type_ids: Integer type IDs (batch,)
                
            Returns:
                Type-aware embeddings (batch, projection_dim)
            """
            # Get type embeddings
            type_emb = self.type_embedding(type_ids)  # (batch, type_dim)
            
            # Concatenate
            combined = torch.cat([base_embeddings, type_emb], dim=-1)
            
            # Project
            output = self.projection(combined)
            
            # Normalize
            if self.layer_norm is not None:
                output = self.layer_norm(output)
            
            return output
    
    
    class ContrastiveTypeModel(nn.Module):
        """
        Full model for contrastive type learning.
        
        Uses InfoNCE-style contrastive loss to learn type-aware embeddings.
        """
        
        def __init__(self, config: TypeProjectionConfig):
            super().__init__()
            self.config = config
            self.projection = TypeProjectionLayer(config)
            self.temperature = config.temperature
        
        def forward(
            self,
            source_emb: torch.Tensor,
            source_types: torch.Tensor,
            target_emb: torch.Tensor,
            target_types: torch.Tensor,
        ) -> tuple[torch.Tensor, torch.Tensor]:
            """
            Compute type-aware embeddings for source and target.
            
            Returns:
                Tuple of (source_projected, target_projected)
            """
            source_proj = self.projection(source_emb, source_types)
            target_proj = self.projection(target_emb, target_types)
            
            return source_proj, target_proj
        
        def compute_similarity(
            self,
            source_proj: torch.Tensor,
            target_proj: torch.Tensor,
        ) -> torch.Tensor:
            """
            Compute cosine similarity between projected embeddings.
            
            Args:
                source_proj: Source projections (batch, dim)
                target_proj: Target projections (batch, dim) or (num_targets, dim)
                
            Returns:
                Similarity scores
            """
            # Normalize
            source_norm = F.normalize(source_proj, p=2, dim=-1)
            target_norm = F.normalize(target_proj, p=2, dim=-1)
            
            # Cosine similarity
            if source_proj.shape[0] == target_proj.shape[0]:
                # Pairwise similarity
                sim = (source_norm * target_norm).sum(dim=-1)
            else:
                # All-pairs similarity
                sim = source_norm @ target_norm.T
            
            return sim / self.temperature
        
        def contrastive_loss(
            self,
            source_proj: torch.Tensor,
            target_proj: torch.Tensor,
            labels: torch.Tensor,  # 1 for match, 0 for non-match
        ) -> torch.Tensor:
            """
            Compute contrastive loss (simplified InfoNCE).
            
            For positive pairs: maximize similarity
            For negative pairs: minimize similarity
            """
            sim = self.compute_similarity(source_proj, target_proj)
            
            # Binary cross entropy with logits
            loss = F.binary_cross_entropy_with_logits(
                sim,
                labels.float(),
                reduction="mean",
            )
            
            return loss
    
    
    class TypeMatchDataset(Dataset):
        """Dataset for type-aware contrastive learning."""
        
        def __init__(
            self,
            pairs: list[TypeMatchPair],
            embedder: Any,  # Any embedding function
            vocab: TypeVocabulary,
        ):
            self.pairs = pairs
            self.embedder = embedder
            self.vocab = vocab
            
            # Pre-compute embeddings
            self._source_texts = [p.source_text for p in pairs]
            self._target_texts = [p.target_text for p in pairs]
            
            # Get embeddings (batched for efficiency)
            print(f"  Pre-computing embeddings for {len(pairs)} pairs...")
            self._source_embs = embedder(self._source_texts)
            self._target_embs = embedder(self._target_texts)
            print("  ✓ Embeddings computed")
        
        def __len__(self) -> int:
            return len(self.pairs)
        
        def __getitem__(self, idx: int) -> dict:
            pair = self.pairs[idx]
            
            return {
                "source_emb": torch.tensor(self._source_embs[idx], dtype=torch.float32),
                "source_type": self.vocab.get_type_id(pair.source_type),
                "target_emb": torch.tensor(self._target_embs[idx], dtype=torch.float32),
                "target_type": self.vocab.get_type_id(pair.target_type),
                "label": 1.0 if pair.is_match else 0.0,
            }


# =============================================================================
# TRAINING DATA GENERATOR
# =============================================================================

class TrainingDataGenerator:
    """
    Generates training pairs for contrastive learning.
    
    Creates positive pairs (matching fields) and negative pairs (non-matching).
    """
    
    # Sample field templates for data augmentation
    FIELD_TEMPLATES = {
        "string": [
            ("{name}", "{name}_text", "{name}_str"),
            ("customer_{name}", "cust_{name}", "{name}_customer"),
            ("user_{name}", "usr_{name}", "{name}_user"),
        ],
        "integer": [
            ("{name}_id", "{name}Id", "{name}_identifier"),
            ("{name}_count", "{name}_num", "num_{name}"),
            ("{name}_code", "{name}_no", "{name}_number"),
        ],
        "decimal": [
            ("{name}_amount", "{name}_value", "{name}_total"),
            ("{name}_price", "{name}_cost", "{name}_rate"),
            ("{name}_balance", "{name}_sum", "{name}_qty"),
        ],
        "date": [
            ("{name}_date", "{name}_dt", "date_{name}"),
            ("{name}_timestamp", "{name}_ts", "{name}_time"),
            ("created_{name}", "updated_{name}", "{name}_on"),
        ],
        "boolean": [
            ("is_{name}", "has_{name}", "{name}_flag"),
            ("{name}_indicator", "{name}_yn", "{name}_bool"),
        ],
    }
    
    # Concepts that commonly appear in schemas
    COMMON_CONCEPTS = [
        "customer", "user", "account", "order", "product", "payment",
        "address", "phone", "email", "name", "status", "created",
        "updated", "amount", "price", "quantity", "description",
        "category", "type", "code", "active", "enabled", "verified",
    ]
    
    def __init__(self, seed: int = 42):
        self.rng = random.Random(seed)
    
    def generate_pairs(
        self,
        num_positive: int = 1000,
        num_negative: int = 1000,
    ) -> list[TypeMatchPair]:
        """
        Generate training pairs.
        
        Args:
            num_positive: Number of positive (matching) pairs
            num_negative: Number of negative (non-matching) pairs
            
        Returns:
            List of TypeMatchPair objects
        """
        pairs = []
        
        # Generate positive pairs
        for _ in range(num_positive):
            pair = self._generate_positive_pair()
            if pair:
                pairs.append(pair)
        
        # Generate negative pairs
        for _ in range(num_negative):
            pair = self._generate_negative_pair()
            if pair:
                pairs.append(pair)
        
        # Shuffle
        self.rng.shuffle(pairs)
        
        return pairs
    
    def _generate_positive_pair(self) -> Optional[TypeMatchPair]:
        """Generate a positive (matching) pair."""
        # Pick a type
        type_category = self.rng.choice(list(self.FIELD_TEMPLATES.keys()))
        templates = self.FIELD_TEMPLATES[type_category]
        
        # Pick a concept
        concept = self.rng.choice(self.COMMON_CONCEPTS)
        
        # Pick two different templates
        template_group = self.rng.choice(templates)
        t1, t2 = self.rng.sample(template_group, 2)
        
        source_name = t1.format(name=concept)
        target_name = t2.format(name=concept)
        
        # Map type category to actual types
        type_mapping = {
            "string": ("string", "varchar"),
            "integer": ("integer", "int"),
            "decimal": ("decimal", "float"),
            "date": ("datetime", "timestamp"),
            "boolean": ("boolean", "bool"),
        }
        
        source_type, target_type = type_mapping.get(type_category, ("string", "string"))
        
        return TypeMatchPair(
            source_text=source_name,
            source_type=source_type,
            target_text=target_name,
            target_type=target_type,
            is_match=True,
        )
    
    def _generate_negative_pair(self) -> Optional[TypeMatchPair]:
        """Generate a negative (non-matching) pair."""
        # Pick two different concepts
        concepts = self.rng.sample(self.COMMON_CONCEPTS, 2)
        
        # Pick two different types
        types = self.rng.sample(list(self.FIELD_TEMPLATES.keys()), 2)
        
        # Generate names
        t1 = self.rng.choice(self.FIELD_TEMPLATES[types[0]])[0]
        t2 = self.rng.choice(self.FIELD_TEMPLATES[types[1]])[0]
        
        source_name = t1.format(name=concepts[0])
        target_name = t2.format(name=concepts[1])
        
        type_mapping = {
            "string": "string",
            "integer": "integer",
            "decimal": "decimal",
            "date": "datetime",
            "boolean": "boolean",
        }
        
        return TypeMatchPair(
            source_text=source_name,
            source_type=type_mapping[types[0]],
            target_text=target_name,
            target_type=type_mapping[types[1]],
            is_match=False,
        )


# =============================================================================
# TYPE PROJECTION MANAGER
# =============================================================================

class TypeProjectionManager:
    """
    High-level manager for type projection model.
    
    Handles training, saving, loading, and inference.
    """
    
    def __init__(
        self,
        config: Optional[TypeProjectionConfig] = None,
        device: str = "cpu",
    ):
        if not TORCH_AVAILABLE:
            raise ImportError(
                "PyTorch is required for type projections. "
                "Install with: pip install torch"
            )
        
        self.config = config or TypeProjectionConfig()
        self.device = torch.device(device)
        self.vocab = TypeVocabulary()
        
        # Initialize model
        self.model = ContrastiveTypeModel(self.config)
        self.model.to(self.device)
        
        self._is_trained = False
    
    def train(
        self,
        pairs: list[TypeMatchPair],
        embedder: Any,
        verbose: bool = True,
    ) -> dict:
        """
        Train the type projection model.
        
        Args:
            pairs: Training pairs
            embedder: Function to compute base embeddings
            verbose: Print training progress
            
        Returns:
            Training metrics
        """
        # Create dataset
        dataset = TypeMatchDataset(pairs, embedder, self.vocab)
        dataloader = DataLoader(
            dataset,
            batch_size=self.config.batch_size,
            shuffle=True,
        )
        
        # Optimizer
        optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=self.config.learning_rate,
        )
        
        # Training loop
        self.model.train()
        metrics = {"losses": [], "accuracies": []}
        
        for epoch in range(self.config.num_epochs):
            epoch_loss = 0.0
            epoch_correct = 0
            epoch_total = 0
            
            for batch in dataloader:
                # Move to device
                source_emb = batch["source_emb"].to(self.device)
                source_type = batch["source_type"].to(self.device)
                target_emb = batch["target_emb"].to(self.device)
                target_type = batch["target_type"].to(self.device)
                labels = batch["label"].to(self.device)
                
                # Forward pass
                optimizer.zero_grad()
                source_proj, target_proj = self.model(
                    source_emb, source_type, target_emb, target_type
                )
                
                # Compute loss
                loss = self.model.contrastive_loss(source_proj, target_proj, labels)
                
                # Backward pass
                loss.backward()
                optimizer.step()
                
                # Track metrics
                epoch_loss += loss.item() * len(labels)
                
                # Accuracy
                with torch.no_grad():
                    sim = self.model.compute_similarity(source_proj, target_proj)
                    preds = (torch.sigmoid(sim) > 0.5).float()
                    epoch_correct += (preds == labels).sum().item()
                    epoch_total += len(labels)
            
            avg_loss = epoch_loss / epoch_total
            accuracy = epoch_correct / epoch_total
            
            metrics["losses"].append(avg_loss)
            metrics["accuracies"].append(accuracy)
            
            if verbose:
                print(f"  Epoch {epoch + 1}/{self.config.num_epochs}: "
                      f"Loss={avg_loss:.4f}, Accuracy={accuracy:.4f}")
        
        self._is_trained = True
        return metrics
    
    def project(
        self,
        base_embedding: np.ndarray,
        data_type: str,
    ) -> np.ndarray:
        """
        Project a base embedding with type information.
        
        Args:
            base_embedding: Base semantic embedding
            data_type: String data type name
            
        Returns:
            Type-aware projected embedding
        """
        self.model.eval()
        
        with torch.no_grad():
            emb = torch.tensor(base_embedding, dtype=torch.float32).unsqueeze(0)
            type_id = torch.tensor([self.vocab.get_type_id(data_type)])
            
            emb = emb.to(self.device)
            type_id = type_id.to(self.device)
            
            projected = self.model.projection(emb, type_id)
            
        return projected.cpu().numpy().squeeze()
    
    def project_batch(
        self,
        base_embeddings: np.ndarray,
        data_types: list[str],
    ) -> np.ndarray:
        """
        Project a batch of embeddings with type information.
        
        Args:
            base_embeddings: Base semantic embeddings (batch, dim)
            data_types: List of type names
            
        Returns:
            Type-aware projected embeddings (batch, dim)
        """
        self.model.eval()
        
        with torch.no_grad():
            embs = torch.tensor(base_embeddings, dtype=torch.float32)
            type_ids = torch.tensor([self.vocab.get_type_id(t) for t in data_types])
            
            embs = embs.to(self.device)
            type_ids = type_ids.to(self.device)
            
            projected = self.model.projection(embs, type_ids)
            
        return projected.cpu().numpy()
    
    def save(self, path: Union[str, Path]) -> None:
        """Save model to file."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        
        torch.save({
            "model_state": self.model.state_dict(),
            "config": self.config,
            "vocab": {
                "type_to_id": self.vocab._type_to_id,
                "next_id": self.vocab._next_id,
            },
            "is_trained": self._is_trained,
        }, path)
    
    def load(self, path: Union[str, Path]) -> None:
        """Load model from file."""
        path = Path(path)
        
        checkpoint = torch.load(path, map_location=self.device)
        
        self.config = checkpoint["config"]
        self.model = ContrastiveTypeModel(self.config)
        self.model.load_state_dict(checkpoint["model_state"])
        self.model.to(self.device)
        
        self.vocab._type_to_id = checkpoint["vocab"]["type_to_id"]
        self.vocab._next_id = checkpoint["vocab"]["next_id"]
        
        self._is_trained = checkpoint["is_trained"]
    
    @property
    def is_trained(self) -> bool:
        """Check if model has been trained."""
        return self._is_trained


# =============================================================================
# EVALUATION
# =============================================================================

def evaluate_type_projection(
    manager: TypeProjectionManager,
    test_pairs: list[TypeMatchPair],
    embedder: Any,
) -> dict:
    """
    Evaluate type projection model.
    
    Args:
        manager: Trained TypeProjectionManager
        test_pairs: Test pairs
        embedder: Function to compute base embeddings
        
    Returns:
        Evaluation metrics
    """
    if not TORCH_AVAILABLE:
        raise ImportError("PyTorch required")
    
    # Get base embeddings
    source_texts = [p.source_text for p in test_pairs]
    target_texts = [p.target_text for p in test_pairs]
    
    source_base = embedder(source_texts)
    target_base = embedder(target_texts)
    
    # Project with types
    source_types = [p.source_type for p in test_pairs]
    target_types = [p.target_type for p in test_pairs]
    
    source_proj = manager.project_batch(source_base, source_types)
    target_proj = manager.project_batch(target_base, target_types)
    
    # Compute similarities
    source_norm = source_proj / np.linalg.norm(source_proj, axis=-1, keepdims=True)
    target_norm = target_proj / np.linalg.norm(target_proj, axis=-1, keepdims=True)
    
    similarities = np.sum(source_norm * target_norm, axis=-1)
    
    # Compute metrics
    labels = np.array([p.is_match for p in test_pairs])
    preds = similarities > 0.5
    
    accuracy = np.mean(preds == labels)
    
    # MRR calculation
    # For each positive pair, rank all targets
    positive_indices = np.where(labels)[0]
    mrr_scores = []
    
    for idx in positive_indices[:100]:  # Limit for efficiency
        source_vec = source_proj[idx]
        all_sims = source_vec @ target_proj.T
        
        # Rank
        sorted_indices = np.argsort(-all_sims)
        rank = np.where(sorted_indices == idx)[0][0] + 1
        mrr_scores.append(1.0 / rank)
    
    mrr = np.mean(mrr_scores) if mrr_scores else 0.0
    
    return {
        "accuracy": float(accuracy),
        "mrr": float(mrr),
        "avg_positive_sim": float(np.mean(similarities[labels])),
        "avg_negative_sim": float(np.mean(similarities[~labels])),
        "separation": float(np.mean(similarities[labels]) - np.mean(similarities[~labels])),
    }


# =============================================================================
# MAIN (for testing)
# =============================================================================

if __name__ == "__main__":
    print("GAP-008: Learned Type Projections")
    print("=" * 50)
    
    if not TORCH_AVAILABLE:
        print("ERROR: PyTorch not installed")
        print("Install with: pip install torch")
        exit(1)
    
    # Test type vocabulary
    print("\n1. Testing Type Vocabulary...")
    vocab = TypeVocabulary()
    
    test_types = ["string", "varchar", "integer", "datetime", "unknown_type"]
    for t in test_types:
        tid = vocab.get_type_id(t)
        print(f"   {t} -> ID {tid}")
    
    print(f"   Vocab size: {vocab.vocab_size}")
    
    # Test training data generation
    print("\n2. Generating Training Data...")
    generator = TrainingDataGenerator()
    pairs = generator.generate_pairs(num_positive=100, num_negative=100)
    print(f"   Generated {len(pairs)} pairs")
    print(f"   Positive: {sum(1 for p in pairs if p.is_match)}")
    print(f"   Negative: {sum(1 for p in pairs if not p.is_match)}")
    
    # Show sample pairs
    print("\n   Sample positive pair:")
    pos_pair = next(p for p in pairs if p.is_match)
    print(f"   Source: {pos_pair.source_text} ({pos_pair.source_type})")
    print(f"   Target: {pos_pair.target_text} ({pos_pair.target_type})")
    
    print("\n   Sample negative pair:")
    neg_pair = next(p for p in pairs if not p.is_match)
    print(f"   Source: {neg_pair.source_text} ({neg_pair.source_type})")
    print(f"   Target: {neg_pair.target_text} ({neg_pair.target_type})")
    
    # Test model creation
    print("\n3. Testing Model Creation...")
    config = TypeProjectionConfig(
        base_embedding_dim=384,
        type_embedding_dim=64,
        projection_dim=384,
    )
    model = ContrastiveTypeModel(config)
    
    # Test forward pass
    batch_size = 4
    source_emb = torch.randn(batch_size, config.base_embedding_dim)
    source_types = torch.randint(0, 25, (batch_size,))
    target_emb = torch.randn(batch_size, config.base_embedding_dim)
    target_types = torch.randint(0, 25, (batch_size,))
    
    source_proj, target_proj = model(source_emb, source_types, target_emb, target_types)
    print(f"   Input shape: {source_emb.shape}")
    print(f"   Output shape: {source_proj.shape}")
    
    # Test similarity
    sim = model.compute_similarity(source_proj, target_proj)
    print(f"   Similarity shape: {sim.shape}")
    
    print("\n✓ Type projection model working!")
    print("\nTo train with real embeddings, use:")
    print("  from sentence_transformers import SentenceTransformer")
    print("  embedder = SentenceTransformer('all-MiniLM-L6-v2')")
    print("  manager = TypeProjectionManager()")
    print("  manager.train(pairs, embedder.encode)")
