"""
predict.py — Amphora NeuroText inference

Models (honest cross-subject holdout val_R — no data leakage)
------
text2roi_whisper_v4.pt    Audio → 56 ROIs  (Whisper-large-v3, 1280d, val R=0.217, holdout R=0.257)
text2roi_combined_v4.pt   Text+Audio → 56 ROIs (whisper|qwen3, 3840d, val R=0.192)
text2roi_qwen3_v8.pt      Text → 56 ROIs  (Qwen3-Embedding-4B, 2560d, val R=0.115)

All v4 models use per-subject z-scoring and a per-subject train/val split.
Previous models (v2/v3) had inflated val_R from within-subject splits — do not compare directly.

Whisper v4 beats TRIBE v2 (Meta AI, Algonauts 2025 winner) by +4.2% on a 23-subject holdout.
Single shared model — no per-subject fine-tuning required.

Quick start
-----------
    from predict import predict_text, predict_audio, top_rois

    # Audio → brain regions (recommended — strongest model)
    roi_map = predict_audio("clip.wav")
    print(top_rois(roi_map, n=5))
    # → [('ACC', 0.44), ('STG', 0.42), ('Thalamus', 0.41), ...]

    # Text → brain regions
    roi_map = predict_text("I am terrified of the dark")
    print(top_rois(roi_map, n=5))
    # → [('Amygdala_L', 0.xx), ('AI', 0.xx), ('dACC', 0.xx), ...]

    # Text + Audio → brain regions
    roi_map = predict_combined("narration text", "clip.wav")
    print(top_rois(roi_map, n=5))
"""
from __future__ import annotations

import sys
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np
import torch
import torch.nn as nn

# ── Canonical 56-ROI schema ───────────────────────────────────────────────────
ROI_NAMES_56: List[str] = [
    "V1", "V2", "V3", "V4", "V3A", "V3B", "LO1", "LO2",
    "MT", "MST", "V7", "IPS1", "FFA-1", "FFA-2", "PPA", "RSC",
    "OFA", "EBA", "IPS2", "IPS3", "IPS4", "IPS5", "SPL1",
    "hIP1", "hIP2", "hIP3", "dlPFC", "vlPFC", "OFC", "ACC",
    "mPFC", "FP1", "FP2", "IFG", "IFGorb", "STG", "STS",
    "MTG", "AG", "PCC", "mPFC_dmn", "LP_L", "LP_R",
    "HPC_L", "HPC_R", "AI", "dACC", "sgACC", "vmPFC",
    "Amygdala_L", "Amygdala_R", "Caudate_L", "Caudate_R",
    "Putamen_L", "Putamen_R", "Thalamus",
]

DEFAULT_CHECKPOINT = "text2roi_whisper_v4.pt"

# ── Shared MLP architecture ───────────────────────────────────────────────────
class Text2ROI(nn.Module):
    def __init__(self, in_dim: int = 1280, hidden: int = 1024,
                 out_dim: int = 56, dropout: float = 0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(hidden),
            nn.Linear(hidden, hidden // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden // 2, out_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# ── Model loader ──────────────────────────────────────────────────────────────
_model_cache: Dict[str, tuple] = {}


def _load_model(checkpoint: str) -> tuple:
    if checkpoint in _model_cache:
        return _model_cache[checkpoint]

    # 1. absolute or relative-to-CWD path
    ckpt_path = Path(checkpoint)
    # 2. same directory as this script (works when bundled in a zip/folder)
    if not ckpt_path.exists():
        ckpt_path = Path(__file__).parent / checkpoint
    # 3. HuggingFace Hub (online fallback)
    if not ckpt_path.exists():
        try:
            from huggingface_hub import hf_hub_download
            ckpt_path = Path(hf_hub_download("ffh92r32rm0/Amphora_NeuroText", checkpoint))
        except Exception as e:
            raise FileNotFoundError(
                f"Checkpoint '{checkpoint}' not found locally or on HuggingFace.\n"
                f"Make sure the .pt file is in the same folder as predict.py.\n"
                f"Original error: {e}"
            )

    state = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
    in_dim = state.get("in_dim", 1280)
    n_roi  = state.get("n_roi", 56)
    model  = Text2ROI(in_dim=in_dim, out_dim=n_roi)
    model.load_state_dict(state["state_dict"])
    model.eval()
    _model_cache[checkpoint] = (model, in_dim)
    return model, in_dim


# ── Qwen3 embedding ───────────────────────────────────────────────────────────
_qwen3_model = None
_qwen3_tokenizer = None


def _embed_text_qwen3(text: str) -> np.ndarray:
    global _qwen3_model, _qwen3_tokenizer
    if _qwen3_model is None:
        from transformers import AutoTokenizer, AutoModel
        _qwen3_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Embedding-4B")
        _qwen3_model = AutoModel.from_pretrained("Qwen/Qwen3-Embedding-4B")
        _qwen3_model.eval()
    inputs = _qwen3_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        out = _qwen3_model(**inputs)
    emb = out.last_hidden_state[:, 0, :].squeeze(0).cpu().numpy()
    return emb.astype(np.float32)


# ── Whisper embedding ─────────────────────────────────────────────────────────
_whisper_model = None
_whisper_processor = None


def _embed_audio_whisper(audio_path: str) -> np.ndarray:
    global _whisper_model, _whisper_processor
    if _whisper_model is None:
        from transformers import WhisperProcessor, WhisperModel
        _whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
        _whisper_model = WhisperModel.from_pretrained("openai/whisper-large-v3")
        _whisper_model.eval()
    import librosa
    audio, sr = librosa.load(audio_path, sr=16000, mono=True)
    inputs = _whisper_processor(audio, sampling_rate=16000, return_tensors="pt")
    with torch.no_grad():
        enc = _whisper_model.encoder(inputs.input_features)
    emb = enc.last_hidden_state.mean(dim=1).squeeze(0).cpu().numpy()
    return emb.astype(np.float32)


# ── Public API ────────────────────────────────────────────────────────────────

def predict_audio(
    audio_path: str,
    checkpoint: str = "text2roi_whisper_v4.pt",
) -> Dict[str, float]:
    """Predict 56 brain ROI activations from an audio file.

    Args:
        audio_path: Path to audio file (wav, mp3, flac, …).
        checkpoint:  Model file. Default: text2roi_whisper_v4.pt (val R=0.217, holdout R=0.257).

    Returns:
        dict mapping ROI name → predicted activation (z-scored units).
    """
    model, in_dim = _load_model(checkpoint)
    emb = _embed_audio_whisper(audio_path)
    if emb.shape[0] != in_dim:
        raise ValueError(f"Audio embedding dim {emb.shape[0]} != model in_dim {in_dim}")
    feat = torch.from_numpy(emb).unsqueeze(0)
    with torch.no_grad():
        pred = model(feat).squeeze(0).numpy()
    return {roi: float(pred[i]) for i, roi in enumerate(ROI_NAMES_56)}


def predict_text(
    text: str,
    checkpoint: str = "text2roi_combined_v4.pt",
) -> Dict[str, float]:
    """Predict 56 brain ROI activations from a text string.

    Uses text2roi_combined_v4.pt by default: qwen3 embedding zero-padded into
    the whisper slot (positions 0:1280 = zeros, 1280:3840 = qwen3 2560d).
    The combined model was trained with modality dropout so text-only works.

    Args:
        text:       Input text string.
        checkpoint: Model file. Default: text2roi_combined_v4.pt (val R=0.192).

    Returns:
        dict mapping ROI name → predicted activation (z-scored units).
    """
    model, in_dim = _load_model(checkpoint)
    emb = _embed_text_qwen3(text)

    if in_dim == 3840:
        # combined model: zero-pad whisper slot
        feat_np = np.concatenate([np.zeros(1280, dtype=np.float32), emb])
    elif in_dim == 2560:
        feat_np = emb
    else:
        raise ValueError(f"Unexpected model in_dim {in_dim} for text inference")

    feat = torch.from_numpy(feat_np).unsqueeze(0)
    with torch.no_grad():
        pred = model(feat).squeeze(0).numpy()
    return {roi: float(pred[i]) for i, roi in enumerate(ROI_NAMES_56)}


def predict_combined(
    text: str,
    audio_path: str,
    checkpoint: str = "text2roi_combined_v4.pt",
) -> Dict[str, float]:
    """Predict 56 brain ROI activations from both text and audio.

    Args:
        text:        Text string.
        audio_path:  Path to audio file.
        checkpoint:  Must be the combined model (in_dim=3840).

    Returns:
        dict mapping ROI name → predicted activation (z-scored units).
    """
    model, in_dim = _load_model(checkpoint)
    if in_dim != 3840:
        raise ValueError("predict_combined requires the combined model (in_dim=3840).")
    w_emb = _embed_audio_whisper(audio_path)
    q_emb = _embed_text_qwen3(text)
    feat_np = np.concatenate([w_emb, q_emb]).astype(np.float32)
    feat = torch.from_numpy(feat_np).unsqueeze(0)
    with torch.no_grad():
        pred = model(feat).squeeze(0).numpy()
    return {roi: float(pred[i]) for i, roi in enumerate(ROI_NAMES_56)}


# ── Utilities ─────────────────────────────────────────────────────────────────

def top_rois(roi_map: Dict[str, float], n: int = 10) -> List[tuple]:
    """Return top-n ROIs sorted by predicted activation."""
    return sorted(roi_map.items(), key=lambda x: -x[1])[:n]


def network_summary(roi_map: Dict[str, float]) -> Dict[str, float]:
    """Return mean activation per brain network."""
    NETWORKS = {
        "Visual":      ["V1","V2","V3","V4","V3A","V3B","LO1","LO2","MT","MST","V7","IPS1","FFA-1","FFA-2","PPA","RSC","OFA","EBA"],
        "Parietal":    ["IPS2","IPS3","IPS4","IPS5","SPL1","hIP1","hIP2","hIP3"],
        "Frontal":     ["dlPFC","vlPFC","OFC","ACC","mPFC","FP1","FP2"],
        "Language":    ["IFG","IFGorb","STG","STS","MTG","AG"],
        "DefaultMode": ["PCC","mPFC_dmn","LP_L","LP_R","HPC_L","HPC_R"],
        "Salience":    ["AI","dACC","sgACC","vmPFC","Amygdala_L","Amygdala_R"],
        "Subcortical": ["Caudate_L","Caudate_R","Putamen_L","Putamen_R","Thalamus"],
    }
    return {
        net: float(np.mean([roi_map[r] for r in rois if r in roi_map]))
        for net, rois in NETWORKS.items()
    }


# ── CLI ───────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import argparse, json

    parser = argparse.ArgumentParser(description="Amphora NeuroText — brain ROI prediction")
    sub = parser.add_subparsers(dest="cmd")

    p_audio = sub.add_parser("audio", help="Audio file → brain ROIs")
    p_audio.add_argument("audio", help="Path to audio file")
    p_audio.add_argument("--model", default="text2roi_whisper_v4.pt")
    p_audio.add_argument("--top", type=int, default=10)

    p_text = sub.add_parser("text", help="Text → brain ROIs")
    p_text.add_argument("text", help="Input text string")
    p_text.add_argument("--model", default="text2roi_combined_v4.pt")
    p_text.add_argument("--top", type=int, default=10)

    p_combo = sub.add_parser("combined", help="Text + Audio → brain ROIs")
    p_combo.add_argument("text")
    p_combo.add_argument("audio")
    p_combo.add_argument("--model", default="text2roi_combined_v4.pt")
    p_combo.add_argument("--top", type=int, default=10)

    args = parser.parse_args()

    if args.cmd == "audio":
        result = predict_audio(args.audio, args.model)
    elif args.cmd == "text":
        result = predict_text(args.text, args.model)
    elif args.cmd == "combined":
        result = predict_combined(args.text, args.audio, args.model)
    else:
        parser.print_help()
        sys.exit(0)

    tops = top_rois(result, args.top)
    nets = network_summary(result)
    print(f"\nTop {args.top} ROIs:")
    for roi, val in tops:
        print(f"  {roi:<16} {val:+.4f}")
    print("\nNetwork summary:")
    for net, val in sorted(nets.items(), key=lambda x: -x[1]):
        print(f"  {net:<16} {val:+.4f}")