"""无监督聚类：特征提取 + KMeans。

用 torch 自己实现 KMeans（kmeans++ 初始化 + Lloyd 迭代），不引 scikit-learn：
一是少一个依赖，二是能直接吃 GPU 上的特征张量。
"""

from __future__ import annotations

import math
from collections.abc import Sequence
from dataclasses import dataclass, field

import torch
from torch import nn
from torch.utils.data import DataLoader

from utils.logging import get_logger

logger = get_logger("engine.cluster")


@dataclass
class ClusterResult:
    labels: torch.Tensor  # (N,) 每个样本的簇下标
    centroids: torch.Tensor  # (K, D)
    distances: torch.Tensor  # (N,) 到所属质心的距离
    inertia: float  # 簇内平方距离和，越小越紧
    silhouette: float = float("nan")  # 轮廓系数，[-1, 1]，越大越好
    sizes: list[int] = field(default_factory=list)
    points: torch.Tensor | None = None  # KMeans 实际使用的空间（归一化 + PCA 之后），供可视化复用


def l2_normalize(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    return x / x.norm(dim=1, keepdim=True).clamp(min=eps)


def pca_reduce(x: torch.Tensor, dim: int, whiten: bool = False, eps: float = 1e-5) -> torch.Tensor:
    """PCA 降维，`whiten` 决定是否把各主成分方差拉平。

    **白化默认关闭**：CNN 特征里只有前若干个主成分携带结构，其余接近噪声；
    白化会把这些噪声方向抬到和信号同权，簇随即塌成"一个大簇 + 几个碎簇"。
    实测 3 个信号维 + 200 个噪声维的合成数据上，dim=3 轮廓 0.93，
    dim=64 + 白化掉到 0.10 且最大簇占 98%。需要白化时请同时把 dim 调小。
    """
    dim = min(dim, x.shape[0], x.shape[1])
    centered = x - x.mean(dim=0, keepdim=True)
    _, singular, components = torch.pca_lowrank(centered, q=dim, center=False)
    projected = centered @ components[:, :dim]
    if not whiten:
        return projected
    return projected / (singular[:dim] / (x.shape[0] ** 0.5) + eps)


def kmeans_plusplus(x: torch.Tensor, k: int, generator: torch.Generator) -> torch.Tensor:
    """kmeans++ 初始化：按到已选质心距离的平方为概率挑下一个，避免初值退化。"""
    n = x.shape[0]
    centroids = [x[torch.randint(n, (1,), generator=generator, device=x.device).item()]]
    closest = torch.cdist(x, centroids[0][None]).squeeze(1) ** 2
    for _ in range(1, k):
        probs = closest / closest.sum().clamp(min=1e-12)
        index = torch.multinomial(probs, 1, generator=generator).item()
        centroids.append(x[index])
        closest = torch.minimum(closest, torch.cdist(x, centroids[-1][None]).squeeze(1) ** 2)
    return torch.stack(centroids)


def kmeans(
    x: torch.Tensor,
    num_clusters: int,
    iters: int = 100,
    n_init: int = 10,
    tol: float = 1e-5,
    seed: int = 42,
) -> tuple[torch.Tensor, torch.Tensor, float]:
    """跑 `n_init` 次取 inertia 最小的一次，返回 (labels, centroids, inertia)。"""
    if num_clusters < 2:
        raise ValueError(f"num_clusters 至少为 2，收到 {num_clusters}")
    if x.shape[0] < num_clusters:
        raise ValueError(f"样本数 {x.shape[0]} 少于簇数 {num_clusters}")

    best: tuple[torch.Tensor, torch.Tensor, float] | None = None
    for run in range(n_init):
        generator = torch.Generator(device=x.device).manual_seed(seed + run)
        centroids = kmeans_plusplus(x, num_clusters, generator)
        previous = float("inf")

        for _ in range(iters):
            distances = torch.cdist(x, centroids)
            labels = distances.argmin(dim=1)
            inertia = float((distances.gather(1, labels[:, None]) ** 2).sum())

            for index in range(num_clusters):
                members = x[labels == index]
                if len(members):
                    centroids[index] = members.mean(dim=0)
                else:
                    # 空簇：把离自己质心最远的点抢过来，避免簇数缩水
                    centroids[index] = x[distances.gather(1, labels[:, None]).squeeze(1).argmax()]
            if abs(previous - inertia) <= tol * max(previous, 1.0):
                break
            previous = inertia

        if best is None or inertia < best[2]:
            best = (labels, centroids.clone(), inertia)

    assert best is not None
    return best


@torch.no_grad()
def silhouette_score(
    x: torch.Tensor, labels: torch.Tensor, max_samples: int = 2000, seed: int = 42
) -> float:
    """轮廓系数。是 O(N²) 的，所以超过 `max_samples` 就抽样估计。"""
    n = x.shape[0]
    if labels.unique().numel() < 2:
        return float("nan")
    if n > max_samples:
        generator = torch.Generator(device="cpu").manual_seed(seed)
        index = torch.randperm(n, generator=generator)[:max_samples].to(x.device)
        x, labels = x[index], labels[index]

    distances = torch.cdist(x, x)
    scores = []
    for index in range(x.shape[0]):
        same = labels == labels[index]
        same[index] = False
        if not same.any():
            continue  # 单点簇按惯例记 0
        a = distances[index][same].mean()
        b = min(
            distances[index][labels == other].mean()
            for other in labels.unique().tolist()
            if other != labels[index].item()
        )
        scores.append(((b - a) / torch.maximum(a, b)).item())
    return float(sum(scores) / len(scores)) if scores else float("nan")


class StageFeatureExtractor(nn.Module):
    """从卷积骨架的多个 stage 取特征，各自全局池化 + L2 归一化后拼接。

    只用最后一层的语义特征对薄膜缺陷不够用——缺陷之间的差异很大程度是**纹理**层面的
    （颗粒感、条纹方向、边缘锐度），这些信息在中间层比在最后一层保留得好。
    """

    def __init__(self, backbone: nn.Module, stages: Sequence[int] | None = None) -> None:
        super().__init__()
        self.backbone = backbone
        layers = getattr(backbone, "features", backbone)
        self.layers = list(layers) if isinstance(layers, nn.Sequential) else None
        self.stages = sorted(stages) if stages else None
        if self.stages and self.layers is None:
            raise ValueError("该骨架不是 nn.Sequential，无法按 stage 取中间特征")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.stages:
            feat = self.backbone(x)
            return feat.mean(dim=(2, 3)) if feat.dim() == 4 else feat

        outputs = []
        for index, layer in enumerate(self.layers):
            x = layer(x)
            if index in self.stages:
                outputs.append(l2_normalize(x.mean(dim=(2, 3))))
        if not outputs:
            raise ValueError(f"stages {self.stages} 超出骨架层数 {len(self.layers)}")
        return torch.cat(outputs, dim=1)


def normalized_mutual_info(labels_a: Sequence, labels_b: Sequence) -> float:
    """两组划分的归一化互信息，[0, 1]。这里用来比较**两次聚类结果**的一致性。"""
    from collections import Counter

    n = len(labels_a)
    if n == 0:
        return float("nan")
    count_a, count_b = Counter(labels_a), Counter(labels_b)
    joint = Counter(zip(labels_a, labels_b, strict=True))

    def entropy(counter: Counter) -> float:
        return -sum((c / n) * math.log(c / n) for c in counter.values() if c)

    mutual = sum(
        (c / n) * math.log((c / n) / ((count_a[a] / n) * (count_b[b] / n)))
        for (a, b), c in joint.items()
        if c
    )
    denominator = (entropy(count_a) * entropy(count_b)) ** 0.5
    return float(mutual / denominator) if denominator > 0 else float("nan")


@torch.no_grad()
def extract_features(
    encoder: nn.Module,
    dataloader: DataLoader,
    device: torch.device,
    use_shape_features: bool = False,
    log_interval: int = 20,
) -> tuple[torch.Tensor, list[dict]]:
    """跑一遍数据集拿 embedding。

    `encoder` 输出特征图，这里做全局平均池化。

    `use_shape_features` 默认关闭：图片的宽高比/尺寸是 ROI 裁剪环节的产物，
    不是缺陷属性，拼进 embedding 只会让聚类去区分"采集方式"而不是缺陷类型。
    留着这个开关只为做对照实验。
    """
    encoder.eval()
    features: list[torch.Tensor] = []
    shapes: list[torch.Tensor] = []
    metas: list[dict] = []

    for step, batch in enumerate(dataloader):
        images = batch["views"][0].to(device, non_blocking=True)
        feat = encoder(images)
        if feat.dim() == 4:
            feat = feat.mean(dim=(2, 3))
        features.append(feat.float().cpu())
        shapes.append(batch["shape"])
        metas.extend(batch["meta"])
        if step % log_interval == 0:
            logger.info("提取特征 [%d/%d]", step, len(dataloader))

    embeddings = torch.cat(features)
    if use_shape_features:
        shape = torch.cat(shapes)
        shape = (shape - shape.mean(dim=0)) / shape.std(dim=0).clamp(min=1e-6)
        # 图像特征先归一化，避免两部分量纲差太多导致形状特征被淹没
        embeddings = torch.cat([l2_normalize(embeddings), shape], dim=1)
    return embeddings, metas


def cluster_features(
    embeddings: torch.Tensor,
    num_clusters: int,
    pca_dim: int = 32,
    whiten: bool = False,
    normalize: bool = True,
    n_init: int = 10,
    iters: int = 100,
    seed: int = 42,
    fit_index: torch.Tensor | None = None,
) -> ClusterResult:
    """L2 归一化 -> PCA 降维 -> KMeans -> 指标。

    `fit_index` 给定时只用这部分样本拟合质心（数据极不均衡时用来做类别平衡采样），
    再把全部样本按最近质心分配。
    """
    x = l2_normalize(embeddings) if normalize else embeddings
    if pca_dim:
        x = pca_reduce(x, pca_dim, whiten=whiten)

    fit_x = x if fit_index is None else x[fit_index]
    _, centroids, _ = kmeans(fit_x, num_clusters, iters=iters, n_init=n_init, seed=seed)

    distances = torch.cdist(x, centroids)
    labels = distances.argmin(dim=1)
    nearest = distances.gather(1, labels[:, None]).squeeze(1)

    return ClusterResult(
        points=x,
        labels=labels,
        centroids=centroids,
        distances=nearest,
        inertia=float((nearest**2).sum()),
        silhouette=silhouette_score(x, labels, seed=seed),
        sizes=torch.bincount(labels, minlength=num_clusters).tolist(),
    )


def centroid_distances(centroids: torch.Tensor) -> torch.Tensor:
    """质心两两欧氏距离 (k, k)。数值上直接回答"簇之间离多远"。"""
    return torch.cdist(centroids, centroids)


def project_2d(
    points: torch.Tensor, centroids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, list[float]]:
    """把聚类空间投到 2 维用于画图，质心用同一组主成分变换。

    用 PCA 而不是 t-SNE/UMAP：后两者为了把簇拉开会**扭曲簇间距离**，
    而这里的目的恰恰是看簇之间离多远。代价是高维结构被压掉，
    所以返回两个主成分的方差占比，图上会标出来供判断可信度。
    """
    mean = points.mean(dim=0, keepdim=True)
    centered = points - mean
    _, singular, components = torch.pca_lowrank(centered, q=min(2, *points.shape), center=False)
    basis = components[:, :2]
    variance = singular[:2] ** 2 / (points.shape[0] - 1)
    total = (centered**2).sum() / (points.shape[0] - 1)
    ratio = (variance / total.clamp(min=1e-12)).tolist()
    return centered @ basis, (centroids - mean) @ basis, ratio


def random_fit_index(num_samples: int, max_samples: int, seed: int = 42) -> torch.Tensor | None:
    """随机抽一部分样本用于拟合质心，纯粹为了省时间（拟合完仍会分配全部样本）。

    不做任何按来源目录的分层——目录名是人工整理的历史产物，不应该影响聚类。
    """
    if not max_samples or num_samples <= max_samples:
        return None
    generator = torch.Generator().manual_seed(seed)
    return torch.randperm(num_samples, generator=generator)[:max_samples].sort().values


def cluster_stability(
    embeddings: torch.Tensor, num_clusters: int, seeds: Sequence[int] = (0, 1), **kwargs
) -> float:
    """换随机种子重跑，看两次划分的 NMI。

    无标签场景下判断簇数是否合理的实用指标：真有结构时不同初始化会收敛到同一套划分，
    强行切分一团连续数据时结果会随种子乱跳。
    """
    runs = [
        cluster_features(embeddings, num_clusters, seed=seed, **kwargs).labels.tolist()
        for seed in seeds
    ]
    return normalized_mutual_info(runs[0], runs[1])