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

import os
import argparse
import torch
import pandas as pd

import torch.nn.functional as F

_PROJ_CACHE = {}

def get_proj(in_dim: int, out_dim: int, device: torch.device):
    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)
    W = W.to(device)
    _PROJ_CACHE[key] = W
    return W

def project_vec(x: torch.Tensor, out_dim: int) -> torch.Tensor:
    in_dim = x.shape[-1]
    if in_dim == out_dim:
        return x
    W = get_proj(in_dim, out_dim, x.device)
    return F.linear(x, W)  

def weighted_sharing(emb_list):
    stacked = torch.cat(emb_list, dim=0)          
    scores = stacked.mean(dim=1)                 
    weights = torch.softmax(scores, dim=0)        
    return (stacked * weights.unsqueeze(1)).sum(dim=0, keepdim=True)  


def load_text_cls(text_embed_base: str, row_i: int) -> torch.Tensor:
    t = torch.load(os.path.join(text_embed_base, f"text_embed_{row_i}.pt"), map_location="cpu", weights_only=True)
    if t.ndim == 3:
        t = t[0]
    if t.ndim != 2:
        raise ValueError(f"Unexpected text tensor shape: {tuple(t.shape)}")
    cls = t[0].unsqueeze(0)  
    return cls


def load_frame_cls(frame_embed_base: str, frame_idx: int) -> torch.Tensor:
    t = torch.load(os.path.join(frame_embed_base, f"frame_embed_{frame_idx}.pt"), map_location="cpu", weights_only=True)
    if t.ndim == 3:
        t = t[0]
    if t.ndim != 2:
        raise ValueError(f"Unexpected frame tensor shape: {tuple(t.shape)}")
    cls = t[0].unsqueeze(0)  
    return cls


def fuse(text_vec: torch.Tensor, frame_vec: torch.Tensor, mode: str, common_dim: int = None) -> torch.Tensor:
    mode = mode.lower()

    Dt = text_vec.shape[-1]
    Df = frame_vec.shape[-1]
    if common_dim is None:
        common_dim = min(Dt, Df)

    text_vec = project_vec(text_vec, common_dim)
    frame_vec = project_vec(frame_vec, common_dim)

    if mode == "mean":
        return (text_vec + frame_vec) / 2.0  
    if mode == "concat":
        return torch.cat([text_vec, frame_vec], dim=1)  
    raise ValueError("mode must be 'mean' or 'concat'")


def process_one_csv(csv_path: str, embed_root: str, out_embed_root: str, mode: str, common_dim: int = None):
    video_id = os.path.splitext(os.path.basename(csv_path))[0]
    df = pd.read_csv(csv_path)

    vid_embed_path = os.path.join(embed_root, video_id)
    text_embed_base = os.path.join(vid_embed_path, "Text")
    frame_embed_base = os.path.join(vid_embed_path, "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)):
        # skip empty sentences
        if pd.isna(df.loc[i, "sentences"]) or len(str(df.loc[i, "sentences"]).strip()) == 0:
            continue

        try:
            text_cls = load_text_cls(text_embed_base, i)  
        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_per_frame = []
        valid_frames = False

        for frame_idx in range(start, end):
            frame_file = os.path.join(frame_embed_base, f"frame_embed_{frame_idx}.pt")
            if not os.path.exists(frame_file):
                continue

            try:
                frame_cls = load_frame_cls(frame_embed_base, frame_idx)  
                valid_frames = True
                fused_per_frame.append(fuse(text_cls, frame_cls, mode, common_dim=common_dim))   
            except Exception as e:
                print(f"[WARN] {video_id}: frame {frame_idx} failed: {e}")
                continue

        if not valid_frames or len(fused_per_frame) == 0:
            count += 1
            continue

        fused_sentence = weighted_sharing(fused_per_frame)  

        torch.save(fused_sentence, os.path.join(out_dir, f"CLS{count}.pt"))
        print(f"[OK] {video_id}: saved {mode} sentence {count}")
        count += 1


def main():
    parser = argparse.ArgumentParser(description="Fuse TEXT+FRAME embeddings using mean or concat (no attention CSVs).")
    parser.add_argument("--manifest", type=str, required=True, help="Path to video_manifest.txt (one CSV per line)")
    parser.add_argument("--embed_dir", type=str, required=True,
                        help="Root input embeddings folder containing <video_id>/Text and <video_id>/Frames")
    parser.add_argument("--out_embed_dir", type=str, required=True,
                        help="Output embeddings folder (writes <video_id>/<mode>/CLS*.pt)")
    parser.add_argument("--mode", type=str, required=True, choices=["mean", "concat"],
                        help="Fusion mode: mean or concat")
    parser.add_argument("--common_dim", type=int, default=None,
                    help="Project text+frame embeddings to this dim before fusion. Default: min(text_dim, frame_dim).")
    args = parser.parse_args()

    with open(args.manifest, "r", encoding="utf-8") as f:
        filenames = [line.strip() for line in f if line.strip() and not line.strip().startswith("#")]

    if not filenames:
        print(f"No files found in manifest: {args.manifest}")
        return

    os.makedirs(args.out_embed_dir, exist_ok=True)

    manifest_dir = os.path.dirname(os.path.abspath(args.manifest))
    print(f"Found {len(filenames)} files in manifest.")
    print(f"embed_dir     : {args.embed_dir}")
    print(f"out_embed_dir : {args.out_embed_dir}")
    print(f"mode          : {args.mode}")

    for k, csv_path in enumerate(filenames, start=1):
        if not os.path.isabs(csv_path):
            candidate = os.path.join(manifest_dir, csv_path)
            if os.path.exists(candidate):
                csv_path = candidate

        if not os.path.exists(csv_path):
            print(f"[SKIP] Missing CSV: {csv_path}")
            continue

        vid = os.path.splitext(os.path.basename(csv_path))[0]
        print(f"\n[{k}/{len(filenames)}] Processing: {vid}")
        try:
            process_one_csv(csv_path, args.embed_dir, args.out_embed_dir, args.mode, common_dim=args.common_dim)
        except Exception as e:
            print(f"[ERROR] {vid} failed: {e}")


if __name__ == "__main__":
    main()