"""
WRAG 2.0 - CLEAN APPROACH
Use frozen base model for features, add weight retrieval layers on top
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
import os
import sys

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from layers.weight_retrieval_layer import WeightRetrievalLayer


class WRAG2TextModel(nn.Module):
    """
    WRAG 2.0 for text classification
    
    Architecture:
    1. Frozen base model (TinyLlama) extracts features
    2. Stack of weight retrieval layers (trainable)
    3. Classification head
    """
    def __init__(self, base_model_name, num_shards=10, k=3, num_classes=2, num_wr_layers=3):
        super().__init__()
        
        print(f"Loading base model: {base_model_name}")
        self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
        self.base_model = AutoModel.from_pretrained(base_model_name, torch_dtype=torch.float32)
        
        # Freeze base model
        for param in self.base_model.parameters():
            param.requires_grad = False
        
        hidden_size = self.base_model.config.hidden_size
        
        print(f"Adding {num_wr_layers} weight retrieval layers")
        print(f"Using {num_shards} shards, k={k}")
        
        # Stack of weight retrieval layers
        self.wr_layers = nn.ModuleList([
            WeightRetrievalLayer(hidden_size, hidden_size, num_shards, k)
            for _ in range(num_wr_layers)
        ])
        
        # Classification head
        self.classifier = nn.Linear(hidden_size, num_classes)
        
        self.num_shards = num_shards
        self.k = k
        
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        total = sum(p.numel() for p in self.parameters())
        print(f"Trainable: {trainable:,} / {total:,} ({trainable/total*100:.1f}%)")
    
    def forward(self, input_ids, attention_mask):
        # Get features from frozen base model
        with torch.no_grad():
            outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
            hidden_states = outputs.last_hidden_state  # [batch, seq, hidden]
        
        # Mean pool to get sequence representation
        pooled = (hidden_states * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.sum(1, keepdim=True)
        pooled = pooled.float()  # Ensure fp32 for training
        
        # Pass through weight retrieval layers
        all_scores = []
        x = pooled
        for wr_layer in self.wr_layers:
            x, scores = wr_layer(x)
            x = F.relu(x)
            all_scores.append(scores)
        
        # Classification
        logits = self.classifier(x)
        
        return logits, all_scores


if __name__ == '__main__':
    print("Testing WRAG 2.0 Text Model")
    model = WRAG2TextModel(
        "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        num_shards=10,
        k=3,
        num_wr_layers=3
    )
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = model.to(device)
    
    # Test
    test_text = ["The capital of France is Paris", "Machine learning is fascinating"]
    inputs = model.tokenizer(test_text, padding=True, truncation=True, return_tensors='pt').to(device)
    
    logits, scores = model(**inputs)
    print(f"✓ Logits shape: {logits.shape}")
    print(f"✓ Scores from {len(scores)} layers")