"""phenome playground — a genome IS the system prompt.

Pick a preset persona or roll a random genome. A compact ~126-char symbolic
string conditions the model's whole personality — verbosity, warmth, humor,
structure, how it pushes back — with no written system prompt at all.

Model: nyxia/phenome-qwen3.5-4b  (Qwen3.5-4B + phenome LoRA, merged)
Idea:  https://huggingface.co/blog/nyxia/genetics-instead-of-system-prompts-for-ai-agents
"""
from __future__ import annotations

import random
from pathlib import Path

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessor, LogitsProcessorList

from phenome.genome import roll_genome
from phenome.prompt import (
    build_training_system_prompt,
    compact_string_to_genome_loci,
    genome_to_compact_string,
)
from phenome.registry import load_registry
from presets import PRESETS

HERE = Path(__file__).resolve().parent
MODEL_ID = "nyxia/phenome-qwen3.5-4b"

# ── load engine + model (module level; .to('cuda') works under ZeroGPU emulation)
REG = load_registry(HERE / "data" / "loci.json")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16, trust_remote_code=True,
).to("cuda")
model.eval()

RANDOM_NAMES = [
    "Vesper", "Orin", "Juno", "Phoenix", "River", "Cassia", "Atlas", "Iris",
    "Nico", "Sable", "Wren", "Lior", "Mira", "Soren", "Ezra", "Marlow",
]

# Recommended Qwen3.5 instruct (non-thinking) sampling. presence_penalty isn't a
# transformers `generate` kwarg (it's a vLLM/serving param) — implement it as a
# LogitsProcessor so the demo matches the model card's recommended settings:
#   temperature=0.7, top_p=0.80, top_k=20, min_p=0.0,
#   presence_penalty=1.5, repetition_penalty=1.0
PRESENCE_PENALTY = 1.5
TOP_P = 0.80
TOP_K = 20
MIN_P = 0.0
REPETITION_PENALTY = 1.0


class PresencePenaltyLogitsProcessor(LogitsProcessor):
    """OpenAI/vLLM-style presence penalty: subtract a flat penalty from the
    logits of any token already generated (output tokens only, not the prompt)."""

    def __init__(self, penalty: float):
        self.penalty = penalty
        self.prompt_len = None

    def __call__(self, input_ids, scores):
        if self.penalty == 0.0:
            return scores
        if self.prompt_len is None:
            self.prompt_len = input_ids.shape[1]
        for i in range(scores.shape[0]):
            gen = input_ids[i, self.prompt_len:]
            if gen.numel():
                scores[i, torch.unique(gen)] -= self.penalty
        return scores


# ── genome helpers ──────────────────────────────────────────────────────────
def _genome_from_alpha6(alpha6: str):
    """Reconstruct a Genome object from its alpha6 string (for prompt building)."""
    proto = roll_genome(REG)
    proto.loci = compact_string_to_genome_loci(alpha6, REG)
    proto.seed = None
    return proto


# intensity dial 0–3 → plain words (1 = baseline, shown bare)
_INTENSITY_WORD = {0: " *(subtle)*", 1: "", 2: " *(strong)*", 3: " *(defining)*"}


def _allele_phrase(tag: str, intensity: int) -> str:
    return tag.replace("_", " ") + _INTENSITY_WORD.get(intensity, "")


def _readable_traits(alpha6: str) -> str:
    """Render the genome as a human-readable Markdown table.

    Each locus carries two alleles (it's diploid). When both alleles are the
    same we show it once at the stronger intensity; otherwise we show the pair
    as "A + B". The raw alpha6 / named string is still available above.
    """
    loci = compact_string_to_genome_loci(alpha6, REG)
    rows = ["| Trait | Expression |", "| --- | --- |"]
    for locus_name, ((a_tag, a_int), (b_tag, b_int)) in loci.items():
        trait = locus_name.replace("_", " ").title()
        if a_tag == b_tag:
            expr = _allele_phrase(a_tag, max(a_int, b_int))
        else:
            expr = f"{_allele_phrase(a_tag, a_int)} + {_allele_phrase(b_tag, b_int)}"
        rows.append(f"| {trait} | {expr} |")
    return "\n".join(rows)


def set_preset(preset_name: str):
    """Return (alpha6, persona_name, traits_display, system_prompt)."""
    if preset_name == "🎲 Random":
        g = roll_genome(
            REG, seed=random.randint(0, 2**31 - 1),
            include_categories=("general", "roleplay", "creative"),
            exclude_categories=("nsfw",), only_default_on=False,
        )
        alpha6 = genome_to_compact_string(g, REG, format="alpha6")
        name = random.choice(RANDOM_NAMES)
    else:
        _emoji, _tag, name, alpha6 = PRESETS[preset_name]
        g = _genome_from_alpha6(alpha6)
    system = build_training_system_prompt(g, REG, name)
    return alpha6, name, _readable_traits(alpha6), system


def reroll():
    return set_preset("🎲 Random")


# ── generation ──────────────────────────────────────────────────────────────
@spaces.GPU(duration=60)
def respond(message, history, alpha6, name, temperature, max_new_tokens):
    if not alpha6:
        a6, name, _, _ = set_preset("🎲 Random")
        alpha6 = a6
    g = _genome_from_alpha6(alpha6)
    system = build_training_system_prompt(g, REG, name)

    messages = [{"role": "system", "content": system}]
    for turn in history:
        messages.append(turn)
    messages.append({"role": "user", "content": message})

    templated = tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True, enable_thinking=False,
    )
    if hasattr(templated, "input_ids"):
        templated = templated["input_ids"]
    if templated and isinstance(templated[0], list):
        templated = templated[0]
    ids = torch.tensor([templated], device=model.device)

    im_end = tokenizer.convert_tokens_to_ids("<|im_end|>")
    eos_ids = [tokenizer.eos_token_id]
    if im_end is not None and im_end != tokenizer.eos_token_id:
        eos_ids.append(im_end)

    with torch.no_grad():
        out = model.generate(
            ids,
            max_new_tokens=int(max_new_tokens),
            do_sample=temperature > 0,
            temperature=float(temperature),
            top_p=TOP_P,
            top_k=TOP_K,
            min_p=MIN_P,
            repetition_penalty=REPETITION_PENALTY,
            logits_processor=LogitsProcessorList([PresencePenaltyLogitsProcessor(PRESENCE_PENALTY)]),
            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
            eos_token_id=eos_ids,
        )
    text = tokenizer.decode(out[0, ids.shape[1]:].tolist(), skip_special_tokens=True).strip()
    return text


# ── UI ──────────────────────────────────────────────────────────────────────
CSS = """
#genome-box textarea { font-family: ui-monospace, monospace; font-size: 12px; }
#traits-table table { font-size: 12.5px; width: 100%; border-collapse: collapse; }
#traits-table td, #traits-table th { padding: 2px 8px; }
#traits-table td:first-child { font-weight: 600; white-space: nowrap; }
.phenome-title { text-align:center; }
footer { visibility: hidden; }
"""

PRESET_CHOICES = ["🎲 Random"] + [f"{e} {k}" for k, (e, _t, _nm, _a6) in PRESETS.items()]
_LABEL_TO_NAME = {f"{e} {k}": k for k, (e, _t, _nm, _a6) in PRESETS.items()}
_LABEL_TO_NAME["🎲 Random"] = "🎲 Random"
DEFAULT_LABEL = "✨ The Method Actor"


def on_preset_label(label):
    return set_preset(_LABEL_TO_NAME[label])


with gr.Blocks(title="phenome playground") as demo:
    gr.Markdown(
        """
        <div class="phenome-title">

        # 🧬 phenome playground
        ### a *genome* is the system prompt

        A ~126-character symbolic string encodes a full persona — verbosity, warmth, humor,
        structure, how hard it pushes back — and a LoRA learned to read it directly.
        No written system prompt. Pick a preset or roll the dice; same model, different genes.

        [model](https://huggingface.co/nyxia/phenome-qwen3.5-4b) ·
        [the idea (blog)](https://huggingface.co/blog/nyxia/genetics-instead-of-system-prompts-for-ai-agents)
        </div>
        """
    )

    alpha6_state = gr.State("")
    name_state = gr.State("")

    with gr.Row():
        with gr.Column(scale=2):
            preset = gr.Dropdown(
                PRESET_CHOICES, value=DEFAULT_LABEL, label="Persona preset",
                info="Curated genomes — or 🎲 Random to roll fresh genes.",
            )
            with gr.Row():
                name_box = gr.Textbox(label="Name", scale=2)
                roll_btn = gr.Button("🎲 Reroll random", scale=1)
            with gr.Accordion("🔬 the genome behind this persona", open=True):
                sysprompt_box = gr.Textbox(
                    label="full system prompt — the ENTIRE prompt the model gets (the genome IS the prompt)",
                    elem_id="genome-box", lines=4, interactive=False,
                )
                gr.Markdown("**decoded traits** — human-readable, *not* sent to the model")
                traits_box = gr.Markdown(elem_id="traits-table")
            with gr.Accordion("⚙️ generation settings", open=False):
                temperature = gr.Slider(0.0, 1.2, value=0.6, step=0.05, label="temperature")
                max_tokens = gr.Slider(32, 512, value=220, step=16, label="max new tokens")

        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=480, label="chat")
            gr.ChatInterface(
                fn=respond,
                chatbot=chatbot,
                additional_inputs=[alpha6_state, name_state, temperature, max_tokens],
                # Prompts chosen to SHOW the genome. The 200-sweep found that
                # "tell me about yourself", "help with an email", and "I'm
                # overwhelmed" flatten every persona into the same boilerplate —
                # dropped. These each surface a different axis.
                examples=[
                    ["We just met at a tavern. Greet me."],
                    ["What do you think about cats?"],
                    ["Explain how a bug in my code might happen."],
                    ["I disagree with you. Push back on me."],
                    ["Be brief — how do I get started?"],
                ],
                cache_examples=False,
            )

    # wire persona changes → state + displays, and clear the chat on persona change
    def _apply(label):
        a6, nm, tr, sp = on_preset_label(label)
        return a6, nm, nm, tr, sp, []

    preset.change(
        _apply, inputs=preset,
        outputs=[alpha6_state, name_state, name_box, traits_box, sysprompt_box, chatbot],
    )

    def _reroll():
        a6, nm, tr, sp = reroll()
        return a6, nm, nm, tr, sp, []

    roll_btn.click(
        _reroll,
        outputs=[alpha6_state, name_state, name_box, traits_box, sysprompt_box, chatbot],
    )

    # editing the name updates state + refreshes the displayed system prompt
    def _rename(new_name, alpha6):
        new_name = new_name or "Assistant"
        g = _genome_from_alpha6(alpha6)
        return new_name, build_training_system_prompt(g, REG, new_name)

    name_box.submit(_rename, inputs=[name_box, alpha6_state], outputs=[name_state, sysprompt_box])

    # initial load → populate the default preset
    def _init():
        a6, nm, tr, sp = on_preset_label(DEFAULT_LABEL)
        return a6, nm, nm, tr, sp

    demo.load(_init, outputs=[alpha6_state, name_state, name_box, traits_box, sysprompt_box])


if __name__ == "__main__":
    demo.launch(css=CSS, theme=gr.themes.Soft())