"""PALIMPSESTE — Creative generation via mixture logits.

The core problem: _logits() collapses K retrieved candidates into ONE HV via
majority sign, then takes argmax. This is structurally incapable of producing
novel tokens — it can only echo the majority.

The solution: Mixture Logits. Instead of one bundle, we use ALL K retrieved
candidates as weighted votes over the vocabulary:

    mixture_logits[t] = sum_k( sim_k * similarity(token_hv[t], value_hv_k) )

A token that's close to MULTIPLE candidates gets boosted beyond what any single
candidate would predict. This is how novel compositions emerge — the model
finds tokens at the intersection of multiple knowledge fragments.

Additionally: anti-repetition via n-gram blocking prevents attractor loops.
"""

from __future__ import annotations

import numpy as np
from dataclasses import dataclass
from typing import Optional

from .hv import HV, bind, similarity, bits_to_signs

__all__ = ["CreativeGenerator", "CreativeResult"]


@dataclass
class CreativeResult:
    """Result of creative generation."""
    text: str
    n_candidates_per_step: list[int]
    novel_tokens: int
    temperature: float
    mixture_diversity: float


class CreativeGenerator:
    """Generate text using mixture logits from multiple retrieved candidates.

    Instead of retrieving one bundled HV, this generator retrieves K matches
    and builds a mixture distribution over the vocabulary. Tokens at the
    intersection of multiple candidates get boosted — enabling novel
    compositions that were never explicitly stored.

    This is how PALIMPSESTE creates, not just repeats.
    """

    def __init__(self, lm, top_k: int = 5, diversity: float = 0.3) -> None:
        """
        Parameters
        ----------
        lm : PalimpsesteForCausalLM
            The language model.
        top_k : int
            Number of retrieved candidates to use per step.
        diversity : float
            How much to boost intersection tokens (0=none, 1=max).
        """
        self.lm = lm
        self.top_k = top_k
        self.diversity = diversity

    def mixture_logits(self, context_ids: list[int]) -> np.ndarray:
        """Compute mixture logits from K retrieved candidates.

        For each of the top-K retrieved traces, compute similarity between
        the trace's value HV and every token HV. Weight by retrieval sim.
        Sum to get mixture logits.

        This produces a RICHER distribution than single-HV retrieval because
        tokens at the intersection of candidates get extra boost.
        """
        assert self.lm.tokenizer is not None

        s = self.lm._state_hv(context_ids)
        q = bind(self.lm._self_hv, s)

        # Retrieve K matches (NOT the bundled result)
        ret = self.lm.phi.retrieve(self.lm.mem, q)

        V = getattr(self.lm.tokenizer, 'vocab_size_actual', None) or self.lm.config.vocab_size
        logits = np.zeros(V, dtype=np.float64)

        if not ret.matches:
            return logits

        # Sort by similarity, take top-K
        scored = sorted(zip(ret.matches, ret.sims), key=lambda x: x[1], reverse=True)
        top_matches = scored[:self.top_k]

        if not top_matches:
            return logits

        # Ensure token bits cache is built
        packed_len = len(top_matches[0][0].value.bits)
        if (self.lm._token_bits_cache is None or
                self.lm._token_bits_cache.shape[0] != V or
                self.lm._token_bits_cache.shape[1] != packed_len):
            self.lm._token_bits_cache = np.empty((V, packed_len), dtype=np.uint8)
            for tid in range(V):
                self.lm._token_bits_cache[tid] = self.lm.tokenizer.token_hv(tid).bits
        tok_bits = self.lm._token_bits_cache

        _POPCOUNT = np.array([bin(i).count('1') for i in range(256)], dtype=np.uint16)
        D = self.lm.config.D

        # For each candidate, compute token logits and add weighted
        for match, sim in top_matches:
            value_bits = match.value.bits
            xored = np.bitwise_xor(tok_bits, value_bits[np.newaxis, :])
            hamming_dists = _POPCOUNT[xored].sum(axis=1)
            candidate_logits = 1.0 - 2.0 * hamming_dists / D

            # Weight by retrieval similarity (map [-1,1] → [0,1])
            weight = (sim + 1.0) / 2.0

            # Add weighted logits
            logits += weight * candidate_logits

            # Diversity boost: tokens that are above threshold in this candidate
            # AND already high in logits get extra boost (intersection effect)
            if self.diversity > 0:
                high_mask = candidate_logits > 0.0
                intersection = logits > 0.0
                boost_mask = high_mask & intersection
                logits[boost_mask] += self.diversity * 0.1

        # Normalize
        if logits.max() > logits.min():
            logits = (logits - logits.min()) / (logits.max() - logits.min())

        return logits

    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 100,
        temperature: float = 0.0,
        ngram_block: int = 3,
        seed: int | None = None,
    ) -> CreativeResult:
        """Generate text creatively using mixture logits.

        Parameters
        ----------
        prompt : str
            Input text to continue.
        max_new_tokens : int
            Maximum tokens to generate.
        temperature : float
            0.0 = greedy on mixture, >0 = sample from mixture.
        ngram_block : int
            Block repetition of n-grams of this length (3 = no trigram repeats).
        seed : int | None
            RNG seed for reproducibility.
        """
        from .lm import BOS, EOS, PAD
        assert self.lm.tokenizer is not None

        tok = self.lm.tokenizer
        rng = np.random.default_rng(seed) if seed is not None else np.random.default_rng()

        ids = tok.encode(prompt, add_bos=True, add_eos=True) + [BOS]
        generated: list[int] = []
        candidates_per_step: list[int] = []
        novel = 0

        for step in range(max_new_tokens):
            logits = self.mixture_logits(ids)

            if logits.max() == 0.0 and logits.min() == 0.0:
                break  # cold start

            # N-gram blocking: penalize tokens that would complete a recent n-gram
            if ngram_block > 0 and len(generated) >= ngram_block - 1:
                recent_ngram = tuple(generated[-(ngram_block - 1):]) if ngram_block > 1 else ()
                # Check if this n-gram appeared before
                for i in range(len(generated) - ngram_block + 1):
                    ng = tuple(generated[i:i + ngram_block - 1]) if ngram_block > 1 else ()
                    if ng == recent_ngram and i + ngram_block - 1 < len(generated):
                        blocked_token = generated[i + ngram_block - 1]
                        logits[blocked_token] *= 0.3  # penalize but don't forbid

            # Select next token
            if temperature <= 1e-6:
                next_id = int(np.argmax(logits))
            else:
                z = logits / temperature
                z -= z.max()
                e = np.exp(z)
                probs = e / e.sum()
                probs[PAD] = 0.0
                s = probs.sum()
                if s <= 0:
                    break
                probs /= s
                next_id = int(rng.choice(len(probs), p=probs))

            # Track if this token is "novel" (different from what standard
            # retrieval would predict)
            standard_logits = self.lm._logits(ids)
            standard_next = int(np.argmax(standard_logits))
            if next_id != standard_next:
                novel += 1

            generated.append(next_id)
            ids.append(next_id)
            candidates_per_step.append(self.top_k)

            if next_id == EOS:
                break

        text = tok.decode(generated)
        avg_diversity = novel / max(len(generated), 1)

        return CreativeResult(
            text=text,
            n_candidates_per_step=candidates_per_step,
            novel_tokens=novel,
            temperature=temperature,
            mixture_diversity=avg_diversity,
        )

    def respond(self, question: str, max_new_tokens: int = 100,
                temperature: float = 0.0) -> str:
        """Respond to a question using creative generation."""
        assert self.lm.tokenizer is not None
        q_ids = self.lm.tokenizer.encode(question, add_bos=True, add_eos=True)
        ctx = q_ids + [1]  # BOS
        result = self.generate(question, max_new_tokens=max_new_tokens,
                               temperature=temperature)
        return result.text