import os

import torch
import torch.nn.functional as F

from run_on_video.data_utils import ClipFeatureExtractor
from run_on_video.model_utils import build_inference_model
from utils.tensor_utils import pad_sequences_1d
from video_detr.span_utils import span_cxw_to_xx

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CKPT_PATH = os.path.join(
    REPO_ROOT, "run_on_video/CLIP_ckpt/qvhighlights_onlyCLIP/model_best.ckpt"
)
CLIP_WEIGHTS_PATH = os.path.join(REPO_ROOT, "checkpoints", "ViT-B-32.pt")

# Must match training config (run_on_video/CLIP_ckpt/qvhighlights_onlyCLIP/opt.json:
# "clip_length": 2, "max_v_l": 75). Do NOT change these without re-checking that file.
CLIP_LEN = 2.0
MAX_CLIPS = 75
MAX_VIDEO_SECONDS = CLIP_LEN * MAX_CLIPS  # 150.0


class VideoDETRPredictor:
    def __init__(self, device=None):
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device
        self.feature_extractor = ClipFeatureExtractor(
            framerate=1.0 / CLIP_LEN,
            size=224,
            centercrop=True,
            model_name_or_path=CLIP_WEIGHTS_PATH,
            device=device,
        )
        self.model = build_inference_model(CKPT_PATH).to(device)
        self.model.eval()

    @torch.no_grad()
    def localize_moment(self, video_path, query):
        """
        Args:
            video_path: path to a video file.
            query: a single natural-language query string.

        Returns: dict with keys query, video_duration, clip_len,
            pred_relevant_windows (list of [start_sec, end_sec, score], sorted
            by score desc), pred_saliency_scores (list of float, one per
            CLIP_LEN-second clip), truncated (bool).
        """
        video_feats = self.feature_extractor.encode_video(video_path)
        video_feats = F.normalize(video_feats, dim=-1, eps=1e-5)
        n_frames = len(video_feats)

        truncated = False
        if n_frames > MAX_CLIPS:
            video_feats = video_feats[:MAX_CLIPS]
            n_frames = MAX_CLIPS
            truncated = True

        tef_st = torch.arange(0, n_frames, 1.0) / n_frames
        tef_ed = tef_st + 1.0 / n_frames
        tef = torch.stack([tef_st, tef_ed], dim=1).to(self.device)
        video_feats = torch.cat([video_feats, tef], dim=1).unsqueeze(0)  # (1, T, d)
        video_mask = torch.ones(1, n_frames).to(self.device)

        query_feats = self.feature_extractor.encode_text([query])
        query_feats, query_mask = pad_sequences_1d(
            query_feats, dtype=torch.float32, device=self.device, fixed_length=None
        )
        query_feats = F.normalize(query_feats, dim=-1, eps=1e-5)

        outputs = self.model(
            src_vid=video_feats,
            src_vid_mask=video_mask,
            src_txt=query_feats,
            src_txt_mask=query_mask,
            vid=None,
            qid=None,
        )

        prob = F.softmax(outputs["pred_logits"], -1)
        scores = prob[0, :, 0]  # foreground label is 0
        pred_spans = outputs["pred_spans"][0]  # (#queries, 2)
        saliency_scores = outputs["saliency_scores"][0].float().cpu().tolist()

        video_duration = n_frames * CLIP_LEN
        spans = span_cxw_to_xx(pred_spans.cpu()) * video_duration
        windows = torch.cat([spans, scores.cpu()[:, None]], dim=1).tolist()
        windows.sort(key=lambda x: x[2], reverse=True)
        windows = [[round(st, 2), round(ed, 2), round(sc, 4)] for st, ed, sc in windows]

        return {
            "query": query,
            "video_duration": video_duration,
            "clip_len": CLIP_LEN,
            "pred_relevant_windows": windows,
            "pred_saliency_scores": saliency_scores,
            "truncated": truncated,
        }