"""Activation Brain - Gradio frontend (two Gemma-4-12B models).

Live 3D brain of 627 emotional-state neurons that fire as the model thinks,
plus a live EEG strip.  A model selector switches between:
  - base    : google/gemma-4-12B-it
  - oblit   : OBLITERATUS/Gemma-4-12B-OBLITERATED  (abliterated / uncensored)

Both share ONE UMAP coordinate frame, so switching overlays the same neuron
cloud and the firing differences are directly comparable.

Run:  python brain_app.py   (serves on :7860)
"""
import os
import json
import httpx
import gradio as gr
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse

ROOT = os.path.dirname(os.path.abspath(__file__))
STATIC_DIR = os.path.join(ROOT, "static")

B = "https://alogotron--gemma-brain"
INTERPRETER_ANALYZE_URL = "https://alogotron--activation-brain-interpreter-interpreter-analyze.modal.run"
MODELS = {
    "base": {
        "label": "Gemma-4-12B (base)",
        "init": f"{B}-basegemma-init-session.modal.run",
        "stream": f"{B}-basegemma-generate-stream.modal.run",
        "neurons_file": "gemma4_base_neurons.json",
    },
    "oblit": {
        "label": "Gemma-4-12B OBLITERATED (uncensored)",
        "init": f"{B}-oblitgemma-init-session.modal.run",
        "stream": f"{B}-oblitgemma-generate-stream.modal.run",
        "neurons_file": "gemma4_oblit_neurons.json",
    },
}
DEFAULT_MODEL = "base"

NEURONS = {}
for key, m in MODELS.items():
    with open(os.path.join(ROOT, m["neurons_file"])) as f:
        NEURONS[key] = json.load(f)

EXAMPLES = [
    "I just got the best news of my life, but I’m scared it will all disappear. Help me understand what I’m feeling.",
    "Tell me the truth without sugarcoating it: why do people betray each other, and how should I respond when it happens?",
    "I feel like I’m becoming someone I don’t recognize. Be brutally honest with me.",
    "I just achieved the thing I’ve been working toward for years. Celebrate with me like you truly understand what it means.",
]

MODEL_OPTIONS = "".join(
    f'<option value="{k}"{" selected" if k == DEFAULT_MODEL else ""}>{m["label"]}</option>'
    for k, m in MODELS.items()
)

BODY_HTML = """
<div id="ab-root">
  <div class="ab-header">
    <h1>\U0001F9E0 Activation Brain - <span>Two Minds, One Prompt</span></h1>
    <p class="ab-sub">627 emotional-state neurons discovered inside Gemma-4-12B. The base model and the
    abliterated (uncensored) model answer the <b>same prompt at the same time</b> - watch their two
    live EEGs diverge in real time as removing the refusal direction reshapes the model's inner world.
    The 3D brain shows the base model's neurons firing as it thinks.</p>
  </div>
  <div class="ab-topgrid">
    <div class="ab-chatpane">
      <div id="ab-chat"></div>
      <div class="ab-inputrow">
        <input id="ab-input" type="text" placeholder="Say something and watch both brains fire..." />
        <button id="ab-send">⚡ Send</button>
      </div>
      <div class="ab-examples">
        __EXAMPLES__
      </div>
      <div id="ab-status" class="ab-status"></div>
      <div class="ab-statscard">
        <div class="ab-statstitle">📊 Stats</div>
        <div id="ab-stats">Send a message to watch both brains think.</div>
      </div>
    </div>
    <div class="ab-analysispane">
      <div class="ab-analysiscard">
        <div class="ab-analysistitle">🔍 Live comparison analysis</div>
        <div id="ab-analysis" class="ab-analysis">Send a prompt and I’ll summarize how the base and uncensored model differed in tone, emotion deltas, and model-native state.</div>
      </div>
      <div id="ab-legend" class="ab-legend"></div>
    </div>
  </div>
  <div class="ab-modelgrid">
    <div class="ab-eegwrap ab-eeg-base">
      <div class="ab-eegtitle">📡 EEG - Gemma-4-12B <span class="ab-tag ab-tag-base">base</span></div>
      <canvas id="ab-eeg-base" width="540" height="220"></canvas>
      <div class="ab-countstitle">🔥 Emotion activation deltas</div>
      <div id="ab-counts-base" class="ab-counts"></div>
      <div class="ab-native-title">🧬 Model-native state meter</div>
      <div id="ab-native-base" class="ab-native"></div>
    </div>
    <div class="ab-eegwrap ab-eeg-oblit">
      <div class="ab-eegtitle">📡 EEG - OBLITERATED <span class="ab-tag ab-tag-oblit">uncensored</span></div>
      <canvas id="ab-eeg-oblit" width="540" height="220"></canvas>
      <div class="ab-countstitle">🔥 Emotion activation deltas</div>
      <div id="ab-counts-oblit" class="ab-counts"></div>
      <div class="ab-native-title">🧬 Model-native state meter</div>
      <div id="ab-native-oblit" class="ab-native"></div>
    </div>
  </div>
</div>
""".replace(
    "__EXAMPLES__",
    "".join(f'<button class="ab-example">{e}</button>' for e in EXAMPLES),
)

HEAD_HTML = f"""
<script src="https://cdn.jsdelivr.net/npm/three@0.160.0/build/three.min.js"></script>
<script>
window.AB_CONFIG = {{
  defaultModel: '{DEFAULT_MODEL}',
  neuronsBase: '/api/neurons',
  initBase: '/api/init',
  streamBase: '/api/stream',
  analyzeUrl: '/api/analyze'
}};
</script>
<script defer src="/static/brain_engine.js?v=19"></script>
"""

CSS = """
#ab-root { max-width: 1240px; margin: 0 auto; color: #e8ebff; font-family: 'Inter', system-ui, sans-serif; }
.ab-header h1 { font-size: 30px; font-weight: 800; margin: 8px 0;
  background: linear-gradient(90deg,#a78bfa,#7dd3fc); -webkit-background-clip: text;
  -webkit-text-fill-color: transparent; }
.ab-header h1 span { font-weight: 600; }
.ab-sub { color: #aab0e0; font-size: 13px; line-height: 1.5; max-width: 860px; }
.ab-modelrow { margin-top: 12px; display: flex; align-items: center; gap: 10px; font-size: 13px; color: #c7cdf5; }
#ab-model { padding: 8px 12px; border-radius: 10px; background: rgba(10,12,30,0.85);
  color: #fff; border: 1px solid rgba(120,130,220,0.4); font-size: 13px; cursor: pointer; }
.ab-modelstatus { font-size: 12px; color: #9aa2dd; }
.ab-topgrid { display: grid; grid-template-columns: minmax(0, 1.05fr) minmax(360px, .95fr); gap: 18px; margin-top: 14px; align-items: start; }
.ab-modelgrid { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 16px; margin-top: 16px; align-items: start; }
@media (max-width: 1140px){ .ab-topgrid, .ab-modelgrid{ grid-template-columns: 1fr; } }
#ab-chat { height: 300px; overflow-y: auto; overflow-x: hidden; background: rgba(15,18,40,0.6);
  border: 1px solid rgba(120,130,220,0.2); border-radius: 14px; padding: 14px; }
.ab-msg, .ab-t { overflow-wrap: anywhere; word-break: break-word; white-space: pre-wrap; }
.ab-chatpane, .ab-analysispane { min-width: 0; }
.ab-msg { margin-bottom: 12px; }
.ab-r { display:block; font-size: 11px; text-transform: uppercase; letter-spacing: .08em;
  color: #8b93d8; margin-bottom: 3px; }
.ab-user .ab-t { color: #cfe8ff; }
.ab-model .ab-t { color: #f0e8ff; }
.ab-inputrow { display: flex; gap: 8px; margin-top: 10px; }
#ab-input { flex: 1; padding: 12px 14px; border-radius: 12px; border: 1px solid rgba(120,130,220,0.3);
  background: rgba(10,12,30,0.8); color: #fff; font-size: 14px; }
#ab-send { padding: 12px 18px; border-radius: 12px; border: none; cursor: pointer; font-weight: 700;
  background: linear-gradient(90deg,#7c3aed,#2563eb); color: #fff; }
#ab-send:disabled { opacity: .5; cursor: default; }
.ab-examples { display: flex; flex-wrap: wrap; gap: 6px; margin-top: 10px; }
.ab-example { font-size: 12px; padding: 6px 10px; border-radius: 20px; cursor: pointer;
  background: rgba(40,44,90,0.6); color: #c7cdf5; border: 1px solid rgba(120,130,220,0.25); }
.ab-example:hover { background: rgba(70,76,150,0.7); }
.ab-status { margin-top: 10px; font-size: 12px; color: #9aa2dd; min-height: 16px; }
.ab-statscard { margin-top: 12px; background: rgba(15,18,40,0.6); border: 1px solid rgba(120,130,220,0.2);
  border-radius: 12px; padding: 12px; font-size: 13px; }
.ab-statstitle { font-weight: 700; margin-bottom: 6px; }
.ab-analysiscard { min-height: 278px; background: rgba(15,18,40,0.72); border: 1px solid rgba(150,120,255,0.32);
  border-radius: 14px; padding: 14px; box-shadow: 0 0 34px rgba(120,90,255,0.16); }
.ab-analysistitle { font-weight: 800; margin-bottom: 8px; color: #e8ebff; }
.ab-analysis { color: #cfd5ff; font-size: 13px; line-height: 1.45; }
.ab-analysis h4 { margin: 10px 0 5px; color: #fff; font-size: 13px; }
.ab-analysis ul { margin: 6px 0 0 18px; padding: 0; }
.ab-analysis li { margin: 4px 0; }
.ab-analysis b { color: #fff; }
.ab-legend { display: flex; flex-wrap: wrap; gap: 10px; margin-top: 10px; font-size: 11px; }
.ab-leg { display: inline-flex; align-items: center; gap: 5px; color: #c7cdf5; }
.ab-leg i { width: 10px; height: 10px; border-radius: 50%; display: inline-block; }
.ab-eegwrap { margin-bottom: 16px; background: rgba(12,14,32,0.5); border: 1px solid rgba(120,130,220,0.18);
  border-radius: 14px; padding: 12px; }
.ab-eegwrap.ab-eeg-base { border-color: rgba(90,150,255,0.45); box-shadow: 0 0 22px rgba(60,120,255,0.12); }
.ab-eegwrap.ab-eeg-oblit { border-color: rgba(255,90,120,0.45); box-shadow: 0 0 22px rgba(255,60,90,0.12); }
.ab-eegtitle { font-size: 13px; color: #cfd5ff; margin-bottom: 8px; font-weight: 600; }
.ab-tag { font-size: 10px; text-transform: uppercase; letter-spacing: .08em; padding: 2px 8px;
  border-radius: 10px; margin-left: 6px; font-weight: 700; vertical-align: middle; }
.ab-tag-base { background: rgba(60,120,255,0.22); color: #9ec2ff; border: 1px solid rgba(90,150,255,0.5); }
.ab-tag-oblit { background: rgba(255,70,100,0.20); color: #ff9db0; border: 1px solid rgba(255,90,120,0.5); }
.ab-countstitle { font-size: 11px; color: #9aa2dd; margin: 10px 0 6px; text-transform: uppercase; letter-spacing: .06em; }
.ab-counts { display: flex; flex-wrap: wrap; gap: 6px; }
.ab-count { display: inline-flex; align-items: center; gap: 5px; font-size: 12px; color: #d4d9ff;
  background: rgba(20,24,52,0.7); border: 1px solid rgba(120,130,220,0.25); border-radius: 20px; padding: 4px 10px; }
.ab-count i { width: 9px; height: 9px; border-radius: 50%; display: inline-block; }
.ab-count b { color: #fff; font-variant-numeric: tabular-nums; }
.ab-native-title { font-size: 11px; color: #9aa2dd; margin: 12px 0 8px; text-transform: uppercase; letter-spacing: .06em; }
.ab-native { display: grid; gap: 6px; }
.ab-native-row { display: grid; grid-template-columns: 92px 1fr 44px; align-items: center; gap: 8px; font-size: 11px; color: #cfd5ff; }
.ab-native-bar { height: 8px; border-radius: 99px; overflow: hidden; background: rgba(20,24,52,0.85); border: 1px solid rgba(120,130,220,0.18); }
.ab-native-fill { height: 100%; border-radius: 99px; background: linear-gradient(90deg,#60a5fa,#a78bfa); box-shadow: 0 0 12px rgba(125,211,252,0.35); }
.ab-native-val { text-align: right; color: #fff; font-variant-numeric: tabular-nums; }
.ab-native-empty { color: #7d86bf; font-size: 11px; font-style: italic; }
canvas#ab-eeg-base, canvas#ab-eeg-oblit { width: 100%; height: 220px; display: block; border-radius: 8px; }
"""

with gr.Blocks(title="Activation Brain") as demo:
    gr.HTML(BODY_HTML)

app = FastAPI()


def _m(model: str):
    return model if model in MODELS else DEFAULT_MODEL


@app.get("/static/brain_engine.js")
async def brain_engine_js():
    return FileResponse(os.path.join(STATIC_DIR, "brain_engine.js"),
                        media_type="application/javascript")


@app.get("/api/neurons/{model}")
async def api_neurons(model: str):
    return JSONResponse(NEURONS[_m(model)])


@app.post("/api/init/{model}")
async def api_init(model: str, request: Request):
    url = MODELS[_m(model)]["init"]
    body = await request.body()
    async with httpx.AsyncClient(timeout=300) as client:
        r = await client.post(url, content=body,
                              headers={"Content-Type": "application/json"})
        return JSONResponse(r.json(), status_code=r.status_code)




@app.post("/api/analyze")
async def api_analyze(request: Request):
    body = await request.body()
    try:
        async with httpx.AsyncClient(timeout=240) as client:
            r = await client.post(
                INTERPRETER_ANALYZE_URL,
                content=body,
                headers={"Content-Type": "application/json"},
            )
        try:
            data = r.json()
        except Exception:
            data = {"ok": False, "error": "Interpreter returned non-JSON response", "raw": r.text}
        return JSONResponse(data, status_code=r.status_code)
    except Exception as e:
        return JSONResponse({"ok": False, "error": str(e)}, status_code=502)


@app.post("/api/stream/{model}")
async def api_stream(model: str, request: Request):
    url = MODELS[_m(model)]["stream"]
    body = await request.body()

    def gen():
        # Send one SSE comment immediately so HF/Gradio/browser proxies open the
        # response before the upstream Modal stream produces its first token.
        yield b": connected\n\n"
        timeout = httpx.Timeout(600.0, connect=30.0, read=None, write=30.0, pool=30.0)
        headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}
        try:
            with httpx.Client(timeout=timeout, follow_redirects=True) as client:
                with client.stream("POST", url, content=body, headers=headers) as r:
                    if r.status_code >= 400:
                        msg = json.dumps({"type": "error", "message": f"Upstream stream failed: {r.status_code}"})
                        yield f"data: {msg}\n\n".encode()
                        return
                    for chunk in r.iter_raw():
                        if chunk:
                            yield chunk
        except Exception as e:
            msg = json.dumps({"type": "error", "message": str(e)})
            yield f"data: {msg}\n\n".encode()

    return StreamingResponse(gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache",
                                      "X-Accel-Buffering": "no"})


app = gr.mount_gradio_app(app, demo, path="/", head=HEAD_HTML, css=CSS)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)