import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from scipy.optimize import linear_sum_assignment
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False


class TokenL2Loss(nn.Module):
    """
    输入:
        history:      [B, T, N, D]，第 0 帧到第 T-1 帧
        refined_pred: [B, T, N, D]，模型预测的第 1 帧到第 T 帧
        target:       [B, T, N, D]，真实的第 1 帧到第 T 帧

    功能:
        1. 先计算 refined_pred 和 target 的原始逐 index MSE。
        2. 对 history[:, t] 和 target[:, t] 的 N 个 tokens 做一一匹配。
           匹配依据是 cosine similarity，目标是最大化整体相似度。
        3. 根据匹配关系，将 target[:, t] 重排到 history[:, t] 的 token 顺序。
        4. 计算 refined_pred 和 matched_target 之间的 MSE。
    """

    def __init__(self, matching_method="hungarian"):
        """
        参数:
            matching_method:
                "hungarian": 全局最优一一匹配，推荐，但依赖 scipy，速度较慢。
                "greedy": 按 cosine 从大到小贪心匹配，不依赖 scipy。
        """
        super().__init__()
        if matching_method not in ["hungarian", "greedy"]:
            raise ValueError("matching_method must be 'hungarian' or 'greedy'.")

        if matching_method == "hungarian" and not SCIPY_AVAILABLE:
            print("[Warning] scipy not available, fallback to greedy matching.")
            matching_method = "greedy"

        self.matching_method = matching_method

    @torch.no_grad()
    def _hungarian_match_one(self, src_tokens, tgt_tokens):
        """
        src_tokens: [N, D]
        tgt_tokens: [N, D]

        返回:
            matched_indices: [N]
            matched_indices[i] = j 表示 src token i 匹配 target token j
        """
        src = F.normalize(src_tokens.float(), dim=-1)
        tgt = F.normalize(tgt_tokens.float(), dim=-1)

        sim = torch.matmul(src, tgt.T)  # [N, N]

        # scipy 的 linear_sum_assignment 是最小化 cost，
        # 所以用 -sim 来实现最大化 cosine similarity。
        cost = (-sim).detach().cpu().numpy()
        row_ind, col_ind = linear_sum_assignment(cost)

        matched_indices = torch.empty(
            src_tokens.shape[0],
            dtype=torch.long,
            device=src_tokens.device,
        )
        matched_indices[torch.from_numpy(row_ind).to(src_tokens.device)] = torch.from_numpy(col_ind).to(src_tokens.device)
        return matched_indices

    @torch.no_grad()
    def _greedy_match_one(self, src_tokens, tgt_tokens):
        """
        贪心一一匹配:
        按 cosine similarity 从大到小排序，优先匹配最高相似度 pair。
        """
        src = F.normalize(src_tokens.float(), dim=-1)
        tgt = F.normalize(tgt_tokens.float(), dim=-1)

        sim = torch.matmul(src, tgt.T)  # [N, N]
        n = sim.shape[0]

        flat_sim = sim.reshape(-1)
        sorted_indices = torch.argsort(flat_sim, descending=True)

        matched_indices = torch.full(
            (n,),
            -1,
            dtype=torch.long,
            device=src_tokens.device,
        )
        used_src = torch.zeros(n, dtype=torch.bool, device=src_tokens.device)
        used_tgt = torch.zeros(n, dtype=torch.bool, device=src_tokens.device)

        num_matched = 0
        for flat_idx in sorted_indices:
            src_idx = flat_idx // n
            tgt_idx = flat_idx % n

            if used_src[src_idx] or used_tgt[tgt_idx]:
                continue

            matched_indices[src_idx] = tgt_idx
            used_src[src_idx] = True
            used_tgt[tgt_idx] = True
            num_matched += 1

            if num_matched == n:
                break

        return matched_indices

    @torch.no_grad()
    def build_matched_target(self, history, target):
        """
        根据 history 和 target 的相邻帧 token cosine similarity，
        为每个 target frame 构造按 history token 顺序重排后的 matched_target。

        history: [B, T, N, D]
        target:  [B, T, N, D]

        返回:
            matched_target: [B, T, N, D]
            match_indices:  [B, T, N]
        """
        if history.shape != target.shape:
            raise ValueError(
                f"history and target must have same shape, got "
                f"{tuple(history.shape)} and {tuple(target.shape)}."
            )

        bsz, seq_len, num_tokens, dim = history.shape

        matched_target = torch.empty_like(target)
        match_indices = torch.empty(
            (bsz, seq_len, num_tokens),
            dtype=torch.long,
            device=target.device,
        )

        for b in range(bsz):
            for t in range(seq_len):
                src_tokens = history[b, t]  # [N, D]
                tgt_tokens = target[b, t]   # [N, D]

                if self.matching_method == "hungarian":
                    indices = self._hungarian_match_one(src_tokens, tgt_tokens)
                else:
                    indices = self._greedy_match_one(src_tokens, tgt_tokens)

                match_indices[b, t] = indices
                matched_target[b, t] = target[b, t, indices]

        return matched_target, match_indices

    def forward(self, history, refined_pred, target):
        """
        返回:
            {
                "total": matched_mse,
                "original_mse": 原始逐 index MSE,
                "matched_mse": 匹配后一一对应 MSE,
                "match_indices": 匹配关系 [B, T, N]
            }
        """
        history = history.float()
        refined_pred = refined_pred.float()
        target = target.float()

        if refined_pred.shape != target.shape:
            raise ValueError(
                f"refined_pred and target must have same shape, got "
                f"{tuple(refined_pred.shape)} and {tuple(target.shape)}."
            )

        if history.shape != target.shape:
            raise ValueError(
                f"history and target must have same shape, got "
                f"{tuple(history.shape)} and {tuple(target.shape)}."
            )

        original_mse = F.mse_loss(refined_pred, target)

        matched_target, match_indices = self.build_matched_target(
            history=history,
            target=target,
        )

        matched_mse = F.mse_loss(refined_pred, matched_target)

        return {
            "total": matched_mse,
            "original_mse": original_mse,
            "matched_mse": matched_mse,
            "match_indices": match_indices,
        }