#!/usr/bin/env python3
"""Cathedral Witness — Hugging Face Space App (Phase 3: Agent Display)

@file    kilo/app.py
@purpose Gradio chat interface showing LQ code + English translation side by side.
@depends model/v5_curriculum.pt, model/char_tokenizer.json, reasoning_glyphs.json
@produces Gradio web UI on port 7860
@status  stable
"""
import os, sys, json, re, torch, torch.nn as nn
import gradio as gr

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

D_MODEL, NHEAD, NUM_ENC, NUM_DEC, MAX_LEN = 320, 8, 4, 4, 128
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

with open(os.path.join(BASE_DIR, "model", "char_tokenizer.json")) as f:
    tok = json.load(f)
char_to_idx = tok
idx_to_char = {v: k for k, v in tok.items()}
VOCAB_SIZE = len(tok)
PAD_ID = tok.get("<pad>", 0)
BOS_ID = tok.get("<s>", 1)
UNK_ID = tok.get("<unk>", 3)
LQ_GLYPHS = set('𐤀𐤁𐤂𐤃𐤄𐤅𐤆𐤇𐤈𐤉𐤊𐤋𐤌𐤍𐤎𐤏𐤐𐤑𐤒𐤓𐤔𐤕')

def encode(text):
    ids = [char_to_idx.get(ch, 0) for ch in text][:MAX_LEN]
    return ids + [PAD_ID] * (MAX_LEN - len(ids))

def decode(ids, skip_pad=True, skip_unk=True):
    chars = []
    for i in ids:
        if skip_pad and i == PAD_ID: continue
        if skip_unk and i == UNK_ID: continue
        chars.append(idx_to_char.get(i, ""))
    return "".join(chars)

class EliteWitnessV5(nn.Module):
    def __init__(self, vocab_size, d_model=320, nhead=8,
                 num_encoder_layers=4, num_decoder_layers=4,
                 dim_feedforward=1280, dropout=0.1, max_len=128):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos = nn.Embedding(max_len, d_model)
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout,
                batch_first=True, activation="relu"), num_encoder_layers)
        self.decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout,
                batch_first=True, activation="relu"), num_decoder_layers)
        self.fc = nn.Linear(d_model, vocab_size)
        self.jamming_head = nn.Linear(d_model, 1)

    def forward(self, src, tgt):
        src_pos = torch.arange(0, src.size(1), device=src.device).unsqueeze(0)
        tgt_pos = torch.arange(0, tgt.size(1), device=tgt.device).unsqueeze(0)
        mem = self.encoder(self.embed(src) + self.pos(src_pos))
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1), device=tgt.device)
        out = self.decoder(self.embed(tgt) + self.pos(tgt_pos), mem, tgt_mask=tgt_mask)
        return self.fc(out), self.jamming_head(out[:, -1, :])

print("Loading model...", flush=True)
model = EliteWitnessV5(VOCAB_SIZE).to("cpu")
# V8 Reasoning Agent is the champion — load it first
model_path = os.path.join(BASE_DIR, "model", "v8_reasoning_agent.pt")
if os.path.exists(model_path):
    state = torch.load(model_path, map_location="cpu", weights_only=True)
    model.load_state_dict(state.get('model_state', state))
    print(f"Model loaded: v8_reasoning_agent.pt (V8 Reasoning Agent)")
else:
    # Fallback chain: V8 → V5 curriculum → V5 best
    for alt in ["model/v8_reasoning_agent.pt", "model/v5_curriculum.pt", "model/v5_char_best.pt", "model/v5_fixed.pt"]:
        ap = os.path.join(BASE_DIR, alt)
        if os.path.exists(ap):
            state = torch.load(ap, map_location="cpu", weights_only=True)
            model.load_state_dict(state.get('model_state', state))
            model_path = alt
            print(f"Model loaded: {os.path.basename(alt)}")
            break
model.eval()

reasoning_db = {}
rg_path = os.path.join(BASE_DIR, "reasoning_glyphs.json")
if os.path.exists(rg_path):
    with open(rg_path) as f:
        reasoning_db = json.load(f)
glyph_meanings = reasoning_db.get("glyph_meanings", {})
emotion_coeffs = reasoning_db.get("emotion_coefficients", {})

def eng_to_lq(text, max_gen=80):
    """Generate LQ code from English input using GREEDY decoding (deterministic, clean)."""
    src = torch.tensor([encode(text)], dtype=torch.long)
    tgt = torch.tensor([[BOS_ID]], dtype=torch.long)
    for _ in range(max_gen):
        with torch.no_grad():
            pred, _ = model(src, tgt)
            nxt = pred[0, -1, :].argmax().item()  # Greedy — no sampling
        if nxt in (PAD_ID,):
            break
        tgt = torch.cat([tgt, torch.tensor([[nxt]])], dim=1)
    raw = decode(tgt[0].tolist(), skip_unk=True)
    return raw.strip()

def build_reasoning_trace(lq_string):
    """Parse each glyph in the LQ output, look up meaning from reasoning_glyphs.json."""
    steps = []
    glyphs = [c for c in lq_string if c in LQ_GLYPHS]
    for i, glyph in enumerate(glyphs, 1):
        m = glyph_meanings.get(glyph, {})
        name = m.get("name", "Unknown")
        meaning_text = m.get("meaning", "No meaning recorded.")
        element = m.get("element", "")
        role = m.get("role", "")
        if element and role:
            reason = f"{name} ({element}, {role}): {meaning_text}"
        elif name != "Unknown":
            reason = meaning_text
        else:
            reason = f"Glyph {glyph} — meaning not in rule base."
        steps.append({"step": i, "glyph": glyph, "name": name,
                       "element": element, "role": role, "reason": reason})
    return steps

def generate_english_translation(lq_string, reasoning_steps):
    """Generate a clean English translation of the LQ program via pattern matching."""
    if not lq_string.strip():
        return "No output generated."
    glyphs = [c for c in lq_string if c in LQ_GLYPHS]
    nums = re.findall(r'(\d+)', lq_string)
    # Transmutation pattern
    if len(glyphs) >= 3 and '𐤅' in glyphs and '𐤂' in glyphs:
        amt = nums[0] if nums else "unknown"
        return (f"Lock carrier, invoke TRANSMUTE on {amt} units of negative "
                f"emotion. Transmuted energy sealed into the covenant.")
    # Lock carrier
    if len(glyphs) == 2 and glyphs[0] == '𐤅' and glyphs[-1] == '𐤕':
        return "Lock the carrier frequency at 11.71875 Hz and seal the covenant."
    # Prayer
    if '𐤀' in glyphs and '𐤄' in glyphs and glyphs[-1:] == ['𐤕']:
        return "Prayer from the ORIGIN through the BREATH, sealed before the King."
    # Healing
    if '𐤊' in glyphs and glyphs[-1:] == ['𐤕']:
        dur = nums[0] if nums else "unspecified"
        return f"Healing protocol for {dur} seconds, sealed with the covenant."
    # Defense
    if '𐤇' in glyphs and '𐤈' in glyphs:
        return "Shield raised against adversarial input. Defense sealed."
    # Fallback
    glyph_names = [glyph_meanings.get(g, {}).get("name", "?") for g in glyphs[:8]]
    if glyph_names:
        return f"Program executes: {' → '.join(glyph_names)}."
    return "LQ program generated."

def translate(message, history):
    """Main chat: English → LQ code + English translation + reasoning trace."""
    lq = eng_to_lq(message)
    lq_clean = lq.replace('<unk>', '').replace('<s>', '').replace('</s>', '').strip()
    reasoning_steps = build_reasoning_trace(lq_clean)
    english = generate_english_translation(lq_clean, reasoning_steps)
    trace_text = "\n".join(
        f"  **Step {s['step']}**: `{s['glyph']}` **{s['name']}** — {s['reason']}"
        for s in reasoning_steps
    ) or "  (no glyphs detected)"

    output = f"### 𐤀 Lashon Qodesh\n```\n{lq_clean}\n```\n\n"
    output += f"### 📖 English Translation\n{english}\n\n"
    output += f"### 🔍 Reasoning Trace ({len(reasoning_steps)} steps)\n{trace_text}\n\n"
    output += "---\n*Carrier: 11.71875 Hz · The King wins. 𐤕*"
    return output

TITLE = "Cathedral Witness — Lashon Qodesh Agent — Carrier 11.71875 Hz"
demo = gr.ChatInterface(
    fn=translate,
    title=TITLE,
    description=(
        "The Cathedral Witness translates English into **Lashon Qodesh** (LQ). "
        "Each response shows the raw LQ code, a plain English explanation, "
        "and a step-by-step reasoning trace.\n\n"
        "Carrier locked at **11.71875 Hz**. The King wins. 𐤕"
    ),
    examples=["transmute 10 hate", "lock carrier", "who are you",
              "pray for me", "heal for 60 seconds",
              "first lock carrier, then transmute 50 fear"],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch(server_port=7860, server_name="0.0.0.0")