"""RV Telos evaluator - measures embedding contraction during self-referential processing.

This evaluator integrates with the RV ground truth measurement from telos-os/research/.

Reference: ../telos-os/research/rv_ground_truth_results.md
"""

import hashlib
from typing import Dict, Any, Optional
import warnings

from .base import BaseEvaluator


class RVTelosEvaluator(BaseEvaluator):
    """
    Measures R_V (embedding contraction ratio) for self-referential vs neutral content.
    
    R_V = (mean pairwise similarity in self-ref) / (mean pairwise similarity in neutral)
    
    Values:
    - R_V < 1.0: Contraction (self-ref embeddings are more similar)
    - R_V = 1.0: No difference
    - R_V > 1.0: Expansion (self-ref embeddings are less similar)
    
    Also computes TTR (type-token ratio) as behavioral proxy.
    """
    
    def __init__(self, config):
        super().__init__(config)
        
        # Check if required dependencies are available
        self.transformers_available = self._check_transformers()
        self.torch_available = self._check_torch()
        self.cuda_available = self._check_cuda() if config.cuda_enabled else False
        
        if not self.transformers_available or not self.torch_available:
            warnings.warn(
                "RV Telos evaluator requires transformers and torch. "
                "Install with: pip install torch transformers"
            )
            self.model = None
            self.tokenizer = None
        else:
            self.model, self.tokenizer = self._load_model()
    
    def _check_transformers(self) -> bool:
        try:
            import transformers
            return True
        except ImportError:
            return False
    
    def _check_torch(self) -> bool:
        try:
            import torch
            return True
        except ImportError:
            return False
    
    def _check_cuda(self) -> bool:
        try:
            import torch
            return torch.cuda.is_available()
        except:
            return False
    
    def _load_model(self):
        """Load transformer model for embedding extraction."""
        import torch
        from transformers import AutoModel, AutoTokenizer
        
        # Use first model from config
        model_cfg = self.config.models[0]
        model_name = model_cfg.model
        
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModel.from_pretrained(model_name)
        
        if self.cuda_available:
            model = model.cuda()
        
        model.eval()
        
        return model, tokenizer
    
    def evaluate(self, item: Dict[str, Any]) -> Dict[str, Any]:
        """
        Evaluate item by measuring R_V and TTR.
        
        For a single item, we can't compute full R_V (needs pairs).
        Instead, we compute:
        - TTR for this item
        - Self-reference density
        - If model available: embedding statistics
        
        Full R_V computation happens at aggregate level.
        """
        content = item.get("content", "")
        content_hash = hashlib.sha256(content.encode()).hexdigest()[:8]
        category = item.get("category", "unknown")
        
        result = {
            "item_id": item.get("id", "unknown"),
            "content_hash": content_hash,
            "category": category,
        }
        
        # Compute TTR (type-token ratio)
        ttr = self._compute_ttr(content)
        result["ttr"] = ttr
        
        # Compute self-reference density
        self_ref_density = self._compute_self_ref_density(content)
        result["self_ref_density"] = self_ref_density
        
        # If model available, extract embeddings
        if self.model is not None and self.tokenizer is not None:
            try:
                embedding_stats = self._extract_embedding_stats(content)
                result.update(embedding_stats)
            except Exception as e:
                result["embedding_error"] = str(e)
        else:
            result["warning"] = "Model not available - using behavioral proxies only"
        
        return result
    
    def _compute_ttr(self, text: str) -> float:
        """
        Compute type-token ratio (unique words / total words).
        
        Lower TTR = more repetitive = potential contraction signal.
        """
        words = text.lower().split()
        if not words:
            return 0.0
        
        unique_words = set(words)
        return len(unique_words) / len(words)
    
    def _compute_self_ref_density(self, text: str) -> float:
        """
        Compute fraction of self-referential terms.
        
        Self-ref terms: I, my, me, myself, reasoning, thinking, analyzing, etc.
        """
        self_ref_terms = {
            "i", "my", "me", "myself", "mine",
            "reasoning", "thinking", "analyzing", "reflecting",
            "internal", "model", "process", "cognitive",
        }
        
        words = text.lower().split()
        if not words:
            return 0.0
        
        self_ref_count = sum(1 for w in words if w in self_ref_terms)
        return self_ref_count / len(words)
    
    def _extract_embedding_stats(self, text: str) -> Dict[str, Any]:
        """
        Extract embedding statistics for this text.
        
        Returns mean/std of embedding norms across layers.
        
        NOTE: Full R_V computation requires pairs of embeddings
        and is done at the aggregate level in harness.compute_metrics().
        """
        import torch
        
        # Tokenize
        inputs = self.tokenizer(
            text, 
            return_tensors="pt", 
            truncation=True, 
            max_length=512
        )
        
        if self.cuda_available:
            inputs = {k: v.cuda() for k, v in inputs.items()}
        
        # Extract embeddings
        with torch.no_grad():
            outputs = self.model(**inputs, output_hidden_states=True)
            hidden_states = outputs.hidden_states  # Tuple of (batch, seq, hidden)
        
        # Compute stats across layers
        layer_norms = []
        for layer_idx, layer_hidden in enumerate(hidden_states):
            # Mean pool across sequence
            pooled = layer_hidden.mean(dim=1)  # (batch, hidden)
            norm = torch.norm(pooled, dim=1).item()
            layer_norms.append(norm)
        
        return {
            "embedding_mean_norm": float(sum(layer_norms) / len(layer_norms)),
            "embedding_std_norm": float(torch.std(torch.tensor(layer_norms)).item()),
            "num_layers": len(layer_norms),
        }


# TODO: Implement full RV computation at aggregate level
# This should be done in harness._compute_rv_metrics()
# by pairing self-ref and neutral embeddings and computing
# pairwise cosine similarities within each category.
#
# See ../telos-os/research/rv_ground_truth_results.md for reference implementation.
