app.py

import gradio as gr
import numpy as np
import re
from collections import defaultdict

# ============================================================
# TRIBE Paper Reader - Cognitive State Highlighter
# ============================================================
# Maps predicted brain activation patterns to cognitive/emotional
# states and highlights research paper text accordingly.
#
# Color Mapping:
#   Red    = Curiosity    (prefrontal + anterior cingulate)
#   Blue   = Questioning  (dorsolateral prefrontal dominance)
#   Green  = Confusion    (high executive load, low coherence)
#   Yellow = Excitement   (emotion + reward circuits)
#   Purple = Insight      (associative cortex coherence)
#   Orange = Frustration  (conflict monitoring, low reward)
# ============================================================

# ---- Cognitive State Definitions ----
STATE_COLORS = {
    "curiosity":    {"hex": "#ff4444", "bg": "rgba(255,68,68,0.15)",  "label": "Curiosity",    "emoji": "🔴"},
    "questioning":  {"hex": "#4488ff", "bg": "rgba(68,136,255,0.15)", "label": "Questioning",  "emoji": "🔵"},
    "confusion":    {"hex": "#44aa44", "bg": "rgba(68,170,68,0.15)",  "label": "Confusion",    "emoji": "🟢"},
    "excitement":   {"hex": "#ffaa00", "bg": "rgba(255,170,0,0.15)",  "label": "Excitement",   "emoji": "🟡"},
    "insight":      {"hex": "#aa44ff", "bg": "rgba(170,68,255,0.15)", "label": "Insight",      "emoji": "🟣"},
    "frustration":  {"hex": "#ff8844", "bg": "rgba(255,136,68,0.15)", "label": "Frustration",  "emoji": "🟠"},
    "neutral":      {"hex": "#888888", "bg": "transparent",          "label": "Neutral",      "emoji": "⚪"},
}

# ---- Text Analysis Heuristics (fallback when TRIBE not available) ----
CURIOSITY_MARKERS = [
    "novel", "new", "previously unknown", "first time", "surprisingly",
    "unexpected", "interesting", "intriguing", "remarkable", "striking",
    "counterintuitive", "paradox", "mystery", "unknown", "unexplored",
    "hypothesis", "propose", "suggest", "may", "might", "could", "possibly",
    "?", "what if", "how does", "why do", "future work", "open question"
]

QUESTIONING_MARKERS = [
    "however", "but", "although", "nevertheless", "yet", "despite",
    "in contrast", "on the other hand", "contradict", "inconsistent",
    "unclear", "debated", "controversial", "challenge", "criticism",
    "limitation", "weakness", "fails to", "does not", "not significantly",
    "p > 0.05", "insufficient", "inconclusive"
]

CONFUSION_MARKERS = [
    "complex", "complicated", "difficult", "challenging", "intricate",
    "non-trivial", "subtle", "nuanced", "ambiguous", "uncertain",
    "unclear", "not well understood", "remains elusive", "poorly understood"
]

EXCITEMENT_MARKERS = [
    "breakthrough", "significant", "substantial", "dramatic", "remarkable",
    "outstanding", "excellent", "superior", "state-of-the-art", "sota",
    "achieves", "improves", "outperforms", "best", "top", "first place",
    "novel contribution", "key result", "important", "crucial"
]

INSIGHT_MARKERS = [
    "thus", "therefore", "hence", "consequently", "as a result",
    "we find that", "our results show", "demonstrate", "reveal",
    "discover", "found that", "this suggests", "this implies",
    "in summary", "to conclude", "overall", "taken together"
]

FRUSTRATION_MARKERS = [
    "failed", "failure", "error", "problem", "issue", "difficulty",
    "bottleneck", "degradation", "drops", "worse", "inferior",
    "not achieve", "unable to", "struggle", "limited by", "constrained"
]


def detect_state_heuristic(sentence):
    """Fallback heuristic-based state detection when TRIBE model unavailable."""
    s_lower = sentence.lower()
    scores = defaultdict(float)

    for marker in CURIOSITY_MARKERS:
        if marker.lower() in s_lower:
            scores["curiosity"] += 1.0
    for marker in QUESTIONING_MARKERS:
        if marker.lower() in s_lower:
            scores["questioning"] += 1.5
    for marker in CONFUSION_MARKERS:
        if marker.lower() in s_lower:
            scores["confusion"] += 1.2
    for marker in EXCITEMENT_MARKERS:
        if marker.lower() in s_lower:
            scores["excitement"] += 1.3
    for marker in INSIGHT_MARKERS:
        if marker.lower() in s_lower:
            scores["insight"] += 1.4
    for marker in FRUSTRATION_MARKERS:
        if marker.lower() in s_lower:
            scores["frustration"] += 1.1

    # Sentence structure heuristics
    if s_lower.endswith("?"):
        scores["questioning"] += 2.0
    if "we hypothesize" in s_lower or "we propose" in s_lower:
        scores["curiosity"] += 1.5
    if "p < 0.001" in s_lower or "p < 0.01" in s_lower:
        scores["excitement"] += 1.0
    if "however" in s_lower or "but" in s_lower:
        scores["questioning"] += 0.8

    if not scores:
        return "neutral", 0.0

    best_state = max(scores, key=scores.get)
    return best_state, scores[best_state]


def split_into_sentences(text):
    """Split text into sentences, preserving structure."""
    # Handle common abbreviations
    text = re.sub(r'(?<=[A-Z])\.(?=[A-Z])', '.', text)
    # Split on sentence boundaries
    sentences = re.split(r'(?<=[.!?])\s+', text)
    return [s.strip() for s in sentences if s.strip()]


def highlight_text(text, use_tribe=False):
    """Analyze text and return HTML with cognitive state highlights."""
    sentences = split_into_sentences(text)
    highlighted = []
    state_counts = defaultdict(int)

    for sent in sentences:
        if not sent:
            continue

        if use_tribe:
            # TODO: Integrate TRIBE v2 prediction here
            # For now, use heuristics
            state, confidence = detect_state_heuristic(sent)
        else:
            state, confidence = detect_state_heuristic(sent)

        state_counts[state] += 1
        color_info = STATE_COLORS.get(state, STATE_COLORS["neutral"])

        # Build highlighted sentence
        tooltip = f"{color_info['emoji']} {color_info['label']} (confidence: {confidence:.1f})"
        html_sent = (
            f'<span style="background-color: {color_info["bg"]}; '
            f'border-bottom: 2px solid {color_info["hex"]}; '
            f'padding: 1px 2px; border-radius: 2px; '
            f'title="{tooltip}">'
            f'{sent}</span>'
        )
        highlighted.append(html_sent)

    # Join with spaces
    full_html = " ".join(highlighted)

    # Build legend
    legend_html = '<div style="margin: 15px 0; padding: 10px; background: #1a1a1a; border-radius: 8px;">'
    legend_html += '<h4 style="color: #fff; margin: 0 0 10px 0;">🧠 Cognitive State Legend</h4>'
    legend_html += '<div style="display: flex; flex-wrap: wrap; gap: 10px;">'
    for state, info in STATE_COLORS.items():
        if state == "neutral":
            continue
        count = state_counts.get(state, 0)
        if count > 0:
            legend_html += (
                f'<span style="background: {info["bg"]}; color: {info["hex"]}; '
                f'padding: 4px 10px; border-radius: 12px; font-size: 0.85em; '
                f'border: 1px solid {info["hex"]};">'
                f'{info["emoji"]} {info["label"]} ({count})</span>'
            )
    legend_html += '</div></div>'

    # Build stats
    total = sum(state_counts.values())
    stats_html = '<div style="margin: 15px 0; padding: 10px; background: #1a1a1a; border-radius: 8px;">'
    stats_html += '<h4 style="color: #fff; margin: 0 0 10px 0;">📊 Reading Profile</h4>'
    stats_html += '<div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 8px;">'
    for state, info in STATE_COLORS.items():
        if state == "neutral":
            continue
        count = state_counts.get(state, 0)
        pct = (count / total * 100) if total > 0 else 0
        bar_width = int(pct * 2)
        stats_html += (
            f'<div style="color: #ccc; font-size: 0.8em;">'
            f'<div style="display: flex; justify-content: space-between;">'
            f'<span>{info["emoji"]} {info["label"]}</span>'
            f'<span>{pct:.0f}%</span></div>'
            f'<div style="background: #333; height: 6px; border-radius: 3px; margin-top: 2px;">'
            f'<div style="background: {info["hex"]}; width: {bar_width}px; height: 100%; border-radius: 3px;">'
            f'</div></div></div>'
        )
    stats_html += '</div></div>'

    return legend_html + stats_html + '<div style="line-height: 1.8; color: #e0e0e0; font-size: 1.05em;">' + full_html + '</div>'


# ---- Gradio UI ----
with gr.Blocks(title="TRIBE Paper Reader", css="""
    body { background-color: #0d0d0d !important; }
    .gradio-container { max-width: 1100px !important; }
""") as demo:
    gr.HTML("""
    <div style="text-align: center; padding: 20px 0; border-bottom: 1px solid #333; margin-bottom: 20px;">
        <h1 style="color: #fff; font-size: 2.2rem; margin: 0;">🧠 TRIBE Paper Reader</h1>
        <p style="color: #888; font-size: 1rem; margin: 8px 0 0 0;">
            Read research papers with cognitive state highlighting powered by brain encoding
        </p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📄 Input")
            input_text = gr.TextArea(
                label="Paste paper text",
                placeholder="Paste abstract, introduction, or full paper text here...",
                lines=20,
                value="""We introduce a novel architecture for visual recognition that achieves state-of-the-art results on ImageNet. However, our method requires significantly more compute than prior approaches. The key insight is that multi-scale feature aggregation improves representation quality, but this comes at the cost of increased model complexity. We hypothesize that future work could reduce this overhead through better optimization. Our experiments demonstrate a 15% improvement over the previous best, but we acknowledge limitations in our evaluation protocol."""
            )
            analyze_btn = gr.Button("🧠 Analyze Paper", variant="primary", size="lg")

        with gr.Column(scale=1):
            gr.Markdown("### ✨ Highlighted Output")
            output_html = gr.HTML(label="Highlighted Text")

    gr.Markdown("""
    <div style="margin-top: 30px; padding: 15px; background: #1a1a1a; border-radius: 8px; color: #888; font-size: 0.85em;">
        <strong>How it works:</strong> Each sentence is analyzed for cognitive markers.
        In the full TRIBE integration, Meta's TRIBE v2 brain encoding model predicts
        whole-brain fMRI responses to infer the reader's likely cognitive state.
        <br><br>
        <strong>Colors:</strong>
        <span style="color:#ff4444">🔴 Curiosity</span> |
        <span style="color:#4488ff">🔵 Questioning</span> |
        <span style="color:#44aa44">🟢 Confusion</span> |
        <span style="color:#ffaa00">🟡 Excitement</span> |
        <span style="color:#aa44ff">🟣 Insight</span> |
        <span style="color:#ff8844">🟠 Frustration</span>
    </div>
    """)

    analyze_btn.click(
        fn=highlight_text,
        inputs=input_text,
        outputs=output_html
    )

if __name__ == "__main__":
    demo.launch()