app.py

import sys
import os
import types

# Python 3.13 audioop stub (not needed on 3.11 but harmless)
if 'audioop' not in sys.modules:
    sys.modules['audioop'] = types.ModuleType('audioop')

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')

model = None

def load_model():
    global model
    if model is not None:
        return "✅ Already loaded!"
    try:
        from tribev2 import TribeModel
        model = TribeModel.from_pretrained("facebook/tribev2", cache_folder="./tribe_cache")
        return "✅ Model loaded!"
    except Exception as e:
        return f"❌ Error: {str(e)}"

REGIONS = [
    ("Visual cortex",           0.00, 0.15, "#378ADD"),
    ("Auditory cortex",         0.15, 0.30, "#D85A30"),
    ("Language (Broca's area)", 0.30, 0.45, "#7F77DD"),
    ("Prefrontal (attention)",  0.45, 0.62, "#1D9E75"),
    ("Temporal (memory)",       0.62, 0.78, "#BA7517"),
    ("Emotion (limbic)",        0.78, 1.00, "#D4537E"),
]

def score_predictions(preds):
    avg = np.mean(np.abs(preds), axis=0)
    global_max = avg.max() + 1e-8
    half = len(avg) // 2
    scores = {}
    for name, s, e, _ in REGIONS:
        start, end = int(half * s), int(half * e)
        scores[name] = round(float(np.mean(avg[start:end]) / global_max * 100), 1)
    return scores, round(sum(scores.values()) / len(scores), 1)

def make_brain_plot(preds):
    try:
        from nilearn import plotting, datasets
        avg = np.mean(np.abs(preds), axis=0)
        avg_norm = (avg - avg.min()) / (avg.max() - avg.min() + 1e-8)
        half = len(avg_norm) // 2
        fsaverage = datasets.fetch_surf_fsaverage("fsaverage5")
        fig, axes = plt.subplots(1, 2, figsize=(14, 5), subplot_kw={"projection": "3d"})
        fig.patch.set_facecolor("#111111")
        plotting.plot_surf_stat_map(fsaverage.infl_left, avg_norm[:half], hemi="left",
            view="lateral", colorbar=True, cmap="hot", title="Left hemisphere", axes=axes[0], figure=fig)
        plotting.plot_surf_stat_map(fsaverage.infl_right, avg_norm[half:], hemi="right",
            view="lateral", colorbar=True, cmap="hot", title="Right hemisphere", axes=axes[1], figure=fig)
        plt.tight_layout()
        plt.savefig("/tmp/brain_map.png", dpi=130, bbox_inches="tight", facecolor="#111111")
        plt.close()
        return "/tmp/brain_map.png"
    except Exception as e:
        print(f"Brain plot error: {e}")
        return None

def make_score_chart(scores, overall):
    fig, ax = plt.subplots(figsize=(9, 4))
    fig.patch.set_facecolor("#1a1a1a")
    ax.set_facecolor("#1a1a1a")
    names  = [r[0] for r in REGIONS]
    colors = [r[3] for r in REGIONS]
    vals   = [scores.get(n, 0) for n in names]
    bars = ax.barh(names, vals, color=colors, height=0.55)
    ax.set_xlim(0, 100)
    ax.axvline(70, color="#888", linestyle="--", linewidth=1, alpha=0.6)
    ax.set_xlabel("Activation score", color="#ccc", fontsize=11)
    ax.set_title(f"Brain region activation  |  Overall: {overall}/100",
                 color="white", fontsize=13, fontweight="bold", pad=12)
    ax.tick_params(colors="#ccc")
    for spine in ax.spines.values():
        spine.set_edgecolor("#333")
    for bar, val in zip(bars, vals):
        ax.text(bar.get_width() + 1, bar.get_y() + bar.get_height() / 2,
                f"{val}", va="center", color="white", fontsize=10, fontweight="bold")
    plt.tight_layout()
    plt.savefig("/tmp/score_chart.png", dpi=130, bbox_inches="tight", facecolor="#1a1a1a")
    plt.close()
    return "/tmp/score_chart.png"

def generate_suggestions(scores, overall):
    tips = []
    if scores.get("Prefrontal (attention)", 100) < 70:
        tips.append("→ Open with a bold question or surprising fact to boost attention")
    if scores.get("Emotion (limbic)", 100) < 70:
        tips.append("→ Add emotional language — 'imagine', 'feel', personal stories")
    if scores.get("Temporal (memory)", 100) < 70:
        tips.append("→ Include specific numbers or data points to improve memorability")
    if scores.get("Visual cortex", 100) < 70:
        tips.append("→ Use more visual language — describe what viewers will 'see'")
    if scores.get("Language (Broca's area)", 100) < 70:
        tips.append("→ Break long sentences into shorter, punchier ones")
    if scores.get("Auditory cortex", 100) < 70:
        tips.append("→ Add rhythm and repetition — the brain responds to sound patterns")
    if not tips:
        tips.append("→ Excellent! Consider adding a strong call-to-action at the end")
    status = "🟢 Strong" if overall >= 75 else "🟡 Good, needs polish" if overall >= 55 else "🔴 Needs work"
    return f"**Overall: {overall}/100 — {status}**\n\n" + "\n".join(tips)

def analyze_script(script_text, progress=gr.Progress()):
    if not script_text or not script_text.strip():
        return None, None, "⚠️ Please paste a script first.", None
    if model is None:
        progress(0.1, desc="Loading TRIBE v2 model (first time ~5 mins)...")
        msg = load_model()
        if "Error" in msg:
            return None, None, msg, None
    try:
        from gtts import gTTS
        progress(0.2, desc="Converting script to speech...")
        tts = gTTS(text=script_text.strip(), lang="en", slow=False)
        tts.save("/tmp/script_audio.mp3")
        progress(0.4, desc="Running TRIBE v2 prediction (1-3 mins)...")
        df = model.get_events_dataframe(audio_path="/tmp/script_audio.mp3")
        preds, segments = model.predict(events=df)
        progress(0.7, desc="Scoring regions...")
        scores, overall = score_predictions(preds)
        progress(0.8, desc="Rendering maps...")
        brain_img   = make_brain_plot(preds)
        score_img   = make_score_chart(scores, overall)
        suggestions = generate_suggestions(scores, overall)
        np.save("/tmp/brain_predictions.npy", preds)
        progress(1.0, desc="Done!")
        return brain_img, score_img, suggestions, "/tmp/brain_predictions.npy"
    except Exception as e:
        return None, None, f"❌ Error:\n{str(e)}", None

css = "#title{text-align:center} #subtitle{text-align:center;color:#888;font-size:14px}"

with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo"), css=css) as demo:
    gr.Markdown("# 🧠 Script Brain Optimizer", elem_id="title")
    gr.Markdown("Paste your script → real fMRI predictions via **TRIBE v2** → iterate", elem_id="subtitle")
    with gr.Row():
        with gr.Column(scale=1):
            script_input = gr.Textbox(label="Your script",
                placeholder="Paste your content script here...", lines=12, max_lines=20)
            with gr.Row():
                clear_btn   = gr.Button("Clear", variant="secondary", scale=1)
                analyze_btn = gr.Button("🧠 Analyze", variant="primary", scale=3)
            suggestions_out = gr.Markdown(value="*Paste a script and click Analyze...*")
            download_out    = gr.File(label="Download predictions (.npy)")
        with gr.Column(scale=2):
            brain_img_out = gr.Image(label="Brain activation map", height=320)
            score_img_out = gr.Image(label="Region scores", height=280)

    analyze_btn.click(fn=analyze_script, inputs=[script_input],
                      outputs=[brain_img_out, score_img_out, suggestions_out, download_out])
    clear_btn.click(fn=lambda: ("", None, None, "*Paste a script and click Analyze...*", None),
                    outputs=[script_input, brain_img_out, score_img_out, suggestions_out, download_out])
    gr.Markdown("---\n*Powered by [TRIBE v2](https://github.com/facebookresearch/tribev2) by Meta FAIR*")

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