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

import os
import json
import argparse
import torch
import torch.nn.functional as F
import pandas as pd


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

_PROJ_CACHE = {}

def get_projection(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_to_dim(x: torch.Tensor, out_dim: int) -> torch.Tensor:
    in_dim = x.shape[-1]
    if in_dim == out_dim:
        return x
    W = get_projection(in_dim, out_dim, x.device)  
    return F.linear(x, W)


def single_sentence_image(text_features, image_features, common_dim: int = None):
    if text_features.dim() == 2:
        text_features = text_features.unsqueeze(0)
    if image_features.dim() == 2:
        image_features = image_features.unsqueeze(0)

    Dt = text_features.shape[-1]
    Di = image_features.shape[-1]

    if common_dim is None:
        common_dim = min(Dt, Di)

    text_features = project_to_dim(text_features, common_dim)
    image_features = project_to_dim(image_features, common_dim)

    scale = common_dim ** 0.5

    text_to_image_attn_raw = torch.matmul(text_features, image_features.transpose(1, 2)) / scale
    image_to_text_attn_raw = torch.matmul(image_features, text_features.transpose(1, 2)) / scale

    text_to_image_attn = F.softmax(text_to_image_attn_raw, dim=-1)
    image_to_text_attn = F.softmax(image_to_text_attn_raw, dim=-1)

    attended_text = torch.matmul(image_to_text_attn, text_features)   
    attended_image = torch.matmul(text_to_image_attn, image_features) 

    return (
        attended_text.squeeze(0),          
        attended_image.squeeze(0),         
        text_to_image_attn_raw.squeeze(0), 
        image_to_text_attn_raw.squeeze(0), 
    )

def combined_coattn(attended_text, attended_image):
    attn_scores = attended_image.mean(dim=1, keepdim=True)  
    attn_weights = F.softmax(attn_scores, dim=0)            
    pooled_text = torch.sum(attn_weights * attended_image, dim=0, keepdim=True)  

    P, D = attended_text.shape
    pooled_text = pooled_text.expand(P, D)  
    gate = torch.sigmoid(pooled_text + attended_text)       
    combined = gate * pooled_text + (1 - gate) * attended_text

    return combined.mean(dim=0) 

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

def norm_t2f(attn_list):
    text_frame_matrix = torch.cat(attn_list, dim=1)  
    row_sums = text_frame_matrix.sum(dim=1, keepdim=True)
    return text_frame_matrix / (row_sums + 1e-9)

def norm_f2t(attn_list):
    frame_text_matrix = torch.cat(attn_list, dim=0)  
    row_sums = frame_text_matrix.sum(dim=1, keepdim=True)
    return frame_text_matrix / (row_sums + 1e-9)


def process_coattention(csv_path, embed_root, out_embed_root, attn_root, common_dim=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_vid_dir = os.path.join(out_embed_root, video_id)
    t2f_out_dir = os.path.join(out_vid_dir, "T2F")
    f2t_out_dir = os.path.join(out_vid_dir, "F2T")
    coattn_out_dir = os.path.join(out_vid_dir, "Co_Attn")
    os.makedirs(t2f_out_dir, exist_ok=True)
    os.makedirs(f2t_out_dir, exist_ok=True)
    os.makedirs(coattn_out_dir, exist_ok=True)

    attn_save_path = os.path.join(attn_root, video_id)
    text_out = os.path.join(attn_save_path, "F2T")
    frame_out = os.path.join(attn_save_path, "T2F")
    os.makedirs(text_out, exist_ok=True)
    os.makedirs(frame_out, exist_ok=True)

    count = 0

    for i in range(len(df)):
        if pd.isna(df.loc[i, "sentences"]) or len(str(df.loc[i, "sentences"]).strip()) == 0:
            continue

        try:
            text_features = safe_torch_load(os.path.join(text_embed_base, f"text_embed_{i}.pt"))
            with open(os.path.join(text_embed_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 for 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_embed, f2t_embed, co_attn_embed = [], [], []
        t2f_attn_list, f2t_attn_list = [], []
        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:
                image_features = safe_torch_load(frame_file)
                valid_frames = True

                attended_text, attended_image, t2i_raw, i2t_raw = single_sentence_image(
                    text_features, image_features, common_dim=common_dim
                )
                combined_sig = combined_coattn(attended_text, attended_image).unsqueeze(0)  # (1,D)

                # use first token as "CLS-like" (works for both)
                cls_t2f = attended_image[0].unsqueeze(0)  # (1,D)
                cls_f2t = attended_text[0].unsqueeze(0)   # (1,D)

                t2f_embed.append(cls_t2f)
                f2t_embed.append(cls_f2t)
                co_attn_embed.append(combined_sig)
                t2f_attn_list.append(t2i_raw)
                f2t_attn_list.append(i2t_raw)

            except Exception as e:
                print(f"[WARN] {video_id}: error at frame {frame_idx}: {e}")
                continue

        if not valid_frames:
            count += 1
            continue

        # Attention CSVs
        try:
            t2f_attn = norm_t2f(t2f_attn_list)
            dft2f = pd.DataFrame(t2f_attn.detach().cpu().numpy(), index=token_labels)

            f2t_attn = norm_f2t(f2t_attn_list)
            dff2t = pd.DataFrame(f2t_attn.detach().cpu().numpy(), columns=token_labels)

            dff2t.to_csv(os.path.join(text_out, f"frame_to_text{count}.csv"), index=False)
            dft2f.to_csv(os.path.join(frame_out, f"text_to_frame{count}.csv"), index=True)
        except Exception as e:
            print(f"[WARN] {video_id}: failed writing attention CSVs for sentence {count}: {e}")

        # Pool per sentence
        t2f_final = weighted_sharing(t2f_embed)
        f2t_final = weighted_sharing(f2t_embed)
        coattn_final = weighted_sharing(co_attn_embed)

        torch.save(t2f_final, os.path.join(t2f_out_dir, f"CLS{count}.pt"))
        torch.save(f2t_final, os.path.join(f2t_out_dir, f"CLS{count}.pt"))
        torch.save(coattn_final, os.path.join(coattn_out_dir, f"CLS{count}.pt"))

        print(f"[OK] {video_id}: saved sentence {count}")
        count += 1


def main():
    parser = argparse.ArgumentParser(description="Compute co-attention + save CLS embeddings and attention CSVs.")
    parser.add_argument("--manifest", type=str, required=True)
    parser.add_argument("--embed_dir", type=str, required=True)
    parser.add_argument("--out_embed_dir", type=str, required=True)
    parser.add_argument("--attn_dir", type=str, required=True)

    # Optional: force a specific shared dim (e.g., 768). If omitted, it uses min(text_dim, image_dim)
    parser.add_argument("--common_dim", type=int, default=None,
                        help="Project text+image embeddings to this dimension before attention (default: min dims).")

    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)
    os.makedirs(args.attn_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"attn_dir       : {args.attn_dir}")
    print(f"common_dim     : {args.common_dim}")

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

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

        video_id = os.path.splitext(os.path.basename(target_csv))[0]
        print(f"\n[{i}/{len(filenames)}] Processing: {video_id}")

        try:
            process_coattention(target_csv, args.embed_dir, args.out_embed_dir, args.attn_dir, common_dim=args.common_dim)
        except Exception as e:
            print(f"[ERROR] {video_id} failed: {e}")


if __name__ == "__main__":
    main()