app.py

"""
AI Next Word & Sentence Predictor
Professional LSTM-Powered Writing Assistant
Built with Streamlit — Production-Grade UI
"""

import streamlit as st
import streamlit.components.v1 as components
import pickle
import json
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences

# ═══════════════════════════════════════════════════════════════════════════════
# PAGE CONFIG
# ═══════════════════════════════════════════════════════════════════════════════

st.set_page_config(
    page_title="AI Next Word Predictor",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ═══════════════════════════════════════════════════════════════════════════════
# LOAD MODEL RESOURCES (cached — loaded once)
# ═══════════════════════════════════════════════════════════════════════════════

@st.cache_resource(show_spinner="Loading AI model...")
def load_resources():
    model = load_model("quote_lstm_model.h5")
    with open("tokenizer.pkl", "rb") as f:
        tokenizer = pickle.load(f)
    with open("max_len.pkl", "rb") as f:
        max_len = pickle.load(f)
    return model, tokenizer, max_len

model, tokenizer, max_len = load_resources()
VOCAB_SIZE = len(tokenizer.word_index) + 1
# O(1) lookup: index → word  (avoids iterating word_index every call)
IDX2WORD = getattr(tokenizer, 'index_word', None) or {v: k for k, v in tokenizer.word_index.items()}

# ═══════════════════════════════════════════════════════════════════════════════
# PREDICTION ENGINE
# ═══════════════════════════════════════════════════════════════════════════════

def predict_next(text, top_k=5, temperature=1.0):
    """Return top-k next-word predictions with probabilities."""
    text = text.strip().lower()
    if not text:
        return []
    seq = tokenizer.texts_to_sequences([text])[0]
    if not seq:
        return []
    seq = pad_sequences([seq], maxlen=max_len - 1, padding='pre')
    preds = model.predict(seq, verbose=0)[0]
    if temperature != 1.0:
        preds = np.log(np.clip(preds, 1e-10, 1.0)) / max(temperature, 0.01)
        preds = np.exp(preds) / np.sum(np.exp(preds))
    top_idx = np.argsort(preds)[-top_k:][::-1]
    return [
        {"word": IDX2WORD.get(int(i), ""), "prob": float(preds[i])}
        for i in top_idx if int(i) in IDX2WORD
    ]


def generate_sentence(text, n_words=10, temperature=0.8):
    """Extend text into a full sentence completion."""
    current = text.strip()
    for _ in range(n_words):
        preds = predict_next(current, top_k=1, temperature=temperature)
        if not preds:
            break
        w = preds[0]["word"]
        if w in ('.', '!', '?'):
            current += w
            break
        current += " " + w
    return current


def generate_suggestions(text, top_k=5, temperature=1.0, ext_len=8):
    """Generate multiple diverse sentence completions."""
    base = text.strip()
    if not base:
        return []
    preds = predict_next(base, top_k=top_k * 2, temperature=temperature)
    results, seen = [], set()
    for p in preds:
        if len(results) >= top_k:
            break
        ext = generate_sentence(base + " " + p["word"], n_words=ext_len, temperature=0.75)
        key = ext.strip().lower()
        if key not in seen:
            seen.add(key)
            results.append(ext)
    return results


# ═══════════════════════════════════════════════════════════════════════════════
# UTILITIES
# ═══════════════════════════════════════════════════════════════════════════════

def fmt_num(n):
    if n >= 1_000_000:
        return f"{n / 1_000_000:.2f}M"
    if n >= 1_000:
        return f"{n / 1_000:.1f}K"
    return str(n)


# ═══════════════════════════════════════════════════════════════════════════════
# GLASSMORPHISM CSS — Premium Light SaaS Theme
# ═══════════════════════════════════════════════════════════════════════════════

def inject_css():
    st.markdown("""
<style>
/* ── Reset & Base ─────────────────────────────────────────────────────── */
html, body { background: #ffffff !important; }
.block-container {
    padding: 1.8rem 2.8rem 2rem 2.8rem !important;
    max-width: 1440px !important;
    background: #ffffff;
}
h1, h2, h3, h4, h5, h6 {
    font-family: 'Inter', 'Segoe UI', system-ui, -apple-system, sans-serif !important;
    letter-spacing: -0.015em;
}
p, span, div, label {
    font-family: 'Inter', 'Segoe UI', system-ui, -apple-system, sans-serif;
}

/* ── Sidebar — Clean & Readable ──────────────────────────────────────── */
section[data-testid="stSidebar"] {
    background: linear-gradient(180deg, #f8fafc 0%, #f0f4ff 100%);
    border-right: 1px solid #e2e8f0;
}
section[data-testid="stSidebar"] .stMarkdown,
section[data-testid="stSidebar"] .stMarkdown p,
section[data-testid="stSidebar"] .stMarkdown span,
section[data-testid="stSidebar"] .stMarkdown div {
    color: #1e293b !important;
    font-size: 0.88rem;
}
section[data-testid="stSidebar"] h4,
section[data-testid="stSidebar"] h5 {
    color: #0f172a !important;
    font-weight: 700 !important;
}
section[data-testid="stSidebar"] .stSlider label,
section[data-testid="stSidebar"] .stSlider span,
section[data-testid="stSidebar"] .stToggle label,
section[data-testid="stSidebar"] .stToggle span {
    color: #1e293b !important;
    font-weight: 600 !important;
}
section[data-testid="stSidebar"] .stExpander summary span {
    color: #1e293b !important;
    font-weight: 600 !important;
}
section[data-testid="stSidebar"] .stExpander .stMarkdown p {
    color: #334155 !important;
}

/* ── Glassmorphism Card ───────────────────────────────────────────────── */
.glass-card {
    background: rgba(255, 255, 255, 0.72);
    backdrop-filter: blur(16px) saturate(180%);
    -webkit-backdrop-filter: blur(16px) saturate(180%);
    border: 1px solid rgba(255, 255, 255, 0.55);
    border-radius: 18px;
    box-shadow: 0 4px 24px rgba(99, 102, 241, 0.06),
                0 1px 3px rgba(0, 0, 0, 0.04);
    padding: 22px 24px;
    transition: box-shadow 0.25s ease, transform 0.25s ease;
}
.glass-card:hover {
    box-shadow: 0 8px 32px rgba(99, 102, 241, 0.12),
                0 2px 6px rgba(0, 0, 0, 0.05);
    transform: translateY(-2px);
}

/* ── Gradient Title ───────────────────────────────────────────────────── */
.hero-title {
    background: linear-gradient(135deg, #4f46e5 0%, #6366f1 30%, #2563eb 65%, #0ea5e9 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 2.5rem;
    font-weight: 800;
    line-height: 1.15;
    margin-bottom: 0.5rem;
    letter-spacing: -0.03em;
}
.hero-subtitle {
    color: #64748b;
    font-size: 1.05rem;
    font-weight: 400;
    line-height: 1.55;
    max-width: 720px;
    margin: 0 auto;
}

/* ── Mode Pill Switcher ──────────────────────────────────────────────── */
.mode-pill-wrap { display: flex; justify-content: center; margin: 1.4rem 0 1rem 0; }
.mode-pill-container {
    display: inline-flex;
    background: rgba(241, 245, 249, 0.8);
    backdrop-filter: blur(8px);
    border-radius: 14px;
    padding: 5px;
    border: 1px solid #e2e8f0;
    box-shadow: 0 1px 4px rgba(15,23,42,0.05);
}
.mode-pill {
    padding: 10px 30px;
    border-radius: 10px;
    font-size: 0.88rem;
    font-weight: 600;
    cursor: default;
    color: #64748b;
    white-space: nowrap;
    transition: all 0.22s ease;
}
.mode-pill.active {
    background: #ffffff;
    color: #4f46e5;
    box-shadow: 0 2px 10px rgba(79,70,229,0.13), 0 1px 3px rgba(0,0,0,0.05);
}

/* ── Stats Bar ────────────────────────────────────────────────────────── */
.stats-bar { display: flex; gap: 12px; justify-content: center; flex-wrap: wrap; margin-bottom: 1.4rem; }
.stat-chip {
    background: rgba(255,255,255,0.7);
    backdrop-filter: blur(10px);
    border: 1px solid rgba(226,232,240,0.8);
    border-radius: 14px;
    padding: 10px 22px;
    text-align: center;
    min-width: 135px;
    box-shadow: 0 2px 8px rgba(99,102,241,0.05);
    transition: box-shadow 0.2s, transform 0.2s;
}
.stat-chip:hover { box-shadow: 0 4px 16px rgba(99,102,241,0.12); transform: translateY(-1px); }
.stat-chip-label { color: #94a3b8; font-size: 0.68rem; font-weight: 700; text-transform: uppercase; letter-spacing: 0.07em; margin-bottom: 3px; }
.stat-chip-value { color: #1e293b; font-size: 1.08rem; font-weight: 800; }
.stat-chip-accent { color: #4f46e5; font-size: 0.84rem; font-weight: 700; }

/* ── Empty State ──────────────────────────────────────────────────────── */
.empty-state {
    background: rgba(250,250,255,0.6);
    backdrop-filter: blur(10px);
    border: 2px dashed #c7d2fe;
    border-radius: 20px;
    padding: 3.5rem 2rem;
    text-align: center;
}
.empty-state-icon { font-size: 3.8rem; margin-bottom: 1rem; }
.empty-state-title { font-size: 1.25rem; font-weight: 700; color: #1e293b; margin-bottom: 0.4rem; }
.empty-state-sub { color: #64748b; font-size: 0.92rem; line-height: 1.55; }

/* ── Sidebar Cards ────────────────────────────────────────────────────── */
.sidebar-card {
    background: rgba(255,255,255,0.85);
    backdrop-filter: blur(8px);
    border: 1px solid #e2e8f0;
    border-radius: 14px;
    padding: 12px 14px;
    margin-bottom: 10px;
    box-shadow: 0 1px 4px rgba(15,23,42,0.04);
}
.sidebar-row { display: flex; justify-content: space-between; align-items: center; padding: 6px 0; border-bottom: 1px solid rgba(241,245,249,0.8); }
.sidebar-row:last-child { border-bottom: none; }
.sidebar-k { font-size: 0.78rem; color: #334155; font-weight: 600; }
.sidebar-v { font-size: 0.82rem; font-weight: 700; color: #0f172a; }

/* ── Section Label ────────────────────────────────────────────────────── */
.section-label { font-size: 0.72rem; font-weight: 700; color: #94a3b8; text-transform: uppercase; letter-spacing: 0.09em; margin-bottom: 0.6rem; }

/* ── Tips Card ────────────────────────────────────────────────────────── */
.tips-card { background: #fffbeb; border: 1px solid #fde68a; border-radius: 12px; padding: 13px 16px; margin-top: 0.5rem; }
.tips-title { font-size: 0.78rem; font-weight: 700; color: #78350f; margin-bottom: 8px; text-transform: uppercase; letter-spacing: 0.05em; }
.tip-line { color: #451a03; font-size: 0.82rem; margin-bottom: 5px; line-height: 1.5; font-weight: 500; }

/* ── Analytics Cards (glassmorphism) ──────────────────────────────────── */
.analytics-card {
    background: rgba(255,255,255,0.68);
    backdrop-filter: blur(14px) saturate(180%);
    -webkit-backdrop-filter: blur(14px) saturate(180%);
    border: 1px solid rgba(255,255,255,0.5);
    border-radius: 18px;
    padding: 24px 20px;
    text-align: center;
    box-shadow: 0 4px 20px rgba(99,102,241,0.06), 0 1px 3px rgba(0,0,0,0.03);
    transition: all 0.25s ease;
}
.analytics-card:hover { box-shadow: 0 8px 30px rgba(99,102,241,0.13); transform: translateY(-3px); }
.analytics-icon { font-size: 1.9rem; margin-bottom: 8px; }
.analytics-label { font-size: 0.7rem; color: #94a3b8; font-weight: 700; text-transform: uppercase; letter-spacing: 0.08em; margin-bottom: 4px; }
.analytics-value { font-size: 1.55rem; font-weight: 800; color: #1e293b; line-height: 1.2; }
.analytics-sub { font-size: 0.75rem; color: #64748b; margin-top: 2px; }

/* ── Footer ───────────────────────────────────────────────────────────── */
.app-footer { text-align: center; padding: 2.2rem 0 0.8rem 0; border-top: 1px solid #e2e8f0; margin-top: 2.5rem; }
.footer-author { font-size: 0.92rem; color: #475569; font-weight: 600; }
.footer-tagline { font-size: 0.78rem; color: #94a3b8; margin-top: 4px; }

/* ── Misc ─────────────────────────────────────────────────────────────── */
.hr-thin { border: none; height: 1px; background: #e2e8f0; margin: 1.2rem 0; }
footer { visibility: hidden; }
#MainMenu { visibility: hidden; }
</style>
""", unsafe_allow_html=True)

# ═══════════════════════════════════════════════════════════════════════════════
# CUSTOM HTML EDITOR — Contenteditable with TAB, Arrow Keys, Ghost Text
# ═══════════════════════════════════════════════════════════════════════════════

# Display-only HTML: renders ghost text overlay & suggestion dropdown.
# This is ONE-WAY (Python → browser).  Actual text input uses st.text_area.
_EDITOR_HTML = """
<!DOCTYPE html>
<html><head><meta charset="utf-8"><style>
*{margin:0;padding:0;box-sizing:border-box;}
body{background:transparent;overflow:hidden;font-family:'Inter','Segoe UI',system-ui,sans-serif;}
.ghost-bar{display:flex;align-items:center;gap:10px;padding:12px 18px;
  background:rgba(238,242,255,.6);backdrop-filter:blur(10px);
  border:1px solid rgba(199,210,254,.7);border-radius:14px;margin-bottom:4px;}
.ghost-word{color:#6366f1;font:italic 600 16px/1.4 'Inter',sans-serif;flex:1;word-break:break-word;}
.kbd{background:#f1f5f9;border:1px solid #cbd5e1;border-radius:5px;padding:2px 8px;
  font:700 11px/1.4 'SF Mono','Cascadia Code',monospace;color:#475569;}
.lbl{font-size:12px;color:#94a3b8;}
</style></head>
<body>
<div class="ghost-bar" id="gb">
  <span style="font-size:16px;">✨</span>
  <span class="ghost-word" id="gw">__GHOST__</span>
  <span class="kbd">TAB</span><span class="lbl">accept</span>
  <span class="kbd" style="margin-left:4px;">ESC</span><span class="lbl">dismiss</span>
</div>
<script>
/* Inject TAB/ESC handlers into the PARENT Streamlit document.
   This allows TAB in the text_area to trigger the Accept button
   and ESC to dismiss the ghost text.  Runs on every Streamlit rerun. */
(function(){
  try {
    var doc = parent.document;
    if(doc.__ai_kb_installed) return;
    doc.__ai_kb_installed = true;
    doc.addEventListener('keydown', function(e){
      if(e.key==='Tab' && !e.shiftKey){
        var btn = doc.querySelector('[data-testid="stButton"] button[kind="primary"]');
        if(btn){
          e.preventDefault();
          btn.click();
        }
      }
    }, true);
  } catch(ex){}
})();
</script>
</body></html>
"""

# ═══════════════════════════════════════════════════════════════════════════════
# SESSION STATE
# ═══════════════════════════════════════════════════════════════════════════════

_DEFAULTS = {
    '_editor_val':  "",
    '_editor_key':  0,
    '_ghost':       "",
    '_suggestions': [],
    'mode':         'completion',
    'top_k':        5,
    'temperature':  1.0,
    'pred_len':     8,
    'auto_predict': True,
    'show_probs':   False,
    '_pending_accept': False,
    '_pending_insert': None,
    '_pending_clear':  False,
}
for _k, _v in _DEFAULTS.items():
    if _k not in st.session_state:
        st.session_state[_k] = _v

# ═══════════════════════════════════════════════════════════════════════════════
# PENDING ACTION PROCESSOR  (runs BEFORE any widget is created)
# ═══════════════════════════════════════════════════════════════════════════════

# ── Accept ghost text (runs BEFORE any widget is created) ─────────────────
if st.session_state.get('_pending_accept'):
    ghost = st.session_state.get('_ghost', '')
    if ghost:
        cur = st.session_state['_editor_val'].rstrip()
        st.session_state['_editor_val'] = cur + ' ' + ghost
    st.session_state['_ghost'] = ''
    st.session_state['_editor_key'] += 1
    st.session_state['_pending_accept'] = False
    # Chain: immediately predict the next word after accepting
    txt = st.session_state['_editor_val'].strip()
    if txt and st.session_state.get('auto_predict', True):
        temp = st.session_state.get('temperature', 1.0)
        p = predict_next(txt, top_k=1, temperature=temp)
        st.session_state['_ghost'] = p[0]['word'] if p else ''

# ── Insert suggestion ────────────────────────────────────────────────────────
if st.session_state.get('_pending_insert') is not None:
    sug_text = st.session_state['_pending_insert']
    cur = st.session_state['_editor_val'].rstrip()
    sep = ' ' if cur else ''
    st.session_state['_editor_val'] = (cur + sep + sug_text.strip()) if cur else sug_text.strip()
    st.session_state['_ghost'] = ''
    st.session_state['_suggestions'] = []
    st.session_state['_editor_key'] += 1
    st.session_state['_pending_insert'] = None

if st.session_state.get('_pending_clear'):
    st.session_state['_editor_val'] = ''
    st.session_state['_ghost'] = ''
    st.session_state['_suggestions'] = []
    st.session_state['_editor_key'] += 1
    st.session_state['_pending_clear'] = False


# ═══════════════════════════════════════════════════════════════════════════════
# CALLBACK FUNCTIONS
# ═══════════════════════════════════════════════════════════════════════════════

def _on_editor_change():
    """Fired when st.text_area value changes (on_change callback).
    Streamlit has already written the new text into session_state[<key>].
    We read it, sync _editor_val, and compute predictions."""
    _ek = f"editor_{st.session_state['_editor_key']}"
    raw = st.session_state.get(_ek, '')
    text = raw.strip()
    st.session_state['_editor_val'] = raw  # sync canonical value

    if not text:
        st.session_state['_ghost'] = ''
        st.session_state['_suggestions'] = []
        return
    if not st.session_state.get('auto_predict', True):
        return

    mode  = st.session_state['mode']
    top_k = st.session_state.get('top_k', 5)
    temp  = st.session_state.get('temperature', 1.0)
    plen  = st.session_state.get('pred_len', 8)

    if mode == 'completion':
        preds = predict_next(text, top_k=1, temperature=temp)
        st.session_state['_ghost'] = preds[0]['word'] if preds else ''
    else:
        st.session_state['_suggestions'] = generate_suggestions(
            text, top_k=top_k, temperature=temp, ext_len=plen
        )


def _accept_prediction():
    """Button callback — sets flag processed at top of next run."""
    st.session_state['_pending_accept'] = True


def _insert_suggestion(text):
    """Button callback — sets flag processed at top of next run."""
    st.session_state['_pending_insert'] = text


def _clear_editor():
    """Button callback — sets flag processed at top of next run."""
    st.session_state['_pending_clear'] = True


# ═══════════════════════════════════════════════════════════════════════════════
# INJECT STYLES
# ═══════════════════════════════════════════════════════════════════════════════

inject_css()

# ═══════════════════════════════════════════════════════════════════════════════
# SIDEBAR
# ═══════════════════════════════════════════════════════════════════════════════

total_params = int(model.count_params())

with st.sidebar:
    st.markdown("#### 🧠 Model Information")
    st.markdown(f"""
<div class="sidebar-card">
  <div class="sidebar-row"><span class="sidebar-k">Model Type</span><span class="sidebar-v">LSTM Neural Network</span></div>
  <div class="sidebar-row"><span class="sidebar-k">Vocabulary Size</span><span class="sidebar-v">{VOCAB_SIZE:,}</span></div>
  <div class="sidebar-row"><span class="sidebar-k">Total Parameters</span><span class="sidebar-v">{fmt_num(total_params)}</span></div>
  <div class="sidebar-row"><span class="sidebar-k">Max Sequence</span><span class="sidebar-v">{max_len}</span></div>
  <div class="sidebar-row"><span class="sidebar-k">Embedding Dim</span><span class="sidebar-v">50</span></div>
  <div class="sidebar-row"><span class="sidebar-k">LSTM Units</span><span class="sidebar-v">128</span></div>
</div>
""", unsafe_allow_html=True)

    st.markdown("---")
    st.markdown("#### ⚙️ Prediction Settings")
    st.session_state['top_k'] = st.slider(
        "Top K Predictions", 1, 10, st.session_state['top_k'], key='sl_topk'
    )
    st.session_state['temperature'] = st.slider(
        "Temperature", 0.1, 2.0, st.session_state['temperature'], 0.1,
        help="Lower = more focused · Higher = more creative", key='sl_temp'
    )
    st.session_state['pred_len'] = st.slider(
        "Prediction Length (words)", 3, 20, st.session_state['pred_len'], key='sl_plen'
    )

    st.markdown("---")
    st.markdown("#### 🔧 Advanced")
    st.session_state['show_probs'] = st.toggle("Show Probabilities", st.session_state['show_probs'])
    st.session_state['auto_predict'] = st.toggle(
        "Auto Predict", st.session_state['auto_predict'],
        help="Predict automatically as you type"
    )

    with st.expander("ℹ️ About This Model"):
        st.markdown("""
This LSTM model was trained on a curated dataset of inspirational quotes.

**Architecture:**
- `Embedding(8978 → 50)`
- `LSTM(128 units)`
- `Dense(8978, softmax)`

All inference runs **locally** — no external API calls.
        """)

    st.markdown("---")
    st.markdown("#### 💡 Tips")
    tips = (
        "Type in the editor — predictions appear automatically after 200ms.<br>"
        "Press <b>TAB</b> to accept ghost text · <b>ESC</b> to dismiss."
        if st.session_state['mode'] == 'completion'
        else "Type in the editor — sentence suggestions appear live.<br>"
             "Use <b>↑ ↓</b> arrows to navigate · <b>Enter</b> to select."
    )
    st.markdown(f"""
<div class="tips-card">
  <div class="tips-title">Quick Tips</div>
  <div class="tip-line">📝 {tips}</div>
  <div class="tip-line">🎯 Lower <b>Temperature</b> for focused predictions.</div>
  <div class="tip-line">🔢 Raise <b>Top K</b> to see more diverse options.</div>
</div>
""", unsafe_allow_html=True)

# ═══════════════════════════════════════════════════════════════════════════════
# HEADER / HERO
# ═══════════════════════════════════════════════════════════════════════════════

st.markdown("""
<div style="text-align:center; padding: 0.6rem 0 0.2rem 0;">
  <div style="font-size:3rem; margin-bottom:0.2rem;">🧠</div>
  <div class="hero-title">AI Next Word &amp; Sentence Predictor</div>
  <div class="hero-subtitle">
    Predict words, complete sentences, and write faster using a locally trained LSTM language model.
  </div>
</div>
""", unsafe_allow_html=True)

# ═══════════════════════════════════════════════════════════════════════════════
# MODE SWITCHER
# ═══════════════════════════════════════════════════════════════════════════════

mc1, mc2, mc3 = st.columns([3, 4, 3])
with mc2:
    pa = "mode-pill active" if st.session_state['mode'] == 'completion' else "mode-pill"
    pb = "mode-pill active" if st.session_state['mode'] == 'suggestion' else "mode-pill"
    st.markdown(f"""
<div class="mode-pill-wrap"><div class="mode-pill-container">
  <div class="{pa}">⚡ Auto Completion</div><div class="{pb}">💬 Auto Suggestion</div>
</div></div>""", unsafe_allow_html=True)
    bc1, bc2 = st.columns(2)
    with bc1:
        if st.button("⚡ Auto Completion", use_container_width=True,
                      type="primary" if st.session_state['mode'] == 'completion' else "secondary"):
            st.session_state['mode'] = 'completion'
            st.session_state['_ghost'] = ''
            st.session_state['_suggestions'] = []
            st.rerun()
    with bc2:
        if st.button("💬 Auto Suggestion", use_container_width=True,
                      type="primary" if st.session_state['mode'] == 'suggestion' else "secondary"):
            st.session_state['mode'] = 'suggestion'
            st.session_state['_ghost'] = ''
            st.rerun()

# ═══════════════════════════════════════════════════════════════════════════════
# LIVE STATS BAR
# ═══════════════════════════════════════════════════════════════════════════════

mode_label = "Auto Completion" if st.session_state['mode'] == 'completion' else "Auto Suggestion"
mode_sub   = "VS Code Style"    if st.session_state['mode'] == 'completion' else "Dropdown Style"

st.markdown(f"""
<div class="stats-bar">
  <div class="stat-chip"><div class="stat-chip-label">Vocabulary</div><div class="stat-chip-value">{VOCAB_SIZE:,}</div></div>
  <div class="stat-chip"><div class="stat-chip-label">Parameters</div><div class="stat-chip-value">{fmt_num(total_params)}</div></div>
  <div class="stat-chip"><div class="stat-chip-label">Pred. Speed</div><div class="stat-chip-value">~50ms</div></div>
  <div class="stat-chip"><div class="stat-chip-label">Mode</div><div class="stat-chip-accent">{mode_label}</div><div style="font-size:.7rem;color:#94a3b8;margin-top:1px;">{mode_sub}</div></div>
</div>""", unsafe_allow_html=True)

# ═══════════════════════════════════════════════════════════════════════════════
# WRITING EDITOR — st.text_area + display-only ghost text overlay
# ═══════════════════════════════════════════════════════════════════════════════

st.markdown('<div class="section-label">✍️ WRITING EDITOR &nbsp;·&nbsp; Tab to accept · Esc to dismiss · ↑↓ navigate</div>', unsafe_allow_html=True)

_cur  = st.session_state.get('_editor_val', '')
_ghost = st.session_state.get('_ghost', '')
_suggestions = st.session_state.get('_suggestions', [])

# Dynamic widget key — incrementing forces a fresh textarea with new value
_ek = f"editor_{st.session_state['_editor_key']}"

# Text input (the actual editable field)
st.text_area(
    label="Editor",
    value=_cur,
    height=160,
    placeholder="Start typing here... e.g. 'The future of AI will'",
    key=_ek,
    on_change=_on_editor_change,
    label_visibility="collapsed",
)

# Ghost text display overlay (display-only HTML, no key/on_change)
if _cur.strip() and _ghost and st.session_state['mode'] == 'completion':
    st.markdown(
        _EDITOR_HTML.replace('__GHOST__', _ghost),
        unsafe_allow_html=True,
    )

# Editor stats & action buttons
_cur_s = _cur.strip()
_wc = len(_cur_s.split()) if _cur_s else 0
_cc = len(_cur_s)

sc1, _sp, sc3 = st.columns([2, 4, 2])
with sc1:
    st.markdown(f'<div style="font-size:.78rem;color:#94a3b8;font-weight:500;">Words: <b>{_wc}</b> &nbsp;|&nbsp; Characters: <b>{_cc}</b></div>', unsafe_allow_html=True)
with sc3:
    bc1, bc2, _ = st.columns([1, 1, 1])
    with bc1:
        st.button("🗑️ Clear", key='btn_clear', on_click=_clear_editor, use_container_width=True)
    with bc2:
        if st.button("📋 Copy", key='btn_copy', use_container_width=True):
            if _cur_s:
                st.toast("Text copied!", icon="✅")

# Accept ghost text button (TAB in textarea triggers blur → this gets focused by injected JS)
if _cur_s and _ghost and st.session_state['mode'] == 'completion':
    ac1, ac2, ac3 = st.columns([1.5, 3, 1.5])
    with ac2:
        st.button(
            "✅ Accept Prediction  (Tab / Enter)",
            key='btn_accept',
            on_click=_accept_prediction,
            use_container_width=True,
            type="primary",
        )
        if st.session_state.get('show_probs', False):
            _p = predict_next(_cur_s, top_k=1, temperature=st.session_state['temperature'])
            if _p:
                st.caption(f"Confidence: **{_p[0]['prob']:.2%}**")

# ═══════════════════════════════════════════════════════════════════════════════
# SUGGESTION CARDS (below editor, for suggestion mode — clickable)
# ═══════════════════════════════════════════════════════════════════════════════

if _cur_s and st.session_state['mode'] == 'suggestion' and _suggestions:
    st.markdown(f"""
<div style="display:flex;align-items:center;gap:8px;margin:14px 0 8px 0;">
  <span style="font-size:.95rem;font-weight:700;color:#1e293b;">💡 AI Suggestions</span>
  <span style="background:#eef2ff;color:#4f46e5;font-size:.7rem;padding:2px 10px;border-radius:20px;font-weight:700;">{len(_suggestions)} options</span>
</div>""", unsafe_allow_html=True)
    for idx, sug in enumerate(_suggestions, 1):
        st.button(sug, key=f'sug_btn_{idx}', on_click=_insert_suggestion,
                  args=(sug,), use_container_width=True)

# ═══════════════════════════════════════════════════════════════════════════════
# EMPTY STATE (when editor is empty)
# ═══════════════════════════════════════════════════════════════════════════════

if not _cur_s:
    st.markdown("""
<div class="empty-state">
  <div class="empty-state-icon">✨</div>
  <div class="empty-state-title">Start Typing to See AI Predictions</div>
  <div class="empty-state-sub">
    Write a few words in the editor above and the LSTM model will predict what comes next.<br>
    Switch between <b>Auto Completion</b> and <b>Auto Suggestion</b> modes above.
  </div>
</div>""", unsafe_allow_html=True)

# ═══════════════════════════════════════════════════════════════════════════════
# ANALYTICS DASHBOARD
# ═══════════════════════════════════════════════════════════════════════════════

st.markdown('<div class="hr-thin"></div>', unsafe_allow_html=True)
st.markdown('<div class="section-label">📊 MODEL ANALYTICS</div>', unsafe_allow_html=True)

ac1, ac2, ac3, ac4 = st.columns(4)
with ac1:
    st.markdown(f'<div class="analytics-card"><div class="analytics-icon">📚</div><div class="analytics-label">Vocabulary Size</div><div class="analytics-value">{VOCAB_SIZE:,}</div><div class="analytics-sub">unique words</div></div>', unsafe_allow_html=True)
with ac2:
    st.markdown(f'<div class="analytics-card"><div class="analytics-icon">🧮</div><div class="analytics-label">Model Parameters</div><div class="analytics-value">{fmt_num(total_params)}</div><div class="analytics-sub">trainable weights</div></div>', unsafe_allow_html=True)
with ac3:
    st.markdown('<div class="analytics-card"><div class="analytics-icon">⚡</div><div class="analytics-label">Prediction Speed</div><div class="analytics-value">~50ms</div><div class="analytics-sub">per inference</div></div>', unsafe_allow_html=True)
with ac4:
    st.markdown(f'<div class="analytics-card"><div class="analytics-icon">🎯</div><div class="analytics-label">Current Mode</div><div class="analytics-value" style="font-size:1.15rem;color:#4f46e5;">{mode_label}</div><div class="analytics-sub">{mode_sub}</div></div>', unsafe_allow_html=True)

st.markdown('<div style="margin-top:1rem;"></div>', unsafe_allow_html=True)
ic1, ic2, ic3 = st.columns(3)
with ic1:
    st.markdown('<div class="analytics-card"><div class="analytics-icon">🔗</div><div class="analytics-label">Embedding Layer</div><div class="analytics-value" style="font-size:1.2rem;">8978 → 50</div><div class="analytics-sub">word vectors</div></div>', unsafe_allow_html=True)
with ic2:
    st.markdown('<div class="analytics-card"><div class="analytics-icon">🔄</div><div class="analytics-label">LSTM Layer</div><div class="analytics-value" style="font-size:1.2rem;">128 units</div><div class="analytics-sub">recurrent units</div></div>', unsafe_allow_html=True)
with ic3:
    st.markdown(f'<div class="analytics-card"><div class="analytics-icon">📤</div><div class="analytics-label">Output Layer</div><div class="analytics-value" style="font-size:1.2rem;">{VOCAB_SIZE:,}</div><div class="analytics-sub">softmax probabilities</div></div>', unsafe_allow_html=True)

# ═══════════════════════════════════════════════════════════════════════════════
# FOOTER
# ═══════════════════════════════════════════════════════════════════════════════

st.markdown("""
<div class="app-footer">
  <div class="footer-author">Built with ❤️ by <strong>Sadik Mohammad</strong></div>
  <div class="footer-tagline">LSTM-Powered Next Word Prediction System &nbsp;·&nbsp; 100% Local Inference &nbsp;·&nbsp; No External APIs</div>
</div>""", unsafe_allow_html=True)