# -*- coding: utf-8 -*-

import argparse
import glob
import json
import math
import os
import traceback
from typing import Any, Dict, List, Tuple
import numpy as np
import pandas as pd

CACHE_DIR = "ENCODERS/hf_cache"

ENCODER_MAP = {
    "bert":       "bert-large-uncased",
    "longformer": "allenai/longformer-large-4096",
    "vit":        "google/vit-base-patch16-224-in21k",
    "dinov2":     "facebook/dinov2-base",
}

PAIR_NAME = {
    ("bert-large-uncased",            "google/vit-base-patch16-224-in21k"): "BERT_ViT",
    ("allenai/longformer-large-4096", "facebook/dinov2-base"):              "LONG_DINO",
    ("bert-large-uncased",            "facebook/dinov2-base"):              "BERT_DINO",
    ("allenai/longformer-large-4096", "google/vit-base-patch16-224-in21k"): "LONG_ViT",
}

ALL_STAGES = [
    "json_to_csv",
    "video_to_frames",
    "download_models",
    "extract_features",
    "compute_coattention",
    "simple_fuse",
    "fuse_embeddings",
    "entropy_scores",
]

_PROJ_CACHE: Dict = {}

def run_path_for(pair: str, prefix: str = None) -> str:
    if prefix:
        return os.path.join(prefix, pair)
    return pair

def run_name_for(text_model: str, vision_model: str) -> str:
    return PAIR_NAME.get(
        (text_model, vision_model),
        f"{text_model.replace('/', '__')}__{vision_model.replace('/', '__')}",
    )


def resolve_encoder(short: str, kind: str) -> str:
    if short in ENCODER_MAP:
        return ENCODER_MAP[short]
    if short in set(ENCODER_MAP.values()):
        return short
    raise ValueError(
        f"Unknown {kind} encoder '{short}'. "
        f"Choose from: {list(ENCODER_MAP.keys())} or pass a full HuggingFace model id."
    )


def _read_manifest(manifest_path: str) -> List[str]:
    with open(manifest_path, "r", encoding="utf-8") as f:
        return [ln.strip() for ln in f if ln.strip() and not ln.strip().startswith("#")]


def _resolve_path(path: str, base_dir: str):
    if not os.path.isabs(path):
        candidate = os.path.join(base_dir, path)
        if os.path.exists(candidate):
            return candidate
    if os.path.exists(path):
        return path
    print(f"  [SKIP] Missing file: {path}")
    return None


def _build_manifest(csv_dir: str, manifest_path: str) -> None:
    csvs = sorted(glob.glob(os.path.join(csv_dir, "*.csv")))
    os.makedirs(os.path.dirname(os.path.abspath(manifest_path)), exist_ok=True)
    with open(manifest_path, "w", encoding="utf-8") as f:
        for p in csvs:
            f.write(p + "\n")
    print(f"  Manifest written: {manifest_path}  ({len(csvs)} entries)")


def _safe_load(path: str):
    import torch
    try:
        return torch.load(path, weights_only=True)
    except TypeError:
        return torch.load(path)


def _get_projection(in_dim: int, out_dim: int, device):
    import torch
    key = (in_dim, out_dim, str(device))
    if key in _PROJ_CACHE:
        return _PROJ_CACHE[key]
    g = torch.Generator(device="cpu")
    g.manual_seed(in_dim * 1000003 + out_dim * 9176)
    W = (torch.randn(out_dim, in_dim, generator=g) / (in_dim ** 0.5)).to(device)
    _PROJ_CACHE[key] = W
    return W


def _project(x, out_dim: int):
    import torch.nn.functional as F
    in_dim = x.shape[-1]
    if in_dim == out_dim:
        return x
    return F.linear(x, _get_projection(in_dim, out_dim, x.device))


def _weighted_sharing(emb_list):
    import torch
    import torch.nn.functional as F
    stacked = torch.cat(emb_list, dim=0)
    weights = F.softmax(stacked.mean(dim=1), dim=0)
    return (stacked * weights.unsqueeze(1)).sum(dim=0, keepdim=True)

# Converting JSON to CSV format (screen-change segments)

def _ts_to_seconds(ts: Dict[str, Any]) -> float:
    if ts is None:
        return math.nan
    return float(ts.get("seconds", 0) or 0) + float(ts.get("nanos", 0) or 0) / 1e9


def _to_int_sec(x: float, mode: str = "floor") -> int:
    if math.isnan(x):
        return 0
    if mode == "ceil":  return int(math.ceil(x))
    if mode == "round": return int(round(x))
    return int(math.floor(x))


def _build_timestamp_list(data: Dict[str, Any], rounding: str = "floor") -> List[List[int]]:
    ann   = data.get("annotation_results", [])
    shots = (ann[0].get("shot_annotations") if ann and isinstance(ann[0], dict) else []) or []
    if len(shots) == 0:
        shots = (ann[1].get("shot_annotations") if ann and isinstance(ann[0], dict) else []) or []
    segments, prev_end = [], None
    for shot in shots:
        raw_end = _ts_to_seconds((shot or {}).get("end_time_offset"))
        if math.isnan(raw_end):
            continue
        end_i   = _to_int_sec(raw_end, mode=rounding)
        start_i = 1 if prev_end is None else prev_end + 1
        if end_i <= start_i:
            end_i = start_i + 1
        segments.append([start_i, end_i])
        prev_end = end_i
    return segments


def _collect_all_words(data: Dict[str, Any]) -> List[Dict[str, Any]]:
    ann    = data.get("annotation_results", [])
    speech = (ann[1].get("speech_transcriptions")
              if len(ann) > 1 and isinstance(ann[1], dict) else []) or []
    if len(speech) == 0:
        speech = (ann[0].get("speech_transcriptions")
                  if len(ann) > 1 and isinstance(ann[1], dict) else []) or []
    words  = []
    for item in speech:
        alts = (item or {}).get("alternatives", []) or []
        if alts:
            words.extend((alts[0] or {}).get("words", []) or [])
    return words


def _assign_words_to_segments(
    all_words: List[Dict[str, Any]],
    segments:  List[List[int]],
    mode: str = "overlap",
) -> List[Dict[str, Any]]:
    words: List[Tuple[float, float, str]] = []
    for w in all_words:
        s   = _ts_to_seconds(w.get("start_time"))
        e   = _ts_to_seconds(w.get("end_time"))
        txt = (w.get("word") or "").strip()
        if not (math.isnan(s) or math.isnan(e) or not txt):
            words.append((s, e, txt))
    words.sort(key=lambda x: x[0])

    out, i, n = [], 0, len(words)
    for seg_start, seg_end in segments:
        sf, ef = float(seg_start), float(seg_end)
        while i < n and words[i][1] <= sf:
            i += 1
        j, seg_words = i, []
        while j < n and words[j][0] < ef:
            ws, we, wt = words[j]
            keep = (ws < ef and we > sf) if mode == "overlap" else (ws >= sf and we <= ef)
            if keep:
                seg_words.append(wt)
            j += 1
        if seg_words:
            out.append({
                "sentences":  " ".join(seg_words),
                "start_time": int(seg_start),
                "end_time":   int(seg_end),
            })
    return out


def stage_json_to_csv(args: argparse.Namespace) -> None:
    print("\n" + "=" * 60)
    print("STAGE 1  JSON -> CSV")
    print("=" * 60)
    os.makedirs(args.csv_dir, exist_ok=True)
    if any(ch in args.json_dir for ch in ["*", "?", "["]):
        json_paths = sorted(glob.glob(args.json_dir))
    else:
        json_paths = sorted(glob.glob(os.path.join(args.json_dir, "*.json")))
    if not json_paths:
        print(f"  [WARN] No JSON files found in: {args.json_dir}")
        return
    print(f"  Found {len(json_paths)} JSON files -> writing CSVs to: {args.csv_dir}")
    n_ok = n_fail = 0
    for jp in json_paths:
        video_id = os.path.splitext(os.path.basename(jp))[0]
        out_csv  = os.path.join(args.csv_dir, f"{video_id}.csv")
        try:
            with open(jp, "r", encoding="utf-8") as f:
                data = json.load(f)
            rows = _assign_words_to_segments(
                _collect_all_words(data),
                _build_timestamp_list(data),
                mode="overlap",
            )
            if not rows:
                print(f"  [SKIP] {video_id}: no text segments found")
                continue
            pd.DataFrame(rows).to_csv(out_csv, index=False)
            n_ok += 1
            print(f"  [OK]   {video_id} -> {out_csv}  ({len(rows)} segments)")
        except Exception as e:
            n_fail += 1
            print(f"  [FAIL] {video_id}: {e}")
    print(f"\n  Stage 1 complete.  Written: {n_ok}   Failed: {n_fail}")


# Converting JSON to CSV format (sentence-based segments)

def _collect_full_transcript(data: Dict[str, Any]) -> List[Dict[str, Any]]:
    ann    = data.get("annotation_results", [])
    speech = (ann[1].get("speech_transcriptions")
              if len(ann) > 1 and isinstance(ann[1], dict) else []) or []

    if len(speech) == 0:
        speech = (ann[0].get("speech_transcriptions")
                  if ann and isinstance(ann[0], dict) else []) or []

    words  = []
    for item in speech:
        alts = (item or {}).get("alternatives", []) or []
        if alts:
            words.extend((alts[0] or {}).get("words", []) or [])
    return words

def _assign_words_to_sentences(all_words: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    import nltk
    try:
        nltk.data.find("tokenizers/punkt")
    except LookupError:
        print("  [INFO] Downloading NLTK punkt tokenizer...")
        nltk.download("punkt", quiet=True)
    try:
        nltk.data.find("tokenizers/punkt_tab")
    except LookupError:
        nltk.download("punkt_tab", quiet=True)

    # Build flat word list with timestamps
    timed_words = []
    for w in all_words:
        s   = _ts_to_seconds(w.get("start_time"))
        e   = _ts_to_seconds(w.get("end_time"))
        txt = (w.get("word") or "").strip()
        if not (math.isnan(s) or math.isnan(e) or not txt):
            timed_words.append({"word": txt, "start": s, "end": e})

    if not timed_words:
        return []

    full_text = " ".join(w["word"] for w in timed_words)
    sentences = nltk.sent_tokenize(full_text)

    rows = []
    word_idx = 0
    n = len(timed_words)

    for sent in sentences:
        sent_words = sent.split()
        if not sent_words:
            continue

        matched = []
        j = word_idx
        for sw in sent_words:
            while j < n:
                tw_clean = timed_words[j]["word"].strip(".,!?;:")
                sw_clean = sw.strip(".,!?;:")
                if tw_clean.lower() == sw_clean.lower():
                    matched.append(timed_words[j])
                    j += 1
                    break
                j += 1

        if not matched:
            continue

        start_time = int(math.floor(matched[0]["start"]))
        end_time   = int(math.ceil(matched[-1]["end"]))
        if end_time <= start_time:
            end_time = start_time + 1

        rows.append({
            "sentences":  sent,
            "start_time": start_time,
            "end_time":   end_time,
        })
        word_idx = j  

    return rows


def stage_json_to_csv_sentences(args: argparse.Namespace) -> None:
    os.makedirs(args.csv_dir, exist_ok=True)

    if any(ch in args.json_dir for ch in ["*", "?", "["]):
        json_paths = sorted(glob.glob(args.json_dir))
    else:
        json_paths = sorted(glob.glob(os.path.join(args.json_dir, "*.json")))

    if not json_paths:
        print(f"  [WARN] No JSON files found in: {args.json_dir}")
        return

    print(f"  Found {len(json_paths)} JSON files -> writing CSVs to: {args.csv_dir}")
    n_ok = n_fail = 0

    for jp in json_paths:
        video_id = os.path.splitext(os.path.basename(jp))[0]
        out_csv  = os.path.join(args.csv_dir, f"{video_id}.csv")
        try:
            with open(jp, "r", encoding="utf-8") as f:
                data = json.load(f)
            all_words = _collect_full_transcript(data)
            rows      = _assign_words_to_sentences(all_words)
            if not rows:
                print(f"  [SKIP] {video_id}: no sentences found")
                continue
            pd.DataFrame(rows).to_csv(out_csv, index=False)
            n_ok += 1
            print(f"  [OK]   {video_id} -> {out_csv}  ({len(rows)} sentences)")
        except Exception as e:
            n_fail += 1
            print(f"  [FAIL] {video_id}: {e}")

    print(f"\n  Stage 1 complete.  Written: {n_ok}   Failed: {n_fail}")

# Converting MP4 to JPEG format

def stage_video_to_frames(args: argparse.Namespace) -> None:
    print("\n" + "=" * 60)
    print("STAGE 2  Video -> JPEG Frames  (1 fps)")
    print("=" * 60)
    import cv2
    os.makedirs(args.frames_dir, exist_ok=True)
    video_files = [
        f for f in os.listdir(args.video_dir)
        if f.lower().endswith((".mp4", ".avi", ".mov", ".mkv"))
    ]
    if not video_files:
        print(f"  [WARN] No video files found in: {args.video_dir}")
        return
    print(f"  Found {len(video_files)} video files -> frames to: {args.frames_dir}")
    for file in video_files:
        file_mod = os.path.splitext(file)[0]
        out_dir  = os.path.join(args.frames_dir, file_mod)
        if os.path.exists(out_dir) and len(os.listdir(out_dir)) > 0:
            print(f"  [SKIP] {file_mod}: frames already exist ({len(os.listdir(out_dir))} files)")
            continue
        os.makedirs(out_dir, exist_ok=True)
        cap = cv2.VideoCapture(os.path.join(args.video_dir, file))
        if not cap.isOpened():
            print(f"  [FAIL] {file_mod}: could not open video")
            continue
        sec, count, frame_rate = 0, 0, 1
        while cap.isOpened():
            cap.set(cv2.CAP_PROP_POS_MSEC, sec * 1000)
            ret, frame = cap.read()
            if not ret:
                break
            cv2.imwrite(os.path.join(out_dir, f"frame{count}.jpg"), frame)
            count += 1
            sec = round(sec + frame_rate, 2)
        cap.release()
        print(f"  [OK]   {file_mod}: {count} frames extracted")
    print("\n  Stage 2 complete.")


# Downloading pre-trained models locally

def stage_download_models(args: argparse.Namespace) -> None:
    print("\n" + "=" * 60)
    print("STAGE 3  Download Pre-trained Models")
    print("=" * 60)
    from huggingface_hub import snapshot_download
    os.makedirs(CACHE_DIR, exist_ok=True)
    print(f"  Cache dir: {CACHE_DIR}\n")
    for short_name, model_id in ENCODER_MAP.items():
        print(f"  [{short_name}]  {model_id}")
        try:
            snapshot_download(model_id, cache_dir=CACHE_DIR, local_dir_use_symlinks=False)
            print(f"  [OK]   {model_id}")
        except Exception as e:
            print(f"  [FAIL] {model_id}: {e}")
    print("\n  Stage 3 complete.")


# Extracting features using the pre-trained models

def _get_max_len(tokenizer, requested: int = 1024) -> int:
    try:
        return min(int(tokenizer.model_max_length), int(requested))
    except Exception:
        return requested


def _load_encoders(text_model_name: str, vision_model_name: str, local_only: bool = True):
    from transformers import AutoTokenizer, AutoModel, AutoImageProcessor
    print(f"  Loading text encoder  : {text_model_name}")
    tokenizer     = AutoTokenizer.from_pretrained(
                        text_model_name, cache_dir=CACHE_DIR, local_files_only=local_only)
    text_model    = AutoModel.from_pretrained(
                        text_model_name, cache_dir=CACHE_DIR, local_files_only=local_only)
    print(f"  Loading vision encoder: {vision_model_name}")
    vis_processor = AutoImageProcessor.from_pretrained(
                        vision_model_name, cache_dir=CACHE_DIR, local_files_only=local_only)
    vis_model     = AutoModel.from_pretrained(
                        vision_model_name, cache_dir=CACHE_DIR, local_files_only=local_only)
    text_model.eval()
    vis_model.eval()
    print(f"  Text   hidden_size : {getattr(text_model.config, 'hidden_size', 'N/A')}")
    print(f"  Vision hidden_size : {getattr(vis_model.config,  'hidden_size', 'N/A')}")
    return tokenizer, text_model, vis_processor, vis_model


def _save_text_embed(text, save_dir, idx, tokenizer, text_model, max_length=1024):
    import torch
    SPECIAL = ",!?...-_*@#_/\\\"'():;~|^=[]{}<>+=`"
    cleaned = str(text).translate(str.maketrans("", "", SPECIAL))
    inputs  = tokenizer(
        cleaned,
        max_length=_get_max_len(tokenizer, max_length),
        truncation=True,
        return_tensors="pt",
    )
    with torch.no_grad():
        outputs = text_model(**inputs)
    torch.save(outputs.last_hidden_state,
               os.path.join(save_dir, f"text_embed_{idx}.pt"))
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze().tolist())
    with open(os.path.join(save_dir, f"text_tokens_{idx}.json"), "w", encoding="utf-8") as f:
        json.dump(tokens, f)


def _save_frame_embed(img_path, save_dir, frame_idx, vis_processor, vis_model):
    import torch
    from PIL import Image
    save_file = os.path.join(save_dir, f"frame_embed_{frame_idx}.pt")
    if os.path.exists(save_file) or not os.path.exists(img_path):
        return
    try:
        image  = Image.open(img_path).convert("RGB")
        inputs = vis_processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = vis_model(**inputs)
        torch.save(outputs.last_hidden_state, save_file)
    except Exception as e:
        print(f"    [WARN] frame {img_path}: {e}")


def _process_video_features(csv_path, frames_root, embed_root,
                             tokenizer, text_model, vis_processor, vis_model):
    video_id      = os.path.splitext(os.path.basename(csv_path))[0]
    df            = pd.read_csv(csv_path)
    text_dir      = os.path.join(embed_root, video_id, "Text")
    frame_dir_out = os.path.join(embed_root, video_id, "Frames")
    os.makedirs(text_dir,      exist_ok=True)
    os.makedirs(frame_dir_out, exist_ok=True)
    for i in range(len(df)):
        if "sentences" not in df.columns:
            raise ValueError(f"'sentences' column missing in {csv_path}")
        if pd.isna(df.loc[i, "sentences"]) or not str(df.loc[i, "sentences"]).strip():
            continue
        _save_text_embed(str(df.loc[i, "sentences"]), text_dir, i, tokenizer, text_model)
        start = int(0 if pd.isna(df.loc[i, "start_time"]) else df.loc[i, "start_time"])
        end   = int(0 if pd.isna(df.loc[i, "end_time"])   else df.loc[i, "end_time"])
        if end == start:
            end += 1
        end += 1
        for frame_idx in range(start, end):
            _save_frame_embed(
                os.path.join(frames_root, video_id, f"frame{frame_idx}.jpg"),
                frame_dir_out, frame_idx, vis_processor, vis_model,
            )


def stage_extract_features(args, text_model_id, vision_model_id):
    print("\n" + "=" * 60)
    print("STAGE 4  Extract Features")
    print(f"  Text encoder  : {text_model_id}")
    print(f"  Vision encoder: {vision_model_id}")
    print("=" * 60)
    pair       = run_name_for(text_model_id, vision_model_id)
    embed_root = os.path.join("Embeddings", run_path_for(pair, args.prefix), "Intermediate")
    os.makedirs(embed_root, exist_ok=True)
    print(f"  Output root: {embed_root}")
    tokenizer, text_model, vis_processor, vis_model = _load_encoders(
        text_model_id, vision_model_id, local_only=args.local_only)
    filenames    = _read_manifest(args.manifest)
    start        = max(args.shard_start, 0)
    end          = len(filenames) if args.shard_end is None else min(args.shard_end, len(filenames))
    manifest_dir = os.path.dirname(os.path.abspath(args.manifest))
    print(f"  Processing {end - start}/{len(filenames)} videos  [shard {start}:{end}]")
    for i in range(start, end):
        csv_path = _resolve_path(filenames[i], manifest_dir)
        if not csv_path:
            continue
        video_id = os.path.splitext(os.path.basename(csv_path))[0]
        print(f"  [{i+1}/{end}] {video_id}")
        try:
            _process_video_features(csv_path, args.frames_dir, embed_root,
                                     tokenizer, text_model, vis_processor, vis_model)
        except Exception as e:
            print(f"  [ERROR] {video_id}: {e}")
            traceback.print_exc()
    print("\n  Stage 4 complete.")


# Computing Co-Attention using the text and visual pre-trained embeddings

def _single_sentence_image(text_feat, image_feat, common_dim=None):
    import torch
    import torch.nn.functional as F
    if text_feat.dim()  == 2: text_feat  = text_feat.unsqueeze(0)
    if image_feat.dim() == 2: image_feat = image_feat.unsqueeze(0)
    if common_dim is None:
        common_dim = min(text_feat.shape[-1], image_feat.shape[-1])
    text_feat  = _project(text_feat,  common_dim)
    image_feat = _project(image_feat, common_dim)
    scale      = common_dim ** 0.5
    t2i_raw    = torch.matmul(text_feat,  image_feat.transpose(1, 2)) / scale
    i2t_raw    = torch.matmul(image_feat, text_feat.transpose(1,  2)) / scale
    att_text   = torch.matmul(F.softmax(i2t_raw, dim=-1), text_feat)
    att_image  = torch.matmul(F.softmax(t2i_raw, dim=-1), image_feat)
    return att_text.squeeze(0), att_image.squeeze(0), t2i_raw.squeeze(0), i2t_raw.squeeze(0)


def _combined_coattn(attended_text, attended_image):
    import torch
    import torch.nn.functional as F
    attn_w      = F.softmax(attended_image.mean(dim=1, keepdim=True), dim=0)
    pooled_text = torch.sum(attn_w * attended_image, dim=0, keepdim=True).expand(
                      *attended_text.shape)
    gate        = torch.sigmoid(pooled_text + attended_text)
    return (gate * pooled_text + (1 - gate) * attended_text).mean(dim=0)


def _norm_t2f(attn_list):
    import torch
    mat = torch.cat(attn_list, dim=1)
    return mat / (mat.sum(dim=1, keepdim=True) + 1e-9)


def _norm_f2t(attn_list):
    import torch
    mat = torch.cat(attn_list, dim=0)
    return mat / (mat.sum(dim=1, keepdim=True) + 1e-9)


def _process_coattention(csv_path, embed_root, out_embed_root, attn_root, common_dim=None):
    import torch
    video_id   = os.path.splitext(os.path.basename(csv_path))[0]
    df         = pd.read_csv(csv_path)
    text_base  = os.path.join(embed_root, video_id, "Text")
    frame_base = os.path.join(embed_root, video_id, "Frames")
    out_vid    = os.path.join(out_embed_root, video_id)
    for sub in ("T2F", "F2T", "Co_Attn"):
        os.makedirs(os.path.join(out_vid, sub), exist_ok=True)
    attn_vid = os.path.join(attn_root, video_id)
    os.makedirs(os.path.join(attn_vid, "F2T"), exist_ok=True)
    os.makedirs(os.path.join(attn_vid, "T2F"), exist_ok=True)
    count = 0
    for i in range(len(df)):
        if pd.isna(df.loc[i, "sentences"]) or not str(df.loc[i, "sentences"]).strip():
            continue
        try:
            text_feat = _safe_load(os.path.join(text_base, f"text_embed_{i}.pt"))
            with open(os.path.join(text_base, f"text_tokens_{i}.json"), "r", encoding="utf-8") as f:
                token_labels = json.load(f)
        except FileNotFoundError:
            print(f"    [SKIP] {video_id}: missing text embeddings row {i}")
            continue
        start = int(0 if pd.isna(df.loc[i, "start_time"]) else df.loc[i, "start_time"])
        end   = int(0 if pd.isna(df.loc[i, "end_time"])   else df.loc[i, "end_time"])
        if end == start: end += 1
        end += 1
        t2f_emb, f2t_emb, co_emb = [], [], []
        t2f_attn, f2t_attn = [], []
        valid = False
        for frame_idx in range(start, end):
            fp = os.path.join(frame_base, f"frame_embed_{frame_idx}.pt")
            if not os.path.exists(fp):
                continue
            try:
                img_feat = _safe_load(fp)
                valid    = True
                att_text, att_img, t2i_raw, i2t_raw = _single_sentence_image(
                    text_feat, img_feat, common_dim)
                t2f_emb.append(att_img[0].unsqueeze(0))
                f2t_emb.append(att_text[0].unsqueeze(0))
                co_emb.append(_combined_coattn(att_text, att_img).unsqueeze(0))
                t2f_attn.append(t2i_raw)
                f2t_attn.append(i2t_raw)
            except Exception as e:
                print(f"    [WARN] {video_id} frame {frame_idx}: {e}")
        if not valid:
            count += 1
            continue
        try:
            pd.DataFrame(
                _norm_f2t(f2t_attn).detach().cpu().numpy(), columns=token_labels
            ).to_csv(os.path.join(attn_vid, "F2T", f"frame_to_text{count}.csv"), index=False)
            pd.DataFrame(
                _norm_t2f(t2f_attn).detach().cpu().numpy(), index=token_labels
            ).to_csv(os.path.join(attn_vid, "T2F", f"text_to_frame{count}.csv"), index=True)
        except Exception as e:
            print(f"    [WARN] {video_id}: attn CSV write failed (sentence {count}): {e}")
        torch.save(_weighted_sharing(t2f_emb), os.path.join(out_vid, "T2F",     f"CLS{count}.pt"))
        torch.save(_weighted_sharing(f2t_emb), os.path.join(out_vid, "F2T",     f"CLS{count}.pt"))
        torch.save(_weighted_sharing(co_emb),  os.path.join(out_vid, "Co_Attn", f"CLS{count}.pt"))
        print(f"    [OK] {video_id}: sentence {count} saved")
        count += 1


def stage_compute_coattention(args, text_model_id, vision_model_id):
    print("\n" + "=" * 60)
    print("STAGE 5  Compute Co-Attention")
    print("=" * 60)
    pair      = run_name_for(text_model_id, vision_model_id)
    embed_dir = os.path.join("Embeddings", run_path_for(pair, args.prefix), "Intermediate")
    attn_dir  = os.path.join("data", "Attention_Files", run_path_for(pair, args.prefix))
    os.makedirs(attn_dir, exist_ok=True)
    filenames    = _read_manifest(args.manifest)
    manifest_dir = os.path.dirname(os.path.abspath(args.manifest))
    print(f"  embed_dir  : {embed_dir}")
    print(f"  attn_dir   : {attn_dir}")
    print(f"  common_dim : {args.common_dim}")
    for i, csv_path in enumerate(filenames, 1):
        csv_path = _resolve_path(csv_path, manifest_dir)
        if not csv_path:
            continue
        video_id = os.path.splitext(os.path.basename(csv_path))[0]
        print(f"  [{i}/{len(filenames)}] {video_id}")
        try:
            _process_coattention(csv_path, embed_dir, embed_dir, attn_dir,
                                  common_dim=args.common_dim)
        except Exception as e:
            print(f"  [ERROR] {video_id}: {e}")
            traceback.print_exc()
    print("\n  Stage 5 complete.")


# Simple fusion (concat and mean)

def _fuse_simple(text_vec, frame_vec, mode: str, common_dim=None):
    import torch
    cdim      = common_dim or min(text_vec.shape[-1], frame_vec.shape[-1])
    text_vec  = _project(text_vec,  cdim)
    frame_vec = _project(frame_vec, cdim)
    if mode == "mean":   return (text_vec + frame_vec) / 2.0
    if mode == "concat": return torch.cat([text_vec, frame_vec], dim=1)
    raise ValueError(f"mode must be 'mean' or 'concat', got '{mode}'")


def _process_simple_fuse(csv_path, embed_root, out_embed_root, mode, common_dim=None):
    import torch
    video_id   = os.path.splitext(os.path.basename(csv_path))[0]
    df         = pd.read_csv(csv_path)
    text_base  = os.path.join(embed_root, video_id, "Text")
    frame_base = os.path.join(embed_root, video_id, "Frames")
    out_dir    = os.path.join(out_embed_root, video_id, mode)
    os.makedirs(out_dir, exist_ok=True)
    count = 0
    for i in range(len(df)):
        if pd.isna(df.loc[i, "sentences"]) or not str(df.loc[i, "sentences"]).strip():
            continue
        try:
            t = torch.load(os.path.join(text_base, f"text_embed_{i}.pt"),
                           map_location="cpu", weights_only=True)
            if t.ndim == 3: t = t[0]
            text_cls = t[0].unsqueeze(0)
        except FileNotFoundError:
            print(f"    [SKIP] {video_id}: missing text_embed_{i}.pt")
            continue
        start = int(0 if pd.isna(df.loc[i, "start_time"]) else df.loc[i, "start_time"])
        end   = int(0 if pd.isna(df.loc[i, "end_time"])   else df.loc[i, "end_time"])
        if end == start: end += 1
        end += 1
        fused_frames, valid = [], False
        for frame_idx in range(start, end):
            fp = os.path.join(frame_base, f"frame_embed_{frame_idx}.pt")
            if not os.path.exists(fp):
                continue
            try:
                fr = torch.load(fp, map_location="cpu", weights_only=True)
                if fr.ndim == 3: fr = fr[0]
                fused_frames.append(
                    _fuse_simple(text_cls, fr[0].unsqueeze(0), mode, common_dim))
                valid = True
            except Exception as e:
                print(f"    [WARN] {video_id} frame {frame_idx}: {e}")
        if not valid or not fused_frames:
            count += 1
            continue
        torch.save(_weighted_sharing(fused_frames), os.path.join(out_dir, f"CLS{count}.pt"))
        print(f"    [OK] {video_id}: {mode} sentence {count} saved")
        count += 1


def stage_simple_fuse(args, text_model_id, vision_model_id):
    print("\n" + "=" * 60)
    print("STAGE 6  Simple Fuse  (Ablation: mean + concat)")
    print("=" * 60)
    pair         = run_name_for(text_model_id, vision_model_id)
    embed_dir    = os.path.join("Embeddings", run_path_for(pair, args.prefix), "Intermediate")
    filenames    = _read_manifest(args.manifest)
    manifest_dir = os.path.dirname(os.path.abspath(args.manifest))
    for mode in ("mean", "concat"):
        print(f"\n  Fusion mode: {mode}")
        for i, csv_path in enumerate(filenames, 1):
            csv_path = _resolve_path(csv_path, manifest_dir)
            if not csv_path:
                continue
            vid = os.path.splitext(os.path.basename(csv_path))[0]
            print(f"    [{i}/{len(filenames)}] {vid}")
            try:
                _process_simple_fuse(csv_path, embed_dir, embed_dir,
                                      mode, common_dim=args.common_dim)
            except Exception as e:
                print(f"    [ERROR] {vid}: {e}")
    print("\n  Stage 6 complete.")

# Final video embeddings (pooling)

def _pool_non_cls(t, modality: str):
    import torch
    t = t.detach().to("cpu")
    if t.ndim == 3 and t.shape[0] > 1:  t = t.mean(dim=0, keepdim=True)
    if t.ndim == 3 and t.shape[0] == 1: t = t.squeeze(0)
    if t.ndim == 1: return t.unsqueeze(0)
    if t.ndim != 2: raise ValueError(f"Unexpected tensor shape {tuple(t.shape)}")
    L, _ = t.shape
    if modality.lower() == "text"   and L > 2: t = t[1:-1, :]
    if modality.lower() == "frames" and L > 1: t = t[1:,    :]
    return t.mean(dim=0, keepdim=True)


def _fuse_vectors(stacked, mode: str, seed: int = 42):
    import torch
    import torch.nn.functional as F
    if stacked.shape[0] == 1:
        return stacked
    if mode == "mean":
        return stacked.mean(dim=0, keepdim=True)
    if mode == "softmax":
        w = F.softmax(stacked.mean(dim=1), dim=0)
        return (stacked * w.unsqueeze(1)).sum(dim=0, keepdim=True)
    if mode == "attn":
        g = torch.Generator(device=stacked.device)
        g.manual_seed(seed)
        q = torch.randn(stacked.shape[1], generator=g, device=stacked.device)
        w = F.softmax(stacked @ q, dim=0)
        return (stacked * w.unsqueeze(1)).sum(dim=0, keepdim=True)
    raise ValueError(f"Unknown fusion mode: {mode}")


def _fuse_one_modality(root_in, root_out, modality,
                       fusion="softmax", seed=42, verbose=False):
    import torch
    out_dir = os.path.join(root_out, fusion, modality)
    os.makedirs(out_dir, exist_ok=True)

    for entry in os.scandir(root_in):
        if not entry.is_dir():
            continue
        in_dir = os.path.join(root_in, entry.name, modality)
        if not os.path.isdir(in_dir):
            if verbose: print(f"    [SKIP] Missing: {in_dir}")
            continue
        files = sorted([f for f in os.listdir(in_dir) if f.endswith(".pt")])
        if not files:
            if verbose: print(f"    [SKIP] No .pt files in: {in_dir}")
            continue
        tensors, ok = [], True
        for fname in files:
            try:
                tensors.append(torch.load(os.path.join(in_dir, fname),
                                          map_location="cpu", weights_only=True))
            except Exception as e:
                ok = False
                print(f"    [WARN] Failed loading {fname}: {e}")
                break
        if not ok:
            continue
        vecs    = [_pool_non_cls(t, modality) for t in tensors]
        stacked = torch.cat(vecs, dim=0)
        fused   = _fuse_vectors(stacked, mode=fusion, seed=seed)
        out_fp  = os.path.join(out_dir, f"{entry.name}.pt")
        torch.save(fused, out_fp)
        if verbose: print(f"    [OK] {fusion}/{modality}: {entry.name} -> {out_fp}")


def stage_fuse_embeddings(args, text_model_id, vision_model_id):
    print("\n" + "=" * 60)
    print("STAGE 7  Fuse Embeddings  (Final pooling per video)")
    print("=" * 60)
    pair     = run_name_for(text_model_id, vision_model_id)
    in_root  = os.path.join("Embeddings", run_path_for(pair, args.prefix), "Intermediate")
    out_root = os.path.join("Embeddings", run_path_for(pair, args.prefix), "Final")
    os.makedirs(out_root, exist_ok=True)
    modalities = ["F2T", "T2F", "Co_Attn", "Text", "Frames", "mean", "concat"]
    fusion_modes = ["softmax", "mean", "attn"]
    print(f"  in_root  : {in_root}")
    print(f"  out_root : {out_root}")
    print(f"  fusion   : {args.fusion}")
    for fusion in fusion_modes:
        print(f"\n  ── Fusion mode: {fusion} ──")
        for m in modalities:
            print(f"    Fusing: {m}")
            _fuse_one_modality(in_root, out_root, m,
                               fusion=fusion, seed=args.seed, verbose=args.verbose)
    print("\n  Stage 7 complete.")


# Calculating entropy scores

def _row_normalize(x) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
    exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
    denom = exp_x.sum(axis=1, keepdims=True)
    denom = np.where(denom == 0, 1.0, denom)
    return exp_x / denom


def _compute_mean_entropy(matrix, eps: float = 1e-12) -> float:
    from scipy.stats import entropy as scipy_entropy

    if matrix.size == 0:
        return 0.0

    matrix = np.nan_to_num(matrix, nan=0.0, posinf=0.0, neginf=0.0)

    matrix = np.clip(matrix, eps, 1)

    row_sums = matrix.sum(axis=1, keepdims=True)
    row_sums = np.where(row_sums == 0, 1.0, row_sums)
    probs = matrix / row_sums

    ent = scipy_entropy(probs.T)
    ent = ent[np.isfinite(ent)]
    if len(ent) == 0:
        return 0.0

    return float(np.mean(ent))


def stage_entropy_scores(args, text_model_id, vision_model_id):
    print("\n" + "=" * 60)
    print("STAGE 8  Entropy Scores")
    print("=" * 60)
    pair     = run_name_for(text_model_id, vision_model_id)
    attn_in  = os.path.join("data", "Attention_Files",  run_path_for(pair, args.prefix))
    attn_out = os.path.join("data", "Attention_Scores", run_path_for(pair, args.prefix))
    os.makedirs(attn_out, exist_ok=True)
    print(f" Input  : {attn_in}")
    print(f" Output : {attn_out}")

    records = []
    nan_videos = []  

    for folder in sorted(os.listdir(attn_in)):
        t2f_path = os.path.join(attn_in, folder, "T2F")
        f2t_path = os.path.join(attn_in, folder, "F2T")
        if not (os.path.isdir(t2f_path) and os.path.isdir(f2t_path)):
            continue

        t2f_ents, f2t_ents = [], []
        n_files = len(os.listdir(t2f_path))

        for i in range(n_files):
            df_t2f = df_f2t = None
            try:
                df_t2f = pd.read_csv(os.path.join(t2f_path, f"text_to_frame{i}.csv"))
                df_f2t = pd.read_csv(os.path.join(f2t_path, f"frame_to_text{i}.csv"))

                if "Unnamed: 0" in df_t2f.columns:
                    df_t2f = df_t2f.drop(columns=["Unnamed: 0"])
                if "Unnamed: 0" in df_f2t.columns:
                    df_f2t = df_f2t.drop(columns=["Unnamed: 0"])

                t2f_val = _compute_mean_entropy(_row_normalize(df_t2f.values))
                f2t_val = _compute_mean_entropy(_row_normalize(df_f2t.values))
                t2f_ents.append(t2f_val if np.isfinite(t2f_val) else 0.0)
                f2t_ents.append(f2t_val if np.isfinite(f2t_val) else 0.0)

            except Exception:
                try:
                    t2f_ents.append(float(np.log(max(df_t2f.shape[1], 1)))
                                    if df_t2f is not None else 0.0)
                    f2t_ents.append(float(np.log(max(df_f2t.shape[1], 1)))
                                    if df_f2t is not None else 0.0)
                except Exception:
                    t2f_ents.append(0.0)
                    f2t_ents.append(0.0)

        ent_t2f = float(np.mean(t2f_ents)) if t2f_ents else 0.0
        ent_f2t = float(np.mean(f2t_ents)) if f2t_ents else 0.0

        if not np.isfinite(ent_t2f) or not np.isfinite(ent_f2t):
            nan_videos.append({
                "video_id":    folder,
                "entropy_t2f": ent_t2f,
                "entropy_f2t": ent_f2t,
            })
            ent_t2f = 0.0 if not np.isfinite(ent_t2f) else ent_t2f
            ent_f2t = 0.0 if not np.isfinite(ent_f2t) else ent_f2t
            print(f"  [WARN] {folder}: NaN entropy detected — filled with 0.0")
        else:
            print(f" [OK] {folder}  entropy_t2f={ent_t2f:.4f}  entropy_f2t={ent_f2t:.4f}")

        records.append([folder, ent_t2f, ent_f2t])

    out_csv = os.path.join(attn_out, "attention_scores.csv")
    pd.DataFrame(records, columns=["video_id", "entropy_t2f", "entropy_f2t"]).to_csv(
        out_csv, index=False)
    print(f"\n  Saved : {out_csv}")

    if nan_videos:
        nan_csv = os.path.join(attn_out, "nan_entropy_report.csv")
        pd.DataFrame(nan_videos).to_csv(nan_csv, index=False)
        print(f" [WARN] {len(nan_videos)} videos had NaN entropy.")
        print(f" Filled with 0.0 in attention_scores.csv")
        print(f" Full report saved : {nan_csv}")
    else:
        print(f" No NaN entropy values detected.")

    print("\n  Stage 8 complete.")


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Full data pipeline: Raw video + JSON -> Embeddings",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("--stages", nargs="+", default=["all"],
                   choices=ALL_STAGES + ["all"],
                   help="Stage(s) to run. Use 'all' for the full pipeline.")
    p.add_argument("--text_encoder",   default="bert",
                   help="bert | longformer | full HF id  (default: bert)")
    p.add_argument("--vision_encoder", default="vit",
                   help="vit | dinov2 | full HF id       (default: vit)")
    p.add_argument("--prefix", default=None,
                   help="Optional prefix inserted before pair name in all output paths. "
                   "Use to keep datasets separate. "
                   "e.g. --prefix COVID saves to "
                   "Embeddings/COVID/BERT_ViT/... and "
                   "data/Attention_Scores/COVID/BERT_ViT/...")
    p.add_argument("--json_dir",    default="data/JSON")
    p.add_argument("--csv_dir",     default="data/Transcripts_Shots")
    p.add_argument(
    "--segmentation",
    default="shots",
    choices=["shots", "sentences"],
    help=(
        "Segmentation method for Stage 1:\n"
        "  shots -> screen change boundaries (our method) -> Transcripts_Shots/\n"
        "  sentences -> NLTK sentence tokenizer  (ablation)   -> Transcripts/\n"
        "(default: shots)"
    ),
    )
    p.add_argument("--video_dir",   default="data/Videos")
    p.add_argument("--frames_dir",  default="data/Video_Frames")
    p.add_argument("--manifest",    default="data/video_manifest.txt")
    p.add_argument("--common_dim",  type=int, default=768)
    p.add_argument("--fusion",      default="softmax",
                   choices=["softmax", "mean", "attn"])
    p.add_argument("--seed",        type=int, default=42)
    p.add_argument("--verbose",     action="store_true")
    p.add_argument("--shard_start", type=int, default=0)
    p.add_argument("--shard_end",   type=int, default=None)
    p.add_argument("--local_only",    action="store_true", default=True)
    p.add_argument("--no-local_only", dest="local_only", action="store_false")
    return p

def main():
    args = build_parser().parse_args()

    _DEFAULT_CSV_DIR  = "data/Transcripts_Shots"
    _DEFAULT_MANIFEST = "data/video_manifest.txt"

    if args.csv_dir == _DEFAULT_CSV_DIR and args.segmentation == "sentences":
        args.csv_dir = "data/Transcripts"
        print(f"[INFO] Segmentation=sentences -> csv_dir auto-set to: {args.csv_dir}")

    if args.segmentation == "sentences" and os.path.basename(args.manifest) == "video_manifest.txt":
        args.manifest = os.path.join(os.path.dirname(args.manifest), "video_manifest_sentences.txt")
        print(f"[INFO] Segmentation=sentences -> manifest auto-set to: {args.manifest}")
        
    if args.segmentation == "sentences" and args.prefix is None:
        args.prefix = "Baseline"
        print(f"[INFO] Segmentation=sentences -> prefix auto-set to: {args.prefix}")
        
    text_model_id   = resolve_encoder(args.text_encoder,   kind="text")
    vision_model_id = resolve_encoder(args.vision_encoder, kind="vision")
    stages          = ALL_STAGES if "all" in args.stages else args.stages
    pair            = run_name_for(text_model_id, vision_model_id)

    print("\n" + "=" * 60)
    print(f"  DATA PIPELINE  |  Encoder pair : {pair}")
    print(f"  Stages         : {', '.join(stages)}")
    print("=" * 60)

    downstream = {"extract_features", "compute_coattention",
                  "simple_fuse", "fuse_embeddings", "entropy_scores"}
    if any(s in stages for s in downstream) and not os.path.exists(args.manifest):
        print(f"\n[INFO] Manifest not found -> auto-building from {args.csv_dir}")
        _build_manifest(args.csv_dir, args.manifest)

    if "json_to_csv" in stages:
        if args.segmentation == "sentences":
            stage_json_to_csv_sentences(args)
        else:
            stage_json_to_csv(args)
    if "video_to_frames"     in stages: stage_video_to_frames(args)
    if "download_models"     in stages: stage_download_models(args)
    if "extract_features"    in stages: stage_extract_features(args, text_model_id, vision_model_id)
    if "compute_coattention" in stages: stage_compute_coattention(args, text_model_id, vision_model_id)
    if "simple_fuse"         in stages: stage_simple_fuse(args, text_model_id, vision_model_id)
    if "fuse_embeddings"     in stages: stage_fuse_embeddings(args, text_model_id, vision_model_id)
    if "entropy_scores"      in stages: stage_entropy_scores(args, text_model_id, vision_model_id)

    print("\n" + "=" * 60)
    print("  PIPELINE COMPLETE")
    print("=" * 60 + "\n")


if __name__ == "__main__":
    main()