"""Gradio app "Biegen": Neural_Bending's tools for Stable Diffusion 1.5, on a map of the network.

The map (bendmap.js) is the control surface: pick tools, choose where they act by clicking the
network, drag their time windows. Every run renders the same seeds without tools too (cached), shows
both images with a wipe, and draws on the map how much each block of the UNet changed.
"""

from __future__ import annotations

import html
import json
import math
import os
import shutil
import tempfile
import threading
from collections import OrderedDict
from dataclasses import replace

os.environ.setdefault("MODEL_FAMILY", "sd15")

import gradio as gr  # noqa: E402

import engine as E  # noqa: E402
import measure as M  # noqa: E402
import safety as S  # noqa: E402
import sites as W  # noqa: E402
import tools_sd15 as T  # noqa: E402
from family import ACTIVE as FAM  # noqa: E402

if FAM.key != "sd15":
    raise SystemExit("Biegen runs on Stable Diffusion 1.5 only (MODEL_FAMILY=sd15).")

HERE = os.path.dirname(os.path.abspath(__file__))
WORKSHOP_PASSWORD = os.environ.get("WORKSHOP_PASSWORD")
OUT_DIR = os.environ.get("BIEGEN_OUT") or tempfile.mkdtemp(prefix="biegen-")
KEEP_RUNS = 40                      # result folders kept on disk, older ones are deleted
PLAIN_CACHE = 12                    # renders without tools kept in memory

PROMPT_A = "a photograph of a horse standing in a field"
PROMPT_B = "carved from cracked white marble, veined stone, museum lighting"
GRID_CHOICES = [(f"{k} × {k}", k * k) for k in range(1, 5)]
GRID_DEFAULT = 4
N_SLOTS = 3
SLOT_COLORS = ("#e8590c", "#1c7ed6", "#2f9e44")    # orange, blue, green: tell apart in light and dark mode
ICONS = {"excise": "✂", "layered": "◫", "positions": "⇄", "blind": "◐", "erosion": "≋", "temperature": "✱",
         "skips": "⤳", "mirror": "⧉", "freeze": "❄", "boil": "♨", "norm": "◩", "vae": "▦", "aphasia": "⌀",
         "guidance": "↺", "bypass": "◌", "shakes": "⌇", "time": "⧗", "ghost": "▞"}
TIME_KNOB = ("blind", "freeze")     # tools whose knob is a moment in the schedule
# Tools without a window act on the whole run. Why, in the words of the timeline:
WHOLE_RUN = {"aphasia": "ändert den Kontext, der in jedem Schritt derselbe ist"}
ENVELOPE = {"shakes": "sine"}       # strength inside the window: half a sine, 0 at both ends
# Typed numbers may go past every slider. These limits stay.
HARD_LIMITS = {"seed": ("Seed", -2 ** 62, 2 ** 62), "steps": ("Schritte", 1, 250)}
# One job: the largest grid at the default steps, with and without tools. Nobody holds the queue longer.
MAX_WORK = 2 * max(n for _, n in GRID_CHOICES)

THEME = gr.themes.Soft(
    primary_hue="indigo", neutral_hue="slate",
    font=[gr.themes.GoogleFont("Atkinson Hyperlegible"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"])

# Gradio's slider snaps a typed value back into its range when the number box loses focus.
FREE_SLIDER_JS = """<script>
window.addEventListener("blur", (e) => {
  const box = e.target;
  if (box instanceof HTMLInputElement && box.type === "number" && box.closest(".free-slider")) {
    e.stopImmediatePropagation();
  }
}, true);
</script>"""


class FreeSlider(gr.Slider):
    """A slider whose number box takes values beyond the slider's range (with FREE_SLIDER_JS).
    No __init__ of its own: Gradio reads the component's settings from the __init__ signature."""

    is_template = True

    def preprocess(self, payload: float) -> float:
        return gr.Number.round_to_precision(payload, self.precision)


def free_slider(minimum: float, maximum: float, **kwargs) -> FreeSlider:
    return FreeSlider(minimum, maximum, elem_classes=["free-slider"], **kwargs)


def read(name: str) -> str:
    with open(os.path.join(HERE, name), encoding="utf-8") as f:
        return f.read()


# ---------------------------------------------------------------- engine

_engine = None
_engine_lock = threading.Lock()
_sec_per_image = None
_tokenizer = None


def get_engine():
    global _engine, _sec_per_image
    with _engine_lock:
        if _engine is None:
            _engine = E.load_engine()
            print(f"[app] warm-up {_engine.warmup():.1f} s", flush=True)
            seconds = _engine.warmup(n=GRID_DEFAULT, steps=E.DEFAULT_STEPS)
            _sec_per_image = seconds / GRID_DEFAULT
            print(f"[app] {GRID_DEFAULT} images at the defaults in {seconds:.1f} s", flush=True)
    return _engine


def tokenizer():
    """CLIP's tokenizer on its own, so the map can show tokens before the model has loaded."""
    global _tokenizer
    if _tokenizer is None:
        from transformers import CLIPTokenizer

        _tokenizer = CLIPTokenizer.from_pretrained(E.MODEL_ID, subfolder="tokenizer")
    return _tokenizer


def sec_per_image() -> float:
    return _sec_per_image or FAM.sec_per_image


def duration(seconds: float) -> str:
    return f"etwa {round(seconds / 60)} Minuten" if seconds >= 90 else f"etwa {round(seconds)} Sekunden"


# ---------------------------------------------------------------- the map's catalogue and state

def catalog() -> dict:
    """Everything bendmap.js needs to know about the tools and the network."""
    tools = []
    for t in T.TOOLS.values():
        places = []
        for label, value in (t.places or (("", ""),)):
            v, start, end = t.defaults_for(value)
            places.append({"value": value, "label": label, "sites": list(W.sites(t.key, value)),
                           "default": v, "start": round(start * 100), "end": round(end * 100)})
        tools.append({"key": t.key, "label": t.label, "group": t.group, "what": t.what, "knob": t.knob,
                      "lo": t.lo, "hi": t.hi, "step": t.step, "neutral": t.neutral, "preset": t.preset,
                      "window": t.window is not None, "time_knob": t.key in TIME_KNOB,
                      "whole": WHOLE_RUN.get(t.key, "ändert Gewichte"), "envelope": ENVELOPE.get(t.key, ""),
                      "after_loop": t.key == "vae", "icon": ICONS.get(t.key, "•"), "places": places})
    return {"tools": tools, "groups": list(dict.fromkeys(t.group for t in T.TOOLS.values())),
            "colors": list(SLOT_COLORS), "max_slots": N_SLOTS}


def default_slot(key: str, place: str = "", word: str = "") -> dict:
    t = T.TOOLS[key]
    place = place or W.places(key)[0]
    value, start, end = t.defaults_for(place)
    return {"tool": key, "place": place, "value": value, "start": round(start * 100), "end": round(end * 100),
            "word": word}


DEFAULT_STATE = {"slots": [default_slot("layered", "up")], "rev": 0}


def tokens(prompt: str) -> list[str]:
    """CLIP's tokens of a prompt as the map shows them, without start, end and padding."""
    tok = tokenizer()
    ids = tok(prompt or "", add_special_tokens=False, truncation=True,
              max_length=tok.model_max_length - 2).input_ids
    return [tok.convert_ids_to_tokens(i).replace("</w>", "") for i in ids]


def server_part(prompt_a: str, prompt_b: str) -> dict:
    return {"prompt": prompt_a or "", "prompt_b": prompt_b or "", "tokens": tokens(prompt_a),
            "tokens_b": tokens(prompt_b), "max_len": tokenizer().model_max_length}


def map_value(server: dict, state: dict | None = None) -> str:
    out = {"server": server}
    if state is not None:
        out["state"] = state
    return json.dumps(out, ensure_ascii=False)


def parse_state(value: str | None) -> dict:
    """The map's state from its JSON value. Anything malformed counts as no tools."""
    try:
        data = json.loads(value) if value else {}
    except (TypeError, ValueError):
        data = {}
    state = data.get("state") if isinstance(data, dict) else None
    if not isinstance(state, dict):
        return {"slots": [], "rev": 0}
    slots = [s for s in state.get("slots") or [] if isinstance(s, dict)][:N_SLOTS]
    rev = state.get("rev", 0)
    return {"slots": slots, "rev": rev if isinstance(rev, int) else 0}


def number(x, default: float) -> float:
    try:
        x = float(x)
    except (TypeError, ValueError):
        return default
    return x if math.isfinite(x) else default


def tool_uses(slots: list[dict]) -> tuple[tuple, list[str]]:
    """ToolUse entries from the map's slots. Windows are clamped to 0..100 % with a note."""
    uses, notes = [], []
    for k, v in enumerate(slots):
        key = v.get("tool")
        if key not in T.TOOLS:
            continue
        t = T.TOOLS[key]
        start, end = number(v.get("start"), 0.0) / 100, number(v.get("end"), 100.0) / 100
        if t.window is not None and not (0 <= start <= 1 and 0 <= end <= 1):
            notes.append(f"Eingriff {k + 1}: das Zeitfenster geht nur von 0 % bis 100 %, "
                         "gerechnet wird mit dem Teil darin.")
        start, end = min(max(start, 0.0), 1.0), min(max(end, 0.0), 1.0)
        place = v.get("place") if v.get("place") in W.places(key) else W.places(key)[0]
        uses.append(T.ToolUse(key, number(v.get("value"), t.neutral), place, start, end,
                              str(v.get("word") or "")[:80]))
    return tuple(uses), notes


def recipe(u: T.ToolUse) -> str:
    """One line per tool, to copy and to compare runs."""
    t = T.TOOLS[u.key]
    parts = [f'"{t.label}"']
    if T.place_label(u):
        parts.append(f'ort="{T.place_label(u)}"')
    if u.key == "excise":
        parts.append(f'wort="{u.word.strip()}"')
    parts.append(f"{t.knob.split(' (')[0].lower()}={E.fmt_num(u.value, 2).rstrip('0').rstrip(',')}")
    if t.window is not None:
        parts.append(f"wirkt={round(u.start * 100)}–{round(u.end * 100)} %")
    return f"biegen({', '.join(parts)})"


# ---------------------------------------------------------------- settings and budget

def settings(prompt_a, prompt_b, seed, steps, guidance, grid, uses) -> E.Settings:
    return E.Settings(prompt_a or "", prompt_b or "", "seed_grid", seed=int(number(seed, 42)),
                      steps=int(round(number(steps, E.DEFAULT_STEPS))), guidance=number(guidance, FAM.g_default),
                      grid=int(grid or GRID_DEFAULT), te_layers=(12,), tools=uses)


def within_limits(s: E.Settings) -> tuple[E.Settings, list[str]]:
    changes, notes = {}, []
    for name, (label, lo, hi) in HARD_LIMITS.items():
        value = getattr(s, name)
        if not lo <= value <= hi:
            changes[name] = min(max(value, lo), hi)
            notes.append(f"{label}: möglich sind {lo} bis {hi}, gerechnet wird mit {changes[name]}.")
    return replace(s, **changes), notes


def work(s: E.Settings) -> float:
    """GPU work of one render in images at the default steps."""
    return (s.grid or GRID_DEFAULT) * s.steps / E.DEFAULT_STEPS


def too_much_work(s: E.Settings, renders: int) -> str | None:
    w = work(s) * renders
    if w <= MAX_WORK:
        return None
    return (f"Das ist zu viel auf einmal: Es würde {duration(w * sec_per_image())} dauern, und alle anderen "
            "müssten warten. Nehmt ein kleineres Raster oder weniger Schritte.")


def plain_key(s: E.Settings) -> E.Settings:
    """Renders without tools do not depend on prompt B."""
    return replace(s, tools=(), prompt_b="")


# ---------------------------------------------------------------- rendering

class Shot:
    """One render: image files, captions and the recorded activations."""

    def __init__(self, paths: list[str], captions: list[str], images: list, acts: dict, seconds: float,
                 warnings: list[str]):
        self.paths, self.captions, self.images, self.acts = paths, captions, images, acts
        self.seconds, self.warnings = seconds, warnings


_plain: OrderedDict = OrderedDict()
_run_counter = 0
_run_lock = threading.Lock()


def run_dir() -> str:
    """A fresh folder for one run's images. Old folders go."""
    global _run_counter
    with _run_lock:
        _run_counter += 1
        path = os.path.join(OUT_DIR, f"run-{_run_counter:05d}")
        os.makedirs(path, exist_ok=True)
        stale = sorted(d for d in os.listdir(OUT_DIR) if d.startswith("run-"))[:-KEEP_RUNS]
    keep = {os.path.dirname(p) for shot in _plain.values() for p in shot.paths[:1]}
    for d in stale:
        if os.path.join(OUT_DIR, d) not in keep:
            shutil.rmtree(os.path.join(OUT_DIR, d), ignore_errors=True)
    return path


def shoot(s: E.Settings, folder: str, tag: str, report) -> Shot:
    eng = get_engine()
    rec = M.Recorder(s.steps)
    try:
        with M.recording(eng, rec):
            res = eng.render(s, E.sample_specs(s), max_seq=77, progress=report)
    except gr.Error:
        raise
    except ValueError as exc:
        raise gr.Error(f"Diese Einstellung geht nicht: {exc}")
    items, hidden = S.filter_items(res.items)
    if hidden:
        res.warnings.append(f"Der Filter hat {hidden} von {len(items)} Bildern ausgeblendet.")
    paths = []
    for k, (im, _) in enumerate(items):
        p = os.path.join(folder, f"{tag}-{k:02d}.jpg")
        im.convert("RGB").save(p, quality=90)
        paths.append(p)
    return Shot(paths, [c for _, c in items], [im for im, _ in items], rec.tensors(), res.seconds, res.warnings)


def image_change(a, b) -> float:
    """Mean absolute pixel difference, 0 (same) to 1 (black against white)."""
    import numpy as np

    x = np.asarray(a.convert("RGB"), dtype=np.float32)
    y = np.asarray(b.convert("RGB"), dtype=np.float32)
    return float(np.abs(x - y).mean() / 255.0) if x.shape == y.shape else 1.0


def file_url(path: str) -> str:
    return "/gradio_api/file=" + path


def compare_html(plain: Shot, bent: Shot | None, changes: list[float]) -> str:
    cols = math.ceil(math.sqrt(len(plain.paths)))
    cards = []
    for k, p in enumerate(plain.paths):
        cap = html.escape(plain.captions[k])
        if bent is None:
            cards.append(f'<figure class="bc-card"><img src="{file_url(p)}" alt="{cap}" loading="lazy">'
                         f"<figcaption>{cap}</figcaption></figure>")
            continue
        d = f"Δ {round(changes[k] * 100)} %"
        cards.append(
            f'<figure class="bc-card" data-both="1"><div class="bc-frame">'
            f'<img class="bc-bent" src="{file_url(bent.paths[k])}" alt="{cap}, gebogen">'
            f'<img class="bc-plain" src="{file_url(p)}" alt="{cap}, ohne Eingriff">'
            f'<div class="bc-line"></div><span class="bc-tag bc-tag-plain">ohne</span>'
            f'<span class="bc-tag bc-tag-bent">gebogen</span></div>'
            f'<figcaption>{cap} <b title="mittlere Pixeländerung gegenüber dem Bild ohne Eingriff">{d}</b>'
            f"</figcaption></figure>")
    hint = ("Fahrt mit der Maus über ein Bild: links vom Strich seht ihr dasselbe Bild ohne Eingriff."
            if bent is not None else "Kein Eingriff aktiv: das sind die Bilder ohne Eingriff.")
    toggle = '<button type="button" class="bc-toggle">Alle ohne Eingriff zeigen</button>' if bent else ""
    return (f'<div class="bc"><div class="bc-bar"><span>{hint}</span>{toggle}</div>'
            f'<div class="bc-grid" style="grid-template-columns: repeat({cols}, minmax(0, 1fr))">'
            f'{"".join(cards)}</div></div>')


EMPTY_RESULT = ('<div class="bc bc-empty">Hier erscheinen die Bilder: jedes einmal gebogen und einmal ohne '
                "Eingriff, zum Vergleichen.</div>")


def run(prompt_a, prompt_b, seed, steps, guidance, grid, map_json, progress=gr.Progress()):
    state = parse_state(map_json)
    print(f"[run] rev {state['rev']}, tools {[v.get('tool') for v in state['slots']]}", flush=True)
    uses, notes = tool_uses(state["slots"])
    s, limit_notes = within_limits(settings(prompt_a, prompt_b, seed, steps, guidance, grid, uses))
    notes = limit_notes + notes
    live = [u for u in uses if T.active(u)]
    key = plain_key(s)
    cached = _plain.get(key)
    renders = (cached is None) + bool(live)
    refusal = too_much_work(s, renders)
    if refusal:
        gr.Warning(refusal)
        return gr.update(), gr.update(), refusal

    total = renders * s.steps * math.ceil(len(E.sample_specs(s)) / max(1, E.MAX_BATCH))
    done = [0]

    def report(i: int, n: int) -> None:
        done[0] += 1
        progress(min(done[0] / max(1, total), 1.0), desc="ohne Eingriff" if cached is None and done[0] <= total / renders
                 else "gebogen")

    folder = run_dir()
    if cached is None:
        cached = shoot(key, folder, "plain", report)
        _plain[key] = cached
        while len(_plain) > PLAIN_CACHE:
            _plain.popitem(last=False)
    else:
        _plain.move_to_end(key)
    bent = shoot(s, folder, "bent", report) if live else None

    server = server_part(prompt_a, prompt_b)
    if bent is not None:
        changes = [image_change(a, b) for a, b in zip(cached.images, bent.images)]
        eff = M.effect(cached.acts, bent.acts)
        eff.update(image=changes, thumb=file_url(bent.paths[0]), plain_thumb=file_url(cached.paths[0]),
                   rev=state["rev"])
        server["effect"] = eff
        seconds = bent.seconds + (cached.seconds if renders == 2 else 0.0)
        lines = [f"{len(bent.paths)} Bilder gebogen in {E.fmt_num(seconds, 1)} s"
                 + ("." if renders == 2 else ", die Bilder ohne Eingriff kamen aus dem Speicher.")]
        lines += notes + bent.warnings
        lines.append("```\n" + "\n".join(recipe(u) for u in live) + "\n```")
    else:
        changes = []
        server["effect"] = {"rev": state["rev"], "steps": 0, "blocks": {}, "latent": [], "impact": {},
                            "first": None, "image": [], "thumb": file_url(cached.paths[0]), "none": True}
        lines = [f"{len(cached.paths)} Bilder ohne Eingriff"
                 + (f" in {E.fmt_num(cached.seconds, 1)} s." if renders else " aus dem Speicher.")]
        lines += notes + cached.warnings
        if uses:
            lines.append("Die Werkzeuge stehen auf dem Wert, bei dem sie nichts ändern.")
    return map_value(server), compare_html(cached, bent, changes), "\n\n".join(lines)


def show_tokens(prompt_a, prompt_b):
    return map_value(server_part(prompt_a, prompt_b))


# ---------------------------------------------------------------- the app

INTRO = """# Biegen
Stable Diffusion 1.5 von innen, von oben nach unten: Text, UNet, VAE-Decoder. Wählt rechts ein Werkzeug: Auf der
Karte leuchtet auf, wo es im Netz eingreift. Klickt auf die Karte, um den Ort zu wechseln, und zieht auf der
Zeitleiste, wann es wirkt. Nach jedem Lauf zeigt die Karte, wie stark sich jeder Block gegenüber dem Lauf ohne
Eingriff verändert hat."""

LEGEND = """**So liest man die Karte.** Oben wird der Text zu 77 Vektoren (Kontext). Darunter das UNet in seiner
U-Form: der Abwärtsweg verkleinert das Bild von 64 × 64 auf 8 × 8 Stellen, der Aufwärtsweg baut es wieder auf, die
Skip-Verbindungen tragen Einzelheiten quer hinüber. In jedem Block rechnen **Res** (Faltungen, hier hört das Netz
auch den Zeitpunkt t), **Selbst** (jede Bildstelle schaut auf alle anderen) und **Quer** (jede Bildstelle schaut auf
den Kontext; nur hier kommt der Text ins Bild). Um das UNet läuft die Schleife: das UNet sagt das Rauschen voraus,
einmal mit und einmal ohne Prompt, die **Guidance** verstärkt den Unterschied, der Sampler zieht das Rauschen ab, und
das Latent geht mit dem nächsten **Zeitpunkt t** in den nächsten Schritt. Nach dem letzten Schritt übersetzt unten
der VAE-Decoder das Latent in Pixel. Ganz unten die Zeitleiste von Rauschen bis Bild.

**Was gemessen wird.** Jeder Lauf rechnet dieselben Seeds auch ohne Eingriff. Die Füllung eines Blocks zeigt, wie
stark sich sein Ausgang im ersten veränderten Schritt unterscheidet: Dort ist das Latent noch gleich, also ändern
sich nur der Ort des Eingriffs und alles danach. Die Linie im Block zeigt dieselbe Änderung über alle Schritte, die
Linie über der Zeitleiste die Änderung des Latents. Gemessen wird am ersten Bild, als relative Änderung."""


def biegen_app() -> gr.Blocks:
    js = read("bendmap.js").replace("__CATALOG__", json.dumps(catalog(), ensure_ascii=False))
    with gr.Blocks(title="Biegen") as demo:
        gr.Markdown(INTRO, elem_id="intro")
        with gr.Row():
            prompt_a = gr.Textbox(PROMPT_A, label="Prompt A", max_lines=2, elem_id="prompt-a",
                                  info="CLIP liest Englisch.")
            prompt_b = gr.Textbox(PROMPT_B, label="Prompt B", max_lines=2, elem_id="prompt-b",
                                  info="Nur für „Zweiter Prompt für die Oberfläche“.")
        bendmap = gr.HTML(map_value(server_part(PROMPT_A, PROMPT_B), DEFAULT_STATE),
                          html_template='<div class="bm-host"></div>', js_on_load=js, apply_default_css=False,
                          elem_id="bendmap")
        with gr.Row():
            seed = gr.Number(value=42, precision=0, label="Seed", scale=1, min_width=120)
            steps = free_slider(4, 50, value=E.DEFAULT_STEPS, step=1, label="Schritte", scale=2)
            guidance = free_slider(1.0, 15.0, value=FAM.g_default, step=0.1, label="Guidance", scale=2)
            grid = gr.Radio(GRID_CHOICES, value=GRID_DEFAULT, label="Raster", scale=3)
            with gr.Column(scale=1, min_width=160):
                go = gr.Button("Biegen", variant="primary", size="lg", elem_id="go")
        status = gr.Markdown(elem_id="status")
        result = gr.HTML(EMPTY_RESULT, js_on_load=read("compare.js"), apply_default_css=False, elem_id="result")
        with gr.Accordion("So liest man die Karte", open=False):
            gr.Markdown(LEGEND)

        for box in (prompt_a, prompt_b):
            box.change(show_tokens, [prompt_a, prompt_b], bendmap, queue=False, show_progress="hidden",
                       trigger_mode="always_last", api_visibility="private")
        go.click(run, [prompt_a, prompt_b, seed, steps, guidance, grid, bendmap], [bendmap, result, status],
                 api_name="biegen", show_progress_on=[result])
    return demo


def launch(**kwargs) -> gr.Blocks:
    get_engine()
    if S.SAFETY_FILTER:
        S.classifier()
    demo = biegen_app()
    demo.queue(max_size=64)
    demo.launch(theme=THEME, css=read("bendmap.css"), head=FREE_SLIDER_JS, allowed_paths=[OUT_DIR],
                auth=("workshop", WORKSHOP_PASSWORD) if WORKSHOP_PASSWORD else None,
                auth_message="Meldet euch mit dem Benutzernamen workshop an. Das Passwort steht an der Tafel.",
                **kwargs)
    return demo


if __name__ == "__main__":
    launch()