"""
TRIBE v2 Neural Engagement Analyzer — HuggingFace Space
Runs on T4 GPU with full WhisperX transcription pipeline.
Falls back to audio-only mode if WhisperX fails.
"""

import json
import os
import subprocess
import sys
import tempfile
import time
import traceback
from pathlib import Path

# Disable hf_transfer — uvx sandboxes don't have it
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"

import gradio as gr
import numpy as np
import pandas as pd
import requests

# Set HF_TOKEN
hf_token = os.environ.get("HF_TOKEN", "")
if hf_token:
    os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token
    print("HF token configured")

FUNCTIONAL_REGIONS = {
    "visual_processing": {
        "description": "Visual cortex",
        "labels": ["G_cuneus", "S_calcarine", "G_occipital_sup", "G_occipital_middle", "Pole_occipital", "G_and_S_occipital_inf", "G_oc-temp_lat-fusifor"],
    },
    "attention": {
        "description": "Prefrontal and parietal",
        "labels": ["G_front_sup", "G_front_middle", "S_front_sup", "G_parietal_sup", "S_intrapariet_and_P_trans", "G_precuneus"],
    },
    "emotional_response": {
        "description": "Limbic and temporal",
        "labels": ["G_cingul-Post-dorsal", "G_temp_sup-Lateral", "S_temporal_sup", "G_and_S_cingul-Ant"],
    },
    "language_processing": {
        "description": "Broca and Wernicke",
        "labels": ["G_front_inf-Opercular", "G_front_inf-Triangul", "G_temp_sup-G_T_transv", "G_temporal_middle", "G_pariet_inf-Angular"],
    },
    "motor_planning": {
        "description": "Motor and premotor",
        "labels": ["G_precentral", "S_precentral-sup-part", "S_central"],
    },
}

WEIGHTS = {"attention": 0.30, "visual_processing": 0.25, "emotional_response": 0.25, "language_processing": 0.15, "motor_planning": 0.05}

model = None
label_map = None


def build_events_without_whisperx(video_path: str) -> pd.DataFrame:
    """
    Build events dataframe from video WITHOUT using WhisperX.
    Extracts audio and creates Video+Audio events only.
    TRIBE v2 can still predict brain responses from video and audio
    features even without word-level transcription.
    """
    from neuralset.events.transforms import ExtractAudioFromVideo, ChunkEvents
    from neuralset.events.utils import standardize_events

    # Create initial video event
    event = {
        "type": "Video",
        "filepath": str(video_path),
        "start": 0,
        "timeline": "default",
        "subject": "default",
    }
    events = pd.DataFrame([event])
    events = standardize_events(events)

    # Extract audio from video
    transform_audio = ExtractAudioFromVideo()
    events = transform_audio(events)

    # Chunk long clips — keep short for faster processing
    transform_chunk_audio = ChunkEvents(event_type_to_chunk="Audio", max_duration=15, min_duration=5)
    transform_chunk_video = ChunkEvents(event_type_to_chunk="Video", max_duration=15, min_duration=5)
    events = transform_chunk_audio(events)
    events = transform_chunk_video(events)

    events = standardize_events(events)
    print(f"Built events: {len(events)} rows (audio_only mode, no WhisperX)")
    return events


def load_model():
    global model, label_map
    if model is not None:
        return

    from tribev2 import TribeModel
    from nilearn.datasets import fetch_atlas_surf_destrieux

    print("Loading TRIBE v2 model...")
    model = TribeModel.from_pretrained("facebook/tribev2", cache_folder="./cache")

    print("Loading Destrieux atlas...")
    atlas = fetch_atlas_surf_destrieux()
    label_map = {}
    labels = atlas.labels

    for i, label_entry in enumerate(labels):
        if isinstance(label_entry, (tuple, list)) and len(label_entry) == 2:
            idx, name = label_entry
        elif isinstance(label_entry, bytes):
            idx, name = i, label_entry.decode("utf-8")
        elif isinstance(label_entry, str):
            idx, name = i, label_entry
        else:
            idx, name = i, str(label_entry)

        if name == "Unknown" or idx == 0:
            continue
        left_verts = np.where(np.array(atlas.map_left) == idx)[0]
        right_verts = np.where(np.array(atlas.map_right) == idx)[0] + len(atlas.map_left)
        if len(left_verts) > 0 or len(right_verts) > 0:
            label_map[name] = np.concatenate([left_verts, right_verts])

    print(f"Loaded {len(label_map)} brain regions. Model ready.")


def compute_regions(preds):
    region_activations = {}
    for rname, rinfo in FUNCTIONAL_REGIONS.items():
        verts = []
        for lbl in rinfo["labels"]:
            if lbl in label_map:
                verts.append(label_map[lbl])
        if not verts:
            region_activations[rname] = np.zeros(preds.shape[0])
            continue
        all_v = np.concatenate(verts)
        all_v = all_v[all_v < preds.shape[1]]
        region_activations[rname] = np.mean(preds[:, all_v], axis=1)
    return region_activations


def compute_engagement(region_activations):
    """
    Compute per-second engagement from region activations.
    Uses z-score normalization instead of min-max to avoid the
    "everything peaks at the last second" artifact.
    """
    n = len(next(iter(region_activations.values())))
    composite = np.zeros(n)
    for region, weight in WEIGHTS.items():
        if region in region_activations:
            vals = region_activations[region]
            mean, std = vals.mean(), vals.std()
            if std > 0:
                # Z-score: how many SDs above/below mean activation
                z = (vals - mean) / std
                # Sigmoid to [0, 1] range — z=0 maps to 0.5
                normalized = 1 / (1 + np.exp(-z))
            else:
                normalized = np.full(n, 0.5)
            composite += weight * normalized
    return composite


def compute_score(engagement):
    """
    Score 1-10 based on variability and peak strength.
    Good ads have high peaks and dynamic variation (not flat).
    """
    if len(engagement) == 0:
        return 5
    peak = float(np.max(engagement))
    valley = float(np.min(engagement))
    mean = float(np.mean(engagement))
    dynamic_range = peak - valley  # 0 to ~1
    # Weight: 60% peak strength, 40% dynamic range
    raw = 0.6 * peak + 0.4 * dynamic_range
    # Scale to 1-10 (0.3 → 1, 0.8 → 10)
    score = 1 + 9 * max(0, min(1, (raw - 0.3) / 0.5))
    return max(1, min(10, round(score)))


def find_peak_region_seconds(region_activations, hrf_offset):
    """
    Find per-region peak seconds, excluding the very last second
    (which is an artifact of the model's attention accumulation).
    """
    summary = {}
    for rname, rinfo in FUNCTIONAL_REGIONS.items():
        if rname not in region_activations:
            continue
        vals = region_activations[rname]
        # Exclude last 2 time points (attention accumulation artifact)
        search_range = vals[:-2] if len(vals) > 3 else vals
        peak_idx = int(np.argmax(search_range))
        mean_val = float(np.mean(vals))
        peak_val = float(np.max(search_range))
        # Relative strength: how much above mean is the peak
        rel_strength = (peak_val - mean_val) / (mean_val + 1e-8)
        summary[rname] = {
            "description": rinfo["description"],
            "mean": round(mean_val, 4),
            "peak": round(peak_val, 4),
            "peak_second": peak_idx + hrf_offset,
            "relative_strength": round(rel_strength, 2),
        }
    return summary


def generate_insights(engagement, region_activations, hrf_offset=0, ad_context=None):
    """Generate data-specific creative insights from brain activation patterns."""
    # Ad context for tailored language
    ctx = ad_context or {}
    offer = ctx.get("offer", "key message")  # "the app", "the service", "the shoe", etc.
    cta_label = ctx.get("cta", "CTA")
    brand = ctx.get("brand", "the brand")

    insights = []
    n = len(engagement)
    if n < 3:
        return insights

    mean_eng = float(np.mean(engagement))
    std_eng = float(np.std(engagement))
    peak_idx = int(np.argmax(engagement))
    low_idx = int(np.argmin(engagement))
    peak_val = float(engagement[peak_idx])
    low_val = float(engagement[low_idx])

    # ── Compute per-region stats ──
    region_stats = {}
    for rname in ["visual_processing", "attention", "emotional_response", "language_processing", "motor_planning"]:
        if rname in region_activations:
            vals = region_activations[rname]
            safe = vals[:-2] if len(vals) > 3 else vals
            region_stats[rname] = {
                "mean": float(np.mean(safe)),
                "std": float(np.std(safe)),
                "peak_idx": int(np.argmax(safe)),
                "peak_val": float(np.max(safe)),
            }

    sorted_regions = sorted(region_stats.items(), key=lambda x: x[1]["mean"], reverse=True) if region_stats else []
    top_region = sorted_regions[0][0] if sorted_regions else None
    bottom_region = sorted_regions[-1][0] if sorted_regions else None

    # ── 1. Engagement Shape Analysis ──
    # Split into thirds and identify the shape
    third = max(1, n // 3)
    start_eng = float(np.mean(engagement[:third]))
    mid_eng = float(np.mean(engagement[third:2*third]))
    end_eng = float(np.mean(engagement[2*third:]))

    if start_eng > mid_eng and start_eng > end_eng:
        shape = "front-loaded"
    elif end_eng > mid_eng and end_eng > start_eng:
        shape = "back-loaded"
    elif mid_eng < start_eng and mid_eng < end_eng:
        shape = "u-shape"
    elif mid_eng > start_eng and mid_eng > end_eng:
        shape = "peak-middle"
    else:
        shape = "flat"

    shape_insights = {
        "front-loaded": {
            "title": "Front-Loaded Engagement",
            "severity": "medium",
            "message": f"Brain engagement is strongest in the first {third}s, then declines. Viewers get hooked early but lose interest before {offer} is reinforced.",
            "action": f"Move {offer} earlier — viewers are most receptive in the opening seconds. Or add a second hook in the final third to re-engage before the {cta_label}."
        },
        "back-loaded": {
            "title": "Slow Build — Late Peak",
            "severity": "medium",
            "message": f"Engagement builds toward the end, peaking after {2*third}s. The opening is the weakest section — viewers in a feed may scroll past before {offer} lands.",
            "action": f"Tease the payoff upfront: show a flash of the climax in the first 2 seconds to buy time. Or restructure to reveal {offer} immediately."
        },
        "u-shape": {
            "title": "Mid-Section Drop",
            "severity": "medium",
            "message": f"Strong opening and closing, but engagement dips in the middle ({third}s-{2*third}s). Attention drops between the hook and the {cta_label}.",
            "action": f"Add a pattern interrupt between {third}s-{2*third}s: scene change, new speaker, text overlay, or tempo shift. Bridge the gap between hook and {cta_label}."
        },
        "peak-middle": {
            "title": "Center-Weighted Engagement",
            "severity": "positive",
            "message": f"The core of the ad ({third}s-{2*third}s) drives the strongest response — this is where {offer} resonates most.",
            "action": f"The structure works. Ensure the {cta_label} appears during or immediately after this peak, not after engagement fades."
        },
        "flat": {
            "title": "Flat Response — Low Contrast",
            "severity": "medium",
            "message": f"Brain engagement is uniform throughout — no strong peaks or valleys. Flat response correlates with lower recall for {brand}.",
            "action": "Create contrast: add a surprising moment, an emotional beat, a visual disruption, or a change in pacing. Memorable ads need at least one peak."
        },
    }
    insights.append({"type": "shape", **shape_insights[shape]})

    # ── 2. Hook Verdict ──
    hook_3s = float(np.mean(engagement[:min(3, n)]))
    hook_vs_avg = (hook_3s - mean_eng) / (mean_eng + 1e-8) * 100

    if hook_vs_avg < -15:
        # Find what's WEAK in the hook
        hook_regions = {r: float(np.mean(region_activations[r][:min(3, n)])) for r in region_stats}
        weakest_hook_region = min(hook_regions, key=hook_regions.get) if hook_regions else None
        weakness_label = {
            "visual_processing": "visuals aren't grabbing attention",
            "attention": "nothing demands focus",
            "emotional_response": "no emotional hook",
            "language_processing": "the copy/audio doesn't land immediately",
            "motor_planning": "no impulse to engage",
        }.get(weakest_hook_region, "the opening lacks stimulus")

        insights.append({
            "type": "hook", "severity": "high",
            "title": f"Weak Hook — {abs(hook_vs_avg):.0f}% Below Average",
            "message": f"The first 3 seconds underperform the rest of the ad. Specifically, {weakness_label} in the opening.",
            "action": "Restructure the first 2 seconds: lead with motion, a face, bold text, or a sound that demands attention. The hook determines 80% of view-through rate."
        })
    elif hook_vs_avg > 15:
        # Find what's STRONG in the hook
        hook_regions = {r: float(np.mean(region_activations[r][:min(3, n)])) for r in region_stats}
        strongest_hook_region = max(hook_regions, key=hook_regions.get) if hook_regions else None
        strength_label = {
            "visual_processing": "strong visual opening",
            "attention": "immediately grabs focus",
            "emotional_response": "emotional hook",
            "language_processing": "compelling opening line/audio",
            "motor_planning": "action-triggering opener",
        }.get(strongest_hook_region, "effective opening")

        insights.append({
            "type": "hook", "severity": "positive",
            "title": f"Strong Hook — {hook_vs_avg:.0f}% Above Average",
            "message": f"The opening outperforms the rest of the ad, driven by {strength_label}. This hook pattern is worth replicating.",
            "action": "Document this hook structure — it works. Test variations of the same opening pattern across other creatives."
        })

    # ── 3. Engagement Drops — find the worst dip ──
    if n >= 6:
        # Find the biggest drop (3-second rolling window)
        window = 3
        rolling = [float(np.mean(engagement[max(0,i-1):i+window-1])) for i in range(n - window + 1)]
        worst_window_idx = int(np.argmin(rolling))
        worst_window_val = rolling[worst_window_idx]
        drop_pct = (mean_eng - worst_window_val) / (mean_eng + 1e-8) * 100

        if drop_pct > 15:
            drop_start = worst_window_idx + hrf_offset
            drop_end = worst_window_idx + window + hrf_offset
            # What region drops the most here?
            drop_regions = {}
            for rname in region_stats:
                region_at_drop = float(np.mean(region_activations[rname][worst_window_idx:worst_window_idx+window]))
                region_overall = region_stats[rname]["mean"]
                drop_regions[rname] = (region_overall - region_at_drop) / (region_overall + 1e-8)
            biggest_drop_region = max(drop_regions, key=drop_regions.get) if drop_regions else None
            drop_label = {
                "visual_processing": "visual interest drops — the scene may be static or repetitive",
                "attention": "cognitive focus fades — nothing new to process",
                "emotional_response": "emotional connection weakens",
                "language_processing": "the message loses clarity or relevance",
                "motor_planning": "action intent drops — the ad feels passive",
            }.get(biggest_drop_region, "engagement drops")

            insights.append({
                "type": "drop", "severity": "high",
                "title": f"Attention Drop at {drop_start}s-{drop_end}s ({drop_pct:.0f}% below avg)",
                "message": f"The weakest moment in the ad. Specifically, {drop_label} during this window.",
                "action": f"Review what happens at {drop_start}s-{drop_end}s in the video. Add a scene change, new audio cue, or text overlay to re-engage."
            })

    # ── 4. Peak Moment — what makes it work ──
    peak_s = peak_idx + hrf_offset
    # Find what region is strongest at the peak
    peak_regions = {r: float(region_activations[r][peak_idx]) for r in region_stats}
    peak_driver = max(peak_regions, key=peak_regions.get) if peak_regions else None
    peak_label = {
        "visual_processing": "a strong visual moment",
        "attention": "high cognitive focus",
        "emotional_response": "an emotional beat",
        "language_processing": "a compelling message or audio cue",
        "motor_planning": "an action trigger",
    }.get(peak_driver, "a peak in brain activity")

    insights.append({
        "type": "peak", "severity": "positive",
        "title": f"Peak at {peak_s}s — Driven by {peak_label.title()}",
        "message": f"The strongest brain response occurs at {peak_s}s, primarily driven by {peak_label}. This is the most impactful moment in the ad.",
        "action": f"Align {offer} with this moment. If the {cta_label} isn't near {peak_s}s, consider moving it closer."
    })

    # ── 5. CTA Window ──
    if n >= 5:
        last_3 = float(np.mean(engagement[-3:]))
        cta_vs_avg = (last_3 - mean_eng) / (mean_eng + 1e-8) * 100
        if cta_vs_avg < -15:
            insights.append({
                "type": "cta", "severity": "high",
                "title": f"{cta_label} Lands on Low Engagement ({cta_vs_avg:.0f}%)",
                "message": f"The final 3 seconds are {abs(cta_vs_avg):.0f}% below average. Viewers' attention has faded when the {cta_label} appears.",
                "action": f"Option A: Move the {cta_label} to the peak moment ({peak_s}s). Option B: Add a re-engagement spike before the {cta_label} — a bold visual, urgency text, or direct question."
            })
        elif cta_vs_avg > 10:
            insights.append({
                "type": "cta", "severity": "positive",
                "title": f"Strong {cta_label} Window",
                "message": f"Engagement is {cta_vs_avg:.0f}% above average at the end. Viewers are attentive when the {cta_label} appears.",
                "action": f"Make sure the {cta_label} is visually prominent and uses action language. This audience is primed to act."
            })

    # ── 6. Region Dominance (only if there's a clear leader) ──
    if len(sorted_regions) >= 2:
        top_name, top_stats = sorted_regions[0]
        second_name, second_stats = sorted_regions[1]
        dominance = (top_stats["mean"] - second_stats["mean"]) / (second_stats["mean"] + 1e-8) * 100

        if dominance > 10:
            labels = {
                "visual_processing": ("Visually-Driven", "imagery and motion", "Test different visual treatments — color, pacing, composition — since the audience responds primarily to what they see."),
                "attention": ("Cognitively Demanding", "focused thinking", "Balance information density with breathing room. Too much cognitive load causes fatigue and drop-off."),
                "emotional_response": ("Emotion-Led", "feelings over logic", "Use emotion-aligned CTAs ('Don't miss out', 'Join the movement') — they outperform rational CTAs for this type of creative."),
                "language_processing": ("Copy-Driven", "words and audio", "The copy is doing the heavy lifting. A/B test different scripts, voiceover styles, or text overlays."),
                "motor_planning": ("Action-Primed", "impulse to act", "This is ideal for performance ads. Place the CTA at the moment of highest motor activation for maximum conversion."),
            }
            if top_name in labels:
                label, driver, action = labels[top_name]
                insights.append({
                    "type": "dominance", "severity": "positive",
                    "title": f"{label} Creative ({dominance:.0f}% stronger than #{2})",
                    "message": f"This ad's engagement is primarily driven by {driver}. The {labels.get(top_name, ('','',''))[0].lower()} signal is {dominance:.0f}% stronger than the next closest region.",
                    "action": action,
                })

    return insights


MAX_DURATION = 60  # Analyze up to 60s of video


def trim_video(input_path: str, max_seconds: int = MAX_DURATION) -> str:
    """Trim video to max_seconds and reduce FPS for faster encoding."""
    output_path = input_path.replace(".mp4", "_trimmed.mp4")
    try:
        result = subprocess.run(
            ["ffmpeg", "-y", "-i", input_path, "-t", str(max_seconds),
             "-c:v", "libx264", "-preset", "ultrafast", "-crf", "28",
             "-vf", "fps=2,scale='min(320,iw)':-2",  # 2fps + 320px width
             "-c:a", "aac", "-b:a", "64k", "-ar", "16000", output_path],
            capture_output=True, text=True, timeout=30,
        )
        if result.returncode == 0 and Path(output_path).exists():
            size = Path(output_path).stat().st_size
            print(f"Trimmed to {max_seconds}s @2fps 320px: {size} bytes")
            return output_path
    except Exception as e:
        print(f"Trim failed ({e}), using original")
    return input_path


def analyze_video(video_url: str, filename: str = "video.mp4", ad_context_json: str = "{}") -> str:
    """Main analysis — uses full TRIBE v2 pipeline with WhisperX on GPU."""
    try:
        # Parse ad context for tailored insights
        try:
            ad_context = json.loads(ad_context_json) if ad_context_json else {}
        except Exception:
            ad_context = {}

        load_model()
        start_time = time.time()
        # HRF offset: TRIBE v2 internally accounts for ~5s hemodynamic delay,
        # but we report seconds as video time (0-based) so users can match
        # insights to actual moments in their video.
        hrf_offset = 0

        # Download video
        print(f"Downloading: {video_url[:100]}")
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
            resp = requests.get(video_url, timeout=120)
            resp.raise_for_status()
            f.write(resp.content)
            video_path = f.name
        print(f"Downloaded {len(resp.content)} bytes")

        # Trim to MAX_DURATION seconds and downscale for faster encoding
        video_path = trim_video(video_path)

        # Build events — try full pipeline (WhisperX + audio + video), fallback to audio-only
        try:
            print("Building events with full pipeline (WhisperX)...")
            df = model.get_events_dataframe(video_path, "default", "default")
            print(f"Events: {len(df)} rows (full pipeline with transcription)")
        except Exception as whisperx_err:
            print(f"WhisperX failed ({whisperx_err}), falling back to audio-only...")
            df = build_events_without_whisperx(video_path)
            print(f"Events: {len(df)} rows (audio-only fallback)")

        # Run brain prediction
        print("Running brain prediction...")
        preds, segments = model.predict(events=df)
        print(f"Predictions: {preds.shape}")

        # Compute engagement
        region_activations = compute_regions(preds)
        engagement = compute_engagement(region_activations)
        mean_eng = float(np.mean(engagement))
        overall_score = compute_score(engagement)

        # Build timeline (exclude last 2 points — attention accumulation artifact)
        valid_len = max(1, len(engagement) - 2)
        timeline = [
            {
                "second": t + hrf_offset,
                "engagement": round(float(engagement[t]), 3),
                "regions": {r: round(float(region_activations[r][t]), 4)
                            for r in FUNCTIONAL_REGIONS if r in region_activations},
            }
            for t in range(valid_len)
        ]

        # Region summary with artifact-corrected peaks
        region_summary = find_peak_region_seconds(region_activations, hrf_offset)

        # Find peak/lowest in the valid range (excluding tail artifact)
        valid_engagement = engagement[:valid_len]
        peak_idx = int(np.argmax(valid_engagement))
        low_idx = int(np.argmin(valid_engagement))

        insights = generate_insights(valid_engagement, region_activations, hrf_offset, ad_context)
        elapsed = time.time() - start_time

        report = {
            "filename": filename,
            "duration_seconds": valid_len + hrf_offset,
            "analyzed_seconds": valid_len,
            "overall_score": overall_score,
            "mean_engagement": round(mean_eng, 3),
            "peak": {
                "value": round(float(valid_engagement[peak_idx]), 3),
                "second": peak_idx + hrf_offset,
            },
            "lowest": {
                "value": round(float(valid_engagement[low_idx]), 3),
                "second": low_idx + hrf_offset,
            },
            "region_summary": region_summary,
            "timeline": timeline,
            "insights": insights,
            "inference_time_seconds": round(elapsed, 1),
            "model": "facebook/tribev2",
        }
        Path(video_path).unlink(missing_ok=True)
        print(f"Analysis complete: score={overall_score}, {len(engagement)} timesteps, {elapsed:.1f}s")
        return json.dumps(report)

    except Exception as e:
        error_msg = f"Error: {type(e).__name__}: {str(e)}"
        print(error_msg)
        traceback.print_exc()
        return json.dumps({"error": error_msg, "traceback": traceback.format_exc()})


app = gr.Blocks(title="TRIBE v2 Neural Engagement Analyzer")
with app:
    gr.Markdown("# TRIBE v2 Neural Engagement Analyzer")
    gr.Markdown("Predict brain response to video content, second by second.")
    with gr.Row():
        video_url_input = gr.Textbox(label="Video URL", placeholder="https://example.com/video.mp4")
        filename_input = gr.Textbox(label="Filename", value="video.mp4")
    ad_context_input = gr.Textbox(label="Ad Context (JSON)", value="{}", visible=False)
    analyze_btn = gr.Button("Analyze", variant="primary")
    output = gr.JSON(label="Neural Engagement Report")
    analyze_btn.click(fn=analyze_video, inputs=[video_url_input, filename_input, ad_context_input], outputs=output)
app.queue()
if __name__ == "__main__":
    app.launch()