#!/usr/bin/env python3
"""
helix_cdc/ginnie/helix_embedding.py

Integration of trained HelixMLP with Echo Ginnie for domain-aware embeddings.

WO-ECHO-FINETUNE-01

This module provides:
    - HelixEmbeddingBackend: Generate domain-aware embeddings using trained HelixMLP
    - HelixRouter: Route queries to appropriate backends based on embedding similarity
    - Domain-specific embedding utilities

Usage:
    from helix_cdc.ginnie.helix_embedding import HelixEmbeddingBackend

    backend = HelixEmbeddingBackend.from_checkpoint('models/echo_helix_v1')
    embedding = backend.embed("What is Se?")
    similarity = backend.similarity(embedding1, embedding2)
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn.functional as F

# Import HelixMLP
from helix_cdc.flowtorch_helix_ops import HelixMLP, compute_se_torch


class HelixLanguageModel(torch.nn.Module):
    """
    Language model wrapper using HelixMLP.
    Must match the architecture from train_echo_helix.py
    """

    def __init__(
        self,
        vocab_size: int = 32000,
        embed_dim: int = 256,
        hidden_dim: int = 512,
        n_helix_layers: int = 4,
        codebook_size: int = 256,
        pattern_dim: int = 64,
        n_mix: int = 32,
    ):
        super().__init__()

        self.vocab_size = vocab_size
        self.embed_dim = embed_dim

        # Token embedding
        self.embedding = torch.nn.Embedding(vocab_size, embed_dim)

        # HelixMLP core
        layer_dims = [embed_dim] + [hidden_dim] * (n_helix_layers - 1) + [embed_dim]
        self.helix_mlp = HelixMLP(
            layer_dims=layer_dims,
            codebook_size=codebook_size,
            pattern_dim=pattern_dim,
            n_mix=n_mix,
            activation='gelu',
        )

        # Output projection
        self.output_proj = torch.nn.Linear(embed_dim, vocab_size)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(input_ids)
        batch_size, seq_len, embed_dim = x.shape
        x = x.view(batch_size * seq_len, embed_dim)
        x = self.helix_mlp(x)
        x = x.view(batch_size, seq_len, embed_dim)
        logits = self.output_proj(x)
        return logits

    def get_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Get embeddings before output projection (for similarity)."""
        x = self.embedding(input_ids)
        batch_size, seq_len, embed_dim = x.shape
        x = x.view(batch_size * seq_len, embed_dim)
        x = self.helix_mlp(x)
        x = x.view(batch_size, seq_len, embed_dim)
        return x


class SimpleTokenizer:
    """Simple character-level tokenizer for embedding generation."""

    def __init__(self, vocab_size: int = 32000):
        self.vocab_size = vocab_size
        self.pad_token_id = 0
        self.eos_token_id = 1
        self.unk_token_id = 2

        self._char_to_id = {'<pad>': 0, '<eos>': 1, '<unk>': 2}
        self._id_to_char = {0: '<pad>', 1: '<eos>', 2: '<unk>'}

        for i, c in enumerate(range(32, 127)):
            self._char_to_id[chr(c)] = i + 3
            self._id_to_char[i + 3] = chr(c)

    def encode(
        self,
        text: str,
        max_length: int = 512,
        return_tensors: bool = True,
    ) -> Union[List[int], torch.Tensor]:
        """Encode text to token IDs."""
        ids = []
        for c in text:
            if c in self._char_to_id:
                ids.append(self._char_to_id[c])
            else:
                ids.append(self.unk_token_id)
        ids.append(self.eos_token_id)

        if len(ids) > max_length:
            ids = ids[:max_length]

        if return_tensors:
            return torch.tensor([ids], dtype=torch.long)
        return ids

    def decode(self, ids: Union[List[int], torch.Tensor]) -> str:
        """Decode token IDs to text."""
        if isinstance(ids, torch.Tensor):
            ids = ids.flatten().tolist()

        chars = []
        for id in ids:
            if id in self._id_to_char:
                char = self._id_to_char[id]
                if char not in ['<pad>', '<eos>', '<unk>']:
                    chars.append(char)
        return ''.join(chars)


class HelixEmbeddingBackend:
    """
    Generate domain-aware embeddings using trained HelixMLP.

    This backend uses a trained HelixLanguageModel to generate
    embeddings that are aware of helix-cdc domain vocabulary.
    """

    def __init__(
        self,
        model: HelixLanguageModel,
        tokenizer: SimpleTokenizer,
        config: Dict[str, Any],
        device: str = 'cpu',
    ):
        self.model = model.to(device)
        self.model.eval()
        self.tokenizer = tokenizer
        self.config = config
        self.device = device
        self.embed_dim = config.get('embed_dim', 256)

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_dir: str,
        device: str = 'cpu',
    ) -> 'HelixEmbeddingBackend':
        """
        Load embedding backend from trained checkpoint.

        Args:
            checkpoint_dir: Path to model directory (from train_echo_helix.py)
            device: Device to load model on

        Returns:
            HelixEmbeddingBackend instance
        """
        checkpoint_dir = Path(checkpoint_dir)

        # Load config
        config_path = checkpoint_dir / 'config.json'
        if not config_path.exists():
            raise FileNotFoundError(f"Config not found: {config_path}")

        with open(config_path, 'r') as f:
            config = json.load(f)

        # Create model
        model = HelixLanguageModel(
            vocab_size=config.get('vocab_size', 32000),
            embed_dim=config.get('embed_dim', 256),
            hidden_dim=config.get('hidden_dim', 512),
            n_helix_layers=config.get('n_helix_layers', 4),
            codebook_size=config.get('codebook_size', 256),
            pattern_dim=config.get('pattern_dim', 64),
            n_mix=config.get('n_mix', 32),
        )

        # Load weights
        model_path = checkpoint_dir / 'model.pt'
        if not model_path.exists():
            raise FileNotFoundError(f"Model not found: {model_path}")

        state_dict = torch.load(model_path, map_location=device)
        model.load_state_dict(state_dict)

        # Create tokenizer
        tokenizer = SimpleTokenizer(vocab_size=config.get('vocab_size', 32000))

        return cls(model, tokenizer, config, device)

    def embed(
        self,
        text: str,
        pooling: str = 'mean',
        normalize: bool = True,
    ) -> torch.Tensor:
        """
        Generate embedding for text.

        Args:
            text: Input text
            pooling: Pooling method ('mean', 'last', 'first')
            normalize: Whether to L2-normalize the output

        Returns:
            Embedding tensor of shape (embed_dim,)
        """
        # Tokenize
        input_ids = self.tokenizer.encode(text).to(self.device)

        # Get embeddings
        with torch.no_grad():
            embeddings = self.model.get_embeddings(input_ids)  # (1, seq_len, embed_dim)

        # Pool
        if pooling == 'mean':
            embedding = embeddings.mean(dim=1).squeeze(0)  # (embed_dim,)
        elif pooling == 'last':
            embedding = embeddings[:, -1, :].squeeze(0)
        elif pooling == 'first':
            embedding = embeddings[:, 0, :].squeeze(0)
        else:
            raise ValueError(f"Unknown pooling: {pooling}")

        # Normalize
        if normalize:
            embedding = F.normalize(embedding, p=2, dim=-1)

        return embedding

    def embed_batch(
        self,
        texts: List[str],
        pooling: str = 'mean',
        normalize: bool = True,
    ) -> torch.Tensor:
        """
        Generate embeddings for a batch of texts.

        Args:
            texts: List of input texts
            pooling: Pooling method
            normalize: Whether to normalize

        Returns:
            Embeddings tensor of shape (batch_size, embed_dim)
        """
        embeddings = []
        for text in texts:
            emb = self.embed(text, pooling=pooling, normalize=normalize)
            embeddings.append(emb)
        return torch.stack(embeddings)

    def similarity(
        self,
        embedding1: torch.Tensor,
        embedding2: torch.Tensor,
    ) -> float:
        """Compute cosine similarity between two embeddings."""
        e1 = F.normalize(embedding1, p=2, dim=-1)
        e2 = F.normalize(embedding2, p=2, dim=-1)
        return float(torch.dot(e1, e2))

    def similarity_matrix(self, embeddings: torch.Tensor) -> torch.Tensor:
        """Compute pairwise similarity matrix."""
        normed = F.normalize(embeddings, p=2, dim=-1)
        return torch.mm(normed, normed.t())

    def get_codebook_se(self) -> float:
        """Get Se of the model's codebook."""
        return compute_se_torch(self.model.helix_mlp.codebook.patterns)

    def is_healthy(self) -> bool:
        """Check if model weights are in healthy Se band."""
        se = self.get_codebook_se()
        return 0.3 <= se <= 0.7


class HelixRouter:
    """
    Route queries to backends based on embedding similarity.

    Uses trained HelixEmbeddingBackend to compute query embeddings
    and route to the most appropriate backend based on domain anchors.
    """

    def __init__(
        self,
        embedding_backend: HelixEmbeddingBackend,
        anchors: Optional[Dict[str, str]] = None,
    ):
        """
        Initialize router.

        Args:
            embedding_backend: HelixEmbeddingBackend for embeddings
            anchors: Dict mapping backend names to anchor texts
        """
        self.backend = embedding_backend

        # Default domain anchors
        if anchors is None:
            anchors = {
                'helix': "HelixCode compression codebook seed generation",
                'se': "Se symbolic entropy measurement log1p H C D formula",
                'fibpi': "FibPi core cycling CPU distribution fibonacci pi",
                'cdna': "CDNA compressed domain neural architecture tensor shard",
                'vault': "chunk vault content addressable storage SHA256 regeneration",
                'flowtorch': "FlowTorch routing GPU CPU QPU braid DLPack",
                'general': "general programming code documentation",
            }

        self.anchors = anchors
        self.anchor_embeddings = {}

        # Precompute anchor embeddings
        self._compute_anchor_embeddings()

    def _compute_anchor_embeddings(self):
        """Compute and cache anchor embeddings."""
        for name, text in self.anchors.items():
            self.anchor_embeddings[name] = self.backend.embed(text)

    def route(self, query: str, threshold: float = 0.5) -> Tuple[str, float]:
        """
        Route a query to the best matching backend.

        Args:
            query: Input query text
            threshold: Minimum similarity threshold

        Returns:
            (backend_name, similarity_score)
        """
        query_embedding = self.backend.embed(query)

        best_backend = 'general'
        best_similarity = 0.0

        for name, anchor_emb in self.anchor_embeddings.items():
            similarity = self.backend.similarity(query_embedding, anchor_emb)
            if similarity > best_similarity:
                best_similarity = similarity
                best_backend = name

        if best_similarity < threshold:
            best_backend = 'general'

        return best_backend, best_similarity

    def route_batch(
        self,
        queries: List[str],
        threshold: float = 0.5,
    ) -> List[Tuple[str, float]]:
        """Route a batch of queries."""
        return [self.route(q, threshold) for q in queries]

    def get_all_similarities(self, query: str) -> Dict[str, float]:
        """Get similarity scores for all backends."""
        query_embedding = self.backend.embed(query)

        similarities = {}
        for name, anchor_emb in self.anchor_embeddings.items():
            similarities[name] = self.backend.similarity(query_embedding, anchor_emb)

        return similarities


def create_helix_router(model_dir: str, device: str = 'cpu') -> HelixRouter:
    """
    Convenience function to create a HelixRouter from a model directory.

    Args:
        model_dir: Path to trained model directory
        device: Device to use

    Returns:
        Configured HelixRouter
    """
    backend = HelixEmbeddingBackend.from_checkpoint(model_dir, device)
    return HelixRouter(backend)


if __name__ == '__main__':
    # Self-test
    import sys

    print("HelixEmbeddingBackend self-test")
    print("=" * 50)

    # Check if a model exists
    default_model = 'models/echo_helix_v1'
    if not os.path.exists(default_model):
        print(f"Note: No trained model at {default_model}")
        print("Run train_echo_helix.py first to create a model")
        print()
        print("Creating minimal test model...")

        # Create a minimal untrained model for testing
        config = {
            'vocab_size': 1000,
            'embed_dim': 64,
            'hidden_dim': 128,
            'n_helix_layers': 2,
            'codebook_size': 64,
            'pattern_dim': 32,
            'n_mix': 16,
        }

        model = HelixLanguageModel(**config)
        tokenizer = SimpleTokenizer(vocab_size=config['vocab_size'])

        backend = HelixEmbeddingBackend(model, tokenizer, config, device='cpu')
    else:
        print(f"Loading model from: {default_model}")
        backend = HelixEmbeddingBackend.from_checkpoint(default_model)

    # Test embedding
    test_texts = [
        "What is Se in the Helix system?",
        "How does FibPi core cycling work?",
        "General programming question",
    ]

    print()
    print("Testing embeddings...")
    for text in test_texts:
        emb = backend.embed(text)
        print(f"  '{text[:40]}...' -> shape={emb.shape}, norm={emb.norm():.4f}")

    # Test similarity
    print()
    print("Testing similarity...")
    emb1 = backend.embed(test_texts[0])
    emb2 = backend.embed(test_texts[1])
    emb3 = backend.embed(test_texts[2])

    print(f"  Helix vs FibPi: {backend.similarity(emb1, emb2):.4f}")
    print(f"  Helix vs General: {backend.similarity(emb1, emb3):.4f}")
    print(f"  FibPi vs General: {backend.similarity(emb2, emb3):.4f}")

    # Test router
    print()
    print("Testing router...")
    router = HelixRouter(backend)

    for text in test_texts:
        route, score = router.route(text)
        print(f"  '{text[:40]}...' -> {route} (score={score:.4f})")

    # Health check
    print()
    print("Health check...")
    se = backend.get_codebook_se()
    healthy = backend.is_healthy()
    print(f"  Codebook Se: {se:.4f}")
    print(f"  In healthy band (0.3-0.7): {healthy}")

    print()
    print("Self-test complete!")
