"""
LLM Encoder: Transform (input, output) pairs into semantic embeddings.

Pipeline:
  1. Claude Haiku describes what transformation was applied (natural language)
  2. sentence-transformers embeds the description into a fixed-dimensional vector
  3. The resulting embeddings capture semantic structure (what the transform DOES)
     rather than geometric correlates (statistical patterns in the data)

This is the Layer 2 replacement for RelationalEncoder/CNNProtoRelNet.
Instead of learning an encoder end-to-end (which hit a 0.53 generalization
ceiling), we leverage the LLM's pre-existing semantic understanding.

Usage:
    from qppg import LLMEncoder

    encoder = LLMEncoder(domain="vector")
    embeddings = encoder.encode(inputs, outputs)  # (N, 384)
"""

import os
import json
import time
import hashlib
import numpy as np
from pathlib import Path
from typing import List, Optional, Union


# ============================================================
# LLM Transform Describer
# ============================================================

# Prompt version — change this when the prompt template changes
# to invalidate cached descriptions
PROMPT_VERSION = "v1"

VECTOR_SYSTEM_PROMPT = """You are a mathematical analyst. Given an input vector and its transformed output, describe WHAT TRANSFORMATION was applied. Focus on the structural relationship between input and output, not the specific numbers. Be concise (2-3 sentences max)."""

VECTOR_USER_TEMPLATE = """Input vector (8 dimensions): [{input_str}]
Output vector (8 dimensions): [{output_str}]

Describe the transformation that maps input to output. Consider:
1. Element-wise relationships: Are elements negated? Scaled? Squared? Clipped?
2. Ordering changes: Are elements reversed? Shifted? Sorted?
3. Global properties: Is the mean preserved? The norm? The variance?
4. Dependencies: Does each output element depend only on the corresponding input, or on multiple input elements?"""

GRID_SYSTEM_PROMPT = """You are a visual pattern analyst. Given an input grid and its transformed output grid, describe WHAT TRANSFORMATION was applied. Focus on the structural operation, not specific cell values. Be concise (2-3 sentences max)."""

GRID_USER_TEMPLATE = """Input grid ({h_in}x{w_in}):
{input_grid}

Output grid ({h_out}x{w_out}):
{output_grid}

Describe the transformation. Consider:
1. Spatial changes: rotation, reflection, translation, scaling, cropping?
2. Color changes: remapping, thresholding, inversion?
3. Object changes: gravity, isolation, outlining, filling?
4. Structural patterns: symmetry creation, tiling, pattern repetition?"""

COLOR_MAP = {
    0: 'black', 1: 'blue', 2: 'red', 3: 'green', 4: 'yellow',
    5: 'gray', 6: 'magenta', 7: 'orange', 8: 'azure', 9: 'maroon'
}


def _cache_key(x_in, x_out, domain):
    """Generate a cache key from input/output pair."""
    data = x_in.tobytes() + x_out.tobytes() + domain.encode() + PROMPT_VERSION.encode()
    return hashlib.sha256(data).hexdigest()[:16]


def _format_grid(grid):
    """Format a 2D integer grid for LLM display."""
    lines = []
    for row in grid:
        lines.append(' '.join(str(int(v)) for v in row))
    return '\n'.join(lines)


class LLMTransformDescriber:
    """
    Generates natural language descriptions of (input, output) transformations
    using Claude Haiku via the Anthropic API.

    Parameters:
        api_key: Anthropic API key. If None, reads from ANTHROPIC_API_KEY env var.
        model: Model to use (default: claude-haiku-4-5-20251001 for cost)
        temperature: Generation temperature (0.0 for deterministic)
        max_tokens: Max response tokens (descriptions should be short)
        cache_path: Path to persist description cache (JSON file)
    """

    def __init__(self, api_key=None, model="claude-haiku-4-5-20251001",
                 temperature=0.0, max_tokens=200, cache_path=None):
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.cache = {}
        self.cache_path = cache_path

        # Load existing cache
        if cache_path and os.path.exists(cache_path):
            with open(cache_path, 'r') as f:
                self.cache = json.load(f)

        # Initialize API client
        self.api_key = api_key or os.environ.get('ANTHROPIC_API_KEY', '')

        self._client = None
        self._init_client()

        # Cost tracking
        self.total_input_tokens = 0
        self.total_output_tokens = 0
        self.total_calls = 0
        self.cache_hits = 0

    def _init_client(self):
        try:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self.api_key)
        except ImportError:
            print("WARNING: anthropic package not installed. pip install anthropic")
            self._client = None

    def _call_api(self, system_prompt, user_prompt, max_retries=5):
        """Call Claude API with exponential backoff retry."""
        for attempt in range(max_retries):
            try:
                response = self._client.messages.create(
                    model=self.model,
                    max_tokens=self.max_tokens,
                    temperature=self.temperature,
                    system=system_prompt,
                    messages=[{"role": "user", "content": user_prompt}]
                )
                self.total_input_tokens += response.usage.input_tokens
                self.total_output_tokens += response.usage.output_tokens
                self.total_calls += 1
                return response.content[0].text
            except Exception as e:
                if attempt < max_retries - 1:
                    wait = 2 ** attempt + np.random.random()
                    time.sleep(wait)
                else:
                    print(f"API call failed after {max_retries} retries: {e}")
                    return f"[API Error: {str(e)[:100]}]"

    def describe_vector_transform(self, x_in, x_out):
        """
        Describe a R^8 vector transformation.

        Args:
            x_in: (8,) numpy array — input vector
            x_out: (8,) numpy array — output vector

        Returns:
            str: Natural language description of the transformation
        """
        key = _cache_key(x_in, x_out, "vector")
        if key in self.cache:
            self.cache_hits += 1
            return self.cache[key]

        input_str = ', '.join(f'{v:.3f}' for v in x_in)
        output_str = ', '.join(f'{v:.3f}' for v in x_out)
        user_prompt = VECTOR_USER_TEMPLATE.format(
            input_str=input_str, output_str=output_str
        )

        description = self._call_api(VECTOR_SYSTEM_PROMPT, user_prompt)
        self.cache[key] = description
        return description

    def describe_grid_transform(self, grid_in, grid_out):
        """
        Describe an ARC grid transformation.

        Args:
            grid_in: (H, W) numpy array — input grid (integer values 0-9)
            grid_out: (H, W) numpy array — output grid

        Returns:
            str: Natural language description of the transformation
        """
        key = _cache_key(grid_in, grid_out, "grid")
        if key in self.cache:
            self.cache_hits += 1
            return self.cache[key]

        h_in, w_in = grid_in.shape
        h_out, w_out = grid_out.shape
        user_prompt = GRID_USER_TEMPLATE.format(
            h_in=h_in, w_in=w_in, input_grid=_format_grid(grid_in),
            h_out=h_out, w_out=w_out, output_grid=_format_grid(grid_out),
        )

        description = self._call_api(GRID_SYSTEM_PROMPT, user_prompt)
        self.cache[key] = description
        return description

    def describe_batch(self, inputs, outputs, domain="vector"):
        """
        Describe a batch of (input, output) pairs.

        Args:
            inputs: list of numpy arrays (inputs)
            outputs: list of numpy arrays (outputs)
            domain: "vector" or "grid"

        Returns:
            List[str]: descriptions for each pair
        """
        describe_fn = self.describe_vector_transform if domain == "vector" \
            else self.describe_grid_transform

        descriptions = []
        for i, (x_in, x_out) in enumerate(zip(inputs, outputs)):
            desc = describe_fn(x_in, x_out)
            descriptions.append(desc)

            # Progress logging every 100 calls
            if (i + 1) % 100 == 0:
                print(f"  Described {i+1}/{len(inputs)} pairs "
                      f"(cache hits: {self.cache_hits}, API calls: {self.total_calls})")

        return descriptions

    def save_cache(self):
        """Persist description cache to disk."""
        if self.cache_path:
            os.makedirs(os.path.dirname(self.cache_path) or '.', exist_ok=True)
            with open(self.cache_path, 'w') as f:
                json.dump(self.cache, f)

    def get_cost_summary(self):
        """Return cumulative API cost tracking."""
        # Haiku pricing: $0.80/1M input, $4.00/1M output
        cost_input = self.total_input_tokens * 0.80 / 1_000_000
        cost_output = self.total_output_tokens * 4.00 / 1_000_000
        return {
            'total_calls': self.total_calls,
            'cache_hits': self.cache_hits,
            'total_input_tokens': self.total_input_tokens,
            'total_output_tokens': self.total_output_tokens,
            'cost_usd': round(cost_input + cost_output, 4),
        }


# ============================================================
# Text Embedder (sentence-transformers)
# ============================================================

class TextEmbedder:
    """
    Embeds text descriptions into fixed-dimensional vectors using
    sentence-transformers. Runs locally on CPU — no API cost.

    Parameters:
        model_name: sentence-transformers model (default: all-MiniLM-L6-v2)
                    Produces 384-dimensional embeddings. Fast and CPU-friendly.
    """

    def __init__(self, model_name="all-MiniLM-L6-v2"):
        self.model_name = model_name
        self._model = None

    def _load_model(self):
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self.model_name)

    def embed(self, texts):
        """
        Embed one or more texts.

        Args:
            texts: str or List[str]

        Returns:
            numpy array of shape (N, d_embed) where d_embed=384 for MiniLM
        """
        self._load_model()
        if isinstance(texts, str):
            texts = [texts]
        embeddings = self._model.encode(texts, show_progress_bar=False)
        return embeddings

    def embed_and_normalize(self, texts):
        """Embed and L2-normalize."""
        embeddings = self.embed(texts)
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        norms = np.clip(norms, 1e-8, None)
        return embeddings / norms

    @property
    def embedding_dim(self):
        self._load_model()
        return self._model.get_sentence_embedding_dimension()


# ============================================================
# LLM Encoder: Complete Pipeline
# ============================================================

class LLMEncoder:
    """
    Complete pipeline: (input, output) → LLM description → text embedding.

    Drop-in replacement for RelationalEncoder/CNNProtoRelNet.
    Uses Claude Haiku for semantic understanding + sentence-transformers
    for embedding. Caches descriptions to minimize API costs.

    Parameters:
        api_key: Anthropic API key (or set ANTHROPIC_API_KEY env var)
        domain: "vector" (R^8) or "grid" (ARC grids)
        embedding_model: sentence-transformers model name
        cache_path: Path to persist LLM description cache
        llm_model: Claude model to use
    """

    def __init__(self, api_key=None, domain="vector",
                 embedding_model="all-MiniLM-L6-v2",
                 cache_path=None, llm_model="claude-haiku-4-5-20251001"):
        self.domain = domain
        self.describer = LLMTransformDescriber(
            api_key=api_key, model=llm_model, cache_path=cache_path
        )
        self.embedder = TextEmbedder(model_name=embedding_model)

    def encode(self, inputs, outputs):
        """
        Encode (input, output) pairs into semantic embeddings.

        Args:
            inputs: list of numpy arrays or (N, d) array
            outputs: list of numpy arrays or (N, d) array

        Returns:
            (N, d_embed) numpy array of L2-normalized embeddings
        """
        # Convert arrays to lists if needed
        if isinstance(inputs, np.ndarray) and inputs.ndim == 2:
            inputs = [inputs[i] for i in range(len(inputs))]
            outputs = [outputs[i] for i in range(len(outputs))]

        # Step 1: Get LLM descriptions
        descriptions = self.describer.describe_batch(inputs, outputs, domain=self.domain)

        # Step 2: Embed descriptions
        embeddings = self.embedder.embed_and_normalize(descriptions)

        return embeddings

    def encode_with_descriptions(self, inputs, outputs):
        """
        Like encode(), but also returns the raw text descriptions.

        Returns:
            embeddings: (N, d_embed) numpy array
            descriptions: List[str]
        """
        if isinstance(inputs, np.ndarray) and inputs.ndim == 2:
            inputs = [inputs[i] for i in range(len(inputs))]
            outputs = [outputs[i] for i in range(len(outputs))]

        descriptions = self.describer.describe_batch(inputs, outputs, domain=self.domain)
        embeddings = self.embedder.embed_and_normalize(descriptions)

        return embeddings, descriptions

    def save_cache(self):
        """Persist description cache to disk."""
        self.describer.save_cache()

    def load_cache(self):
        """Load cache (done automatically in constructor if cache_path exists)."""
        pass

    def get_cost_summary(self):
        """Return API cost summary."""
        return self.describer.get_cost_summary()

    @property
    def embedding_dim(self):
        return self.embedder.embedding_dim
