#!/usr/bin/env python3
import os
import argparse
import torch
import torch.nn.functional as F


def pool_non_cls(t: torch.Tensor, modality: str) -> torch.Tensor:
    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)} for modality={modality}")

    L, D = t.shape

    if modality.lower() == "text":
        if L > 2:
            t = t[1:-1, :]  
    elif modality.lower() == "frames":
        if L > 1:
            t = t[1:, :]  
    else:
        pass

    return t.mean(dim=0, keepdim=True)  

def attention_fusion(stacked: torch.Tensor, seed: int = 42) -> torch.Tensor:
    N, D = stacked.shape
    g = torch.Generator(device=stacked.device)
    g.manual_seed(seed)

    q = torch.randn(D, generator=g, device=stacked.device)  
    scores = stacked @ q 
    weights = F.softmax(scores, dim=0)  
    return (stacked * weights.unsqueeze(1)).sum(dim=0, keepdim=True)  

def fuse_vectors(stacked: torch.Tensor, mode: str, seed: int = 42) -> torch.Tensor:

    mode = mode.lower()
    if stacked.shape[0] == 1:
        return stacked  

    if mode == "mean":
        return stacked.mean(dim=0, keepdim=True)

    if mode == "softmax":
        scores = stacked.mean(dim=1)          
        weights = F.softmax(scores, dim=0)    
        return (stacked * weights.unsqueeze(1)).sum(dim=0, keepdim=True)

    if mode == "attn":
        return attention_fusion(stacked, seed=seed)

    raise ValueError(f"Unknown fusion mode: {mode}")

def weighted_sharing(tensors, modality: str, fusion: str = "softmax", seed: int = 42):
    vecs = [pool_non_cls(t, modality) for t in tensors]  
    stacked = torch.cat(vecs, dim=0)  
    return fuse_vectors(stacked, mode=fusion, seed=seed)

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

    for entry in os.scandir(root_in):
        if not entry.is_dir():
            continue

        vid = entry.name
        in_dir = os.path.join(root_in, vid, modality)
        if not os.path.isdir(in_dir):
            if verbose:
                print(f"[SKIP] Missing modality folder: {in_dir}")
            continue

        files = sorted([f for f in os.listdir(in_dir) if f.endswith(ext)])
        if not files:
            if verbose:
                print(f"[SKIP] No {ext} files in: {in_dir}")
            continue

        tensors = []
        ok = True
        for f in files:
            fp = os.path.join(in_dir, f)
            try:
                tensors.append(torch.load(fp, map_location="cpu", weights_only=True))
            except Exception as e:
                ok = False
                print(f"[WARN] Failed loading {fp}: {e}")
                break

        if not ok:
            continue

        fused = weighted_sharing(tensors, modality, fusion=fusion, seed=seed)
        out_fp = os.path.join(out_dir, f"{vid}{ext}")
        torch.save(fused, out_fp)

        if verbose:
            print(f"[OK] {modality}: {vid} -> {out_fp}")


def main():
    parser = argparse.ArgumentParser(
        description="Fuse embeddings per videoid into a single embedding per modality."
    )

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--F2T", action="store_true", help="Fuse only F2T embeddings")
    group.add_argument("--T2F", action="store_true", help="Fuse only T2F embeddings")
    group.add_argument("--Co_Attn", action="store_true", help="Fuse only Co_Attn embeddings")
    group.add_argument("--Text", action="store_true", help="Fuse only TEXT embeddings (non-CLS pooled)")
    group.add_argument("--Frames", action="store_true", help="Fuse only FRAME embeddings (non-CLS pooled)")
    group.add_argument("--mean", action="store_true", help="Fuse late fused embeddings")
    group.add_argument("--concat", action="store_true", help="Fuse early fused embeddings")
    group.add_argument("--all", action="store_true", help="Fuse all modalities")

    parser.add_argument("--in-root", required=True, help="Root input folder (e.g., .../Embeddings/Intermediate)")
    parser.add_argument("--out-root", required=True, help="Root output folder (e.g., .../Embeddings/Final)")
    parser.add_argument("--ext", default=".pt", help="File extension to load/save (default: .pt)")
    parser.add_argument(
    "--fusion",
    type=str,
    default="softmax",
    choices=["softmax", "mean", "attn"],
    help="Fusion method across segments: softmax (default), mean, or attn (fixed attention pooling)."
    )
    parser.add_argument(
    "--seed",
    type=int,
    default=42,
    help="Seed used for deterministic attention pooling when --fusion attn."
    )
    parser.add_argument("--verbose", action="store_true", help="Print progress")

    args = parser.parse_args()

    if args.all:
        modalities = ["F2T", "T2F", "Co_Attn", "Text", "Frames", "mean", "concat"]
    else:
        modalities = []
        if args.F2T: modalities.append("F2T")
        if args.T2F: modalities.append("T2F")
        if args.Co_Attn: modalities.append("Co_Attn")
        if args.Text: modalities.append("Text")
        if args.Frames: modalities.append("Frames")

    for m in modalities:
        if args.verbose:
            print(f"\n=== Fusing modality: {m} ===")
        fuse_one_modality(args.in_root, args.out_root, m, ext=args.ext, fusion=args.fusion, seed=args.seed, verbose=args.verbose)

    if args.verbose:
        print("\nDone.")


if __name__ == "__main__":
    main()
