"""
谱计算核心模块 - 基于 SVCCA 对偶性理论

核心功能:
- 计算 Gram 矩阵的谱（利用对偶性处理不同维度 d）
- 计算有效秩 (Effective Rank)
- 计算两个谱之间的 Wasserstein-1 距离

数学基础:
    H^T H 与 HH^T 有相同的非零特征值（对偶性定理）

    有效秩定义:
        r_eff = exp(S), 其中 S = -Σ σ'_i log σ'_i
        σ'_i = σ_i / Σ σ_j  (归一化奇异值)
"""

import torch
import numpy as np
from typing import Tuple, Optional, Union


def normalize_representations(H: torch.Tensor, method: str = "per_sample_l2") -> torch.Tensor:
    """
    对表示矩阵进行归一化，消除 scale 影响

    Args:
        H: 表示矩阵，shape = (d, N)
        method: 归一化方法
            - "per_sample_l2": 对每个样本（列）做 L2 归一化（推荐）
            - "global_frobenius": 全局 Frobenius 范数归一化
            - "none": 不归一化

    Returns:
        归一化后的表示矩阵
    """
    if method == "none":
        return H

    if method == "global_frobenius":
        return H / (H.norm(p="fro") + 1e-10)

    if method == "per_sample_l2":
        # 对每列（每个样本）做 L2 归一化
        norm = H.norm(dim=0, keepdim=True) + 1e-10
        return H / norm

    raise ValueError(f"Unknown normalize method: {method}")


def compute_gram_spectrum(
    H: torch.Tensor,
    normalize: str = "per_sample_l2",
    return_eigvals: bool = False
) -> Union[float, Tuple[float, np.ndarray]]:
    """
    利用对偶性，通过 Gram 矩阵计算谱和有效秩

    对于 H ∈ R^{d×N}:
        - 若 d <= N: 计算 d×d 协方差矩阵 C = HH^T / N
        - 若 d > N:  计算 N×N Gram 矩阵 G = H^T H / N（对偶性保证谱相同）

    Args:
        H: 表示矩阵，shape = (d, N)，d=特征维度，N=样本数
        normalize: 归一化方法，见 normalize_representations
        return_eigvals: 是否返回特征值数组

    Returns:
        r_eff: 有效秩（标量）
        eigvals: (可选) 归一化后的特征值分布
    """
    # Step 1: 归一化
    H_norm = normalize_representations(H, method=normalize)
    d, N = H_norm.shape

    # Step 2: 选择较小的矩阵计算（对偶性应用）
    if d <= N:
        # 使用 d×d 协方差矩阵
        C = (H_norm @ H_norm.T) / N
    else:
        # 使用 N×N Gram 矩阵
        C = (H_norm.T @ H_norm) / N

    # Step 3: 特征值分解（对称矩阵用 eigvalsh）
    eigvals_raw = torch.linalg.eigvalsh(C)

    # Step 4: 过滤数值噪声，只保留正特征值
    # 按降序排列
    eigvals_raw = eigvals_raw.flip(0)
    mask = eigvals_raw > 1e-9
    eigvals = eigvals_raw[mask]

    if len(eigvals) == 0:
        # 所有特征值都是数值噪声
        return (1.0, np.array([])) if return_eigvals else 1.0

    # Step 5: 归一化为概率分布
    prob = eigvals / eigvals.sum()

    # Step 6: 计算谱熵
    entropy = -(prob * torch.log(prob + 1e-10)).sum()

    # Step 7: 有效秩 = exp(熵)
    r_eff = torch.exp(entropy).item()

    if return_eigvals:
        return r_eff, prob.cpu().numpy()
    return r_eff


def compute_spectrum_array(
    H: torch.Tensor,
    normalize: str = "per_sample_l2"
) -> Tuple[np.ndarray, np.ndarray]:
    """
    计算完整的谱信息（不仅有效秩）

    Returns:
        eigvals: 原始特征值（降序）
        prob: 归一化概率分布
    """
    H_norm = normalize_representations(H, method=normalize)
    d, N = H_norm.shape

    # 对偶性选择
    if d <= N:
        C = (H_norm @ H_norm.T) / N
    else:
        C = (H_norm.T @ H_norm) / N

    eigvals_raw = torch.linalg.eigvalsh(C).flip(0)
    mask = eigvals_raw > 1e-9
    eigvals = eigvals_raw[mask].cpu().numpy()
    prob = eigvals / eigvals.sum()

    return eigvals, prob


def wasserstein1_distance(p: np.ndarray, q: np.ndarray) -> float:
    """
    计算两个谱概率分布之间的 Wasserstein-1 距离（Earth Mover's Distance）

    对于 1D 分布，W1 距离 = CDF 差的 L1 范数

    Args:
        p, q: 概率分布数组（会自动归一化）

    Returns:
        W1 距离
    """
    # 截断/补零到相同长度
    n = max(len(p), len(q))
    p_ = np.pad(p, (0, n - len(p)), mode='constant').astype(np.float64)
    q_ = np.pad(q, (0, n - len(q)), mode='constant').astype(np.float64)

    # 重新归一化（确保是概率分布）
    p_ = p_ / (p_.sum() + 1e-10)
    q_ = q_ / (q_.sum() + 1e-10)

    # 1D Wasserstein = CDF 差的积分
    cdf_p = np.cumsum(p_)
    cdf_q = np.cumsum(q_)

    return float(np.sum(np.abs(cdf_p - cdf_q)))


def compute_spectrum_distance_matrix(
    prob_dict: dict,
    distance_fn: str = "wasserstein1"
) -> Tuple[np.ndarray, list]:
    """
    计算多个模型谱之间的距离矩阵

    Args:
        prob_dict: {model_name: prob_distribution}
        distance_fn: 距离函数名

    Returns:
        distance_matrix: n×n 距离矩阵
        model_names: 模型名称列表
    """
    model_names = list(prob_dict.keys())
    n = len(model_names)
    dist_matrix = np.zeros((n, n))

    dist_fn = wasserstein1_distance if distance_fn == "wasserstein1" else None

    for i in range(n):
        for j in range(i + 1, n):
            d = dist_fn(prob_dict[model_names[i]], prob_dict[model_names[j]])
            dist_matrix[i, j] = d
            dist_matrix[j, i] = d

    return dist_matrix, model_names


def effective_rank_from_eigvals(eigvals: np.ndarray) -> float:
    """
    从特征值直接计算有效秩

    Args:
        eigvals: 特征值数组（无需预先归一化）

    Returns:
        有效秩
    """
    eigvals = np.array(eigvals)
    eigvals = eigvals[eigvals > 1e-10]  # 过滤零值

    if len(eigvals) == 0:
        return 1.0

    # 归一化为概率分布
    prob = eigvals / eigvals.sum()

    # 谱熵
    entropy = -np.sum(prob * np.log(prob + 1e-10))

    return float(np.exp(entropy))


# ============= 可视化辅助函数 =============

def compute_top_k_ratio(prob: np.ndarray, k: int = 10) -> float:
    """
    计算前 k 个特征值的能量占比

    Args:
        prob: 归一化特征值分布（降序）
        k: 前 k 个

    Returns:
        能量占比
    """
    return float(prob[:k].sum())


def spectrum_summary(r_eff: float, prob: np.ndarray) -> dict:
    """
    生成谱的摘要统计

    Returns:
        包含各种统计指标的字典
    """
    return {
        "r_eff": r_eff,
        "num_eigvals": len(prob),
        "top1_ratio": float(prob[0]) if len(prob) > 0 else 0,
        "top5_ratio": float(prob[:5].sum()) if len(prob) >= 5 else float(prob.sum()),
        "top10_ratio": float(prob[:10].sum()) if len(prob) >= 10 else float(prob.sum()),
        "entropy": float(np.log(r_eff)),  # S = log(r)
    }