"""
NCA 3D Brain — Inference utilities
Generate text and predict next words with the 3D cellular automaton.
"""
import torch
import torch.nn.functional as F
import json
from model import NCA3D_Fatigue, PUNCT_MAP, PUNCT_INV, PAD_ID, EOS_ID

REP_PENALTY = 1.3
MAX_WORDS = 50


def id_to_word(wid, num2word):
    """Convert word ID to string."""
    if wid < 30000:
        return num2word.get(wid, f"[{wid}]")
    return {PAD_ID: "<PAD>", EOS_ID: "<EOS>"}.get(wid, PUNCT_INV.get(wid, f"[{wid}]"))


def words_to_ids(words, word2num):
    """Convert list of words to IDs, handling punctuation."""
    ids = []
    for w in words:
        if w in PUNCT_MAP:
            ids.append(PUNCT_MAP[w])
        elif w in word2num:
            ids.append(word2num[w])
    return ids


def predict(model, word2num, num2word, context_words, n_steps=15):
    """
    Predict the next word given a context.

    Returns:
        (predicted_word, top5_list) where top5_list is [(word, probability), ...]
    """
    ids = words_to_ids(context_words, word2num)
    if not ids:
        return "???", []
    with torch.no_grad():
        inp = torch.tensor([ids], dtype=torch.long)
        logits = model(inp, n_steps)
        probs = F.softmax(logits[0].float(), dim=-1)
        pred_id = logits.argmax(-1).item()
        top5 = logits[0].topk(5)
    pred_word = id_to_word(pred_id, num2word)
    top5_list = [
        (id_to_word(i.item(), num2word), f"{probs[i].item()*100:.1f}%")
        for i in top5.indices
    ]
    return pred_word, top5_list


def generate(model, word2num, num2word, seed_words, max_words=10, n_steps=15, temperature=None):
    """
    Generate text autoregressively from seed words.

    Args:
        model: NCA3D_Fatigue model
        word2num: word → id dictionary
        num2word: id → word dictionary
        seed_words: list of starting words
        max_words: maximum words to generate
        n_steps: propagation steps per prediction
        temperature: sampling temperature (None = greedy/argmax)

    Returns:
        list of all words (seed + generated)
    """
    words = list(seed_words)
    generated_ids = []
    for _ in range(max_words):
        ids = words_to_ids(words, word2num)
        if not ids or len(ids) > MAX_WORDS:
            break
        with torch.no_grad():
            inp = torch.tensor([ids], dtype=torch.long)
            logits = model(inp, n_steps).float()
            # Repetition penalty
            for prev_id in generated_ids:
                logits[0, prev_id] /= REP_PENALTY
            if temperature and temperature != 0:
                logits = logits / temperature
                probs = F.softmax(logits[0], dim=-1)
                pred_id = torch.multinomial(probs, 1).item()
            else:
                pred_id = logits.argmax(-1).item()
        if pred_id == PAD_ID or pred_id == EOS_ID:
            break
        word = id_to_word(pred_id, num2word)
        if word.startswith("[") or word == "???":
            break
        words.append(word)
        generated_ids.append(pred_id)
    return words


def load_model(model_path="model_phase4c_v5_fatigue_best.pt", device="cpu"):
    """Load the NCA 3D Brain model."""
    model = NCA3D_Fatigue()
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    return model


def load_dictionary(dict_path="word_dictionary_30k.json"):
    """Load the 30K word dictionary."""
    word2num = {k: int(v) for k, v in json.load(open(dict_path)).items()}
    num2word = {v: k for k, v in word2num.items()}
    return word2num, num2word


if __name__ == "__main__":
    print("Loading NCA 3D Brain v5...")
    model = load_model()
    word2num, num2word = load_dictionary()
    params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"Model loaded: {params:.1f}M parameters\n")

    # Demo: next word prediction
    test_contexts = [
        ["the", "little", "girl"],
        ["he", "said"],
        ["she", "wanted", "to"],
        ["in", "the", "morning"],
        ["the", "old", "man"],
        ["the", "dog", "ate", "the"],
    ]

    print("=== Next Word Prediction ===\n")
    for ctx in test_contexts:
        pred, top5 = predict(model, word2num, num2word, ctx)
        top5_str = ", ".join([f"{w} ({p})" for w, p in top5])
        print(f"  '{' '.join(ctx)}' -> '{pred}'")
        print(f"    top 5: {top5_str}\n")

    # Demo: text generation
    print("=== Text Generation ===\n")
    seeds = [
        ["the", "little", "girl"],
        ["he", "said", "that"],
        ["she", "wanted", "to"],
        ["in", "the", "morning"],
        ["the", "old", "man"],
        ["one", "day"],
    ]

    for seed in seeds:
        result = generate(model, word2num, num2word, seed, max_words=10)
        print(f"  {' '.join(result)}")

    print("\n=== Generation with temperature ===\n")
    for temp in [0.5, 1.0, 1.5]:
        print(f"  Temperature {temp}:")
        for seed in seeds[:3]:
            result = generate(model, word2num, num2word, seed, max_words=10, temperature=temp)
            print(f"    {' '.join(result)}")
        print()