"""UCE Gemma: normalized embeddings or plain text to 80,513 labeled concepts."""
import json
from pathlib import Path

import numpy as np
import torch
from safetensors.torch import load_file
from scipy import sparse

if __package__:
    from .concept_metadata import ConceptMetadata, METADATA_FILES
else:
    from concept_metadata import ConceptMetadata, METADATA_FILES


def _directory(path, revision=None, cache_dir=None, local_files_only=False, *, with_embedder=False):
    directory = Path(path).expanduser()
    if not directory.is_dir():
        from huggingface_hub import snapshot_download
        directory = Path(snapshot_download(
            repo_id=str(path), revision=revision, cache_dir=cache_dir,
            local_files_only=local_files_only,
            allow_patterns=METADATA_FILES + ['encoder.safetensors']
                           + (['gemma_embedder/*'] if with_embedder else []),
        ))
    return directory


def _batch_size(value):
    if not isinstance(value, int) or isinstance(value, bool) or value < 1:
        raise ValueError('batch_size must be a positive integer')


class UCEGemmaEncoder(torch.nn.Module):
    """Gemma embedding-to-concept encoder; float32 sigmoid scores, top 128."""
    def __init__(self, config, tensors, labels, *, prevalences=None, metadata=None):
        super().__init__()
        self.config = dict(config)
        if config.get('model_type') != 'uce_gemma_80k' or config.get('format_version') != 1:
            raise ValueError('Unsupported UCE Gemma package')
        self.input_dim = config['input_dim']
        self.num_features = config['num_features']
        self.top_k = config['top_k']
        if config.get('input_unit_norm') is not True or not 0 < self.top_k <= self.num_features:
            raise ValueError('Invalid encoder configuration')
        if set(tensors) != {'weight', 'bias'}:
            raise ValueError('Expected weight and bias only')
        if tensors['weight'].shape != (self.num_features, self.input_dim) or tensors['bias'].shape != (self.num_features,):
            raise ValueError('Invalid encoder weight shape')
        for name, tensor in tensors.items():
            if tensor.dtype != torch.float32 or not torch.isfinite(tensor).all():
                raise ValueError('Expected finite float32 weights')
            self.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False))
        self.labels = dict(labels)
        self.metadata = metadata or ConceptMetadata(config, labels, prevalences=prevalences)
        self._prevalences = self.metadata._prevalences
        if set(self.labels) != set(range(self.num_features)) or any(not label.strip() for label in self.labels.values()):
            raise ValueError('Every output coordinate must have one nonempty label')
        self.eval()

    @classmethod
    def from_pretrained(cls, model_id_or_path, *, revision=None, device='cpu', cache_dir=None, local_files_only=False):
        directory = _directory(model_id_or_path, revision, cache_dir, local_files_only)
        config = json.loads((directory / 'config.json').read_text())
        metadata = ConceptMetadata.from_pretrained(directory)
        if not Path(model_id_or_path).expanduser().is_dir():
            metadata.provenance.update(
                source_repo_id=str(model_id_or_path), requested_revision=revision,
                resolved_revision=directory.name if directory.parent.name == 'snapshots' else None)
        return cls(config, load_file(str(directory / 'encoder.safetensors')), metadata.labels,
                   metadata=metadata).to(device)

    def get_prevalence(self, concept_id):
        """Return saved historical Gemini baseline counts, fractions, and bounds."""
        return self.metadata.get_prevalence(concept_id)

    @torch.inference_mode()
    def forward(self, embeddings):
        """Return device tensors (concept IDs, scores), each [batch, 128]."""
        if self.weight.dtype != torch.float32 or self.bias.dtype != torch.float32:
            raise ValueError('Keep the concept encoder in float32')
        if isinstance(embeddings, torch.Tensor):
            if embeddings.is_complex() or embeddings.dtype == torch.bool:
                raise ValueError('Expected real embeddings')
            x = embeddings.detach().to(device=self.weight.device, dtype=torch.float32)
        else:
            a = np.asarray(embeddings)
            if a.dtype.kind not in 'fiu':
                raise ValueError('Expected real embeddings')
            x = torch.from_numpy(np.array(a, dtype=np.float32, copy=True)).to(self.weight.device)
        if x.ndim != 2 or x.shape[1] != self.input_dim or not torch.isfinite(x).all():
            raise ValueError(f'Expected finite embeddings of shape (N, {self.input_dim})')
        if not torch.allclose(torch.linalg.vector_norm(x, dim=1), torch.ones(len(x), device=x.device), atol=2e-5, rtol=0):
            raise ValueError('Provide unit-normalized Gemma document embeddings')
        with torch.autocast(device_type=self.weight.device.type, enabled=False):
            logits = torch.nn.functional.linear(x, self.weight, self.bias)
        if not torch.isfinite(logits).all():
            raise ValueError('Nonfinite encoder scores')
        # Rank logits first: sigmoid saturation must not change top-k membership.
        values, indices = logits.topk(self.top_k, dim=-1)
        return indices, torch.sigmoid(values)

    def encode_iter(self, embeddings, *, batch_size=64):
        """Yield float32 SciPy CSR batches with 80,513 concept columns."""
        _batch_size(batch_size)
        if not isinstance(embeddings, (np.ndarray, torch.Tensor)):
            raise TypeError('Expected a NumPy array or Torch tensor')
        if embeddings.ndim != 2 or embeddings.shape[1] != self.input_dim:
            raise ValueError(f'Expected shape (N, {self.input_dim})')
        for start in range(0, len(embeddings), batch_size):
            ids, vals = self(embeddings[start:start + batch_size])
            ids, vals = ids.cpu().numpy().astype(np.int32), vals.cpu().numpy()
            rowptr = np.arange(0, vals.size + 1, self.top_k, dtype=np.int32)
            result = sparse.csr_matrix((vals.ravel(), ids.ravel(), rowptr), shape=(len(vals), self.num_features))
            result.eliminate_zeros()
            result.sort_indices()
            yield result

    def encode(self, embeddings, *, batch_size=64):
        batches = list(self.encode_iter(embeddings, batch_size=batch_size))
        return sparse.vstack(batches, format='csr') if batches else sparse.csr_matrix((0, self.num_features), dtype=np.float32)

    def explain(self, row, *, top_n=20, include_prevalence=False):
        """Describe an encoded row using saved labels; no LLM/API calls."""
        if not sparse.issparse(row) or row.shape != (1, self.num_features):
            raise ValueError('Pass one encoded sparse row')
        if not isinstance(top_n, int) or isinstance(top_n, bool) or top_n < 0:
            raise ValueError('top_n must be a nonnegative integer')
        row = row.tocsr(copy=True)
        row.sum_duplicates(); row.eliminate_zeros()
        order = np.lexsort((row.indices, -row.data))[:top_n]
        concepts = [dict(concept_id=int(row.indices[i]), label=self.labels[int(row.indices[i])], activation=float(row.data[i])) for i in order]
        if include_prevalence:
            for concept in concepts:
                concept['baseline_prevalence'] = self.get_prevalence(concept['concept_id'])
        return concepts


class UCEGemma(torch.nn.Module):
    """Plain strings -> text-only Gemma backbone -> concept scores.

    Includes no vision/audio towers, multimodal projection weights, or processor.
    Uses the original bfloat16 text inference and float32 concept encoder.
    """
    def __init__(self, encoder, text_model, tokenizer):
        super().__init__()
        self.encoder = encoder
        self.text_model = text_model
        self.tokenizer = tokenizer
        self.text_config = encoder.config['text_pipeline']
        self.config = encoder.config
        self.num_features = encoder.num_features
        self.labels = encoder.labels
        self.metadata = encoder.metadata
        self.eval()

    @classmethod
    def from_pretrained(cls, model_id_or_path, *, revision=None, device='cpu', cache_dir=None, local_files_only=False):
        from transformers import AutoTokenizer, EmbeddingGemma2TextModel
        directory = _directory(model_id_or_path, revision, cache_dir, local_files_only, with_embedder=True)
        encoder = UCEGemmaEncoder.from_pretrained(directory, device=device)
        # Record the Hub revision before the local path loses its repository context.
        if not Path(model_id_or_path).expanduser().is_dir():
            encoder.metadata.provenance.update(
                source_repo_id=str(model_id_or_path), requested_revision=revision,
                resolved_revision=directory.name if directory.parent.name == 'snapshots' else None)
        if 'text_pipeline' not in encoder.config or not (directory / 'gemma_embedder').is_dir():
            raise ValueError('Load the full package to encode text')
        text_model = EmbeddingGemma2TextModel.from_pretrained(
            directory / 'gemma_embedder', dtype=torch.bfloat16, local_files_only=True,
        ).to(device).eval()
        tokenizer = AutoTokenizer.from_pretrained(directory / 'gemma_embedder', local_files_only=True, trust_remote_code=False)
        return cls(encoder, text_model, tokenizer)

    @staticmethod
    def _texts(texts):
        if isinstance(texts, str):
            texts = [texts]
        if not isinstance(texts, (list, tuple)) or any(not isinstance(t, str) for t in texts):
            raise TypeError('Expected a string or a list/tuple of plain strings')
        return texts

    @torch.inference_mode()
    def embed(self, texts, *, batch_size=32):
        """Return normalized float32 [N,768] embeddings; overlength inputs raise."""
        texts = self._texts(texts)
        _batch_size(batch_size)
        batches = []
        device = next(self.text_model.parameters()).device
        for start in range(0, len(texts), batch_size):
            inputs = self.tokenizer([self.text_config['prompt'] + t for t in texts[start:start+batch_size]],
                                    padding=True, truncation=False, return_tensors='pt')
            if inputs['input_ids'].shape[1] > self.text_config['max_seq_length']:
                raise ValueError('Text exceeds 8192 tokens including prompt; split the document explicitly')
            ids, mask = inputs['input_ids'].to(device), inputs['attention_mask'].to(device)
            # Original multimodal wrapper replaces placeholder IDs by PAD even for text-only input.
            model_ids = ids.clone()
            for token_id in self.text_config['placeholder_token_ids']:
                model_ids.masked_fill_(ids == token_id, self.tokenizer.pad_token_id)
            hidden = self.text_model(input_ids=model_ids, attention_mask=mask).last_hidden_state
            expanded = mask.unsqueeze(-1).expand_as(hidden).to(hidden.dtype)
            pooled = (hidden * expanded).sum(1) / expanded.sum(1).clamp_min(1e-9)
            # Match SentenceTransformer Normalize and encode(normalize_embeddings=True).
            pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
            pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
            array = pooled.float().cpu().numpy()
            array /= np.linalg.norm(array, axis=1, keepdims=True)
            batches.append(array)
        return np.concatenate(batches) if batches else np.empty((0, self.encoder.input_dim), np.float32)

    def encode_iter(self, texts, *, batch_size=32):
        texts = self._texts(texts)
        _batch_size(batch_size)
        for start in range(0, len(texts), batch_size):
            embeddings = self.embed(texts[start:start+batch_size], batch_size=batch_size)
            yield self.encoder.encode(embeddings, batch_size=batch_size)

    def encode(self, texts, *, batch_size=32):
        batches = list(self.encode_iter(texts, batch_size=batch_size))
        return sparse.vstack(batches, format='csr') if batches else sparse.csr_matrix((0, self.encoder.num_features), dtype=np.float32)

    def forward(self, texts):
        return self.encode(texts)

    def explain(self, row, *, top_n=20, include_prevalence=False):
        return self.encoder.explain(row, top_n=top_n, include_prevalence=include_prevalence)

    def get_prevalence(self, concept_id):
        """Look up one concept's historical Gemini corpus baseline statistics."""
        return self.encoder.get_prevalence(concept_id)

def main():
    import argparse
    parser = argparse.ArgumentParser(description='Universal Concept Encoder: Gemma text to 80,513 concepts')
    parser.add_argument('--model', default=str(Path(__file__).resolve().parent), help='Local package or Hub repository ID')
    parser.add_argument('--text', required=True, help='Plain text to encode')
    parser.add_argument('--device', default='cpu')
    parser.add_argument('--top-n', type=int, default=20)
    args = parser.parse_args()
    model = UCEGemma.from_pretrained(args.model, device=args.device)
    print(json.dumps(model.explain(model.encode(args.text)[0], top_n=args.top_n), ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()