#!/usr/bin/env python
"""把未分类的样本做无监督聚类，产出可人工核对的分组结果。

    # 默认 4 类，ImageNet 预训练特征（不训练，CPU 可跑）
    uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml

    # 改簇数 / 换成在本地数据上自监督训练的特征（Mac 上自动走 MPS）
    # 断点续训：--set method=simclr simclr.resume=runs/simclr/simclr_last.pth
    uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
        --set num_clusters=6 method=simclr

产出（work_dir 下）：
    assignments.csv           每张图的簇号、到质心距离、原目录
    summary.json              簇大小、轮廓系数、簇 × 原目录交叉表
    clusters/cluster_00/...   按簇建的软链目录，直接用看图软件翻
    cluster_00_preview.png    每簇代表样本（离质心最近）的拼图
"""

from __future__ import annotations

import argparse
import csv
import json
import shutil
import sys
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch  # noqa: E402

from datasets.build import build_dataset  # noqa: E402
from datasets.unlabeled import unlabeled_collate  # noqa: E402
from engine.cluster import (  # noqa: E402
    StageFeatureExtractor,
    centroid_distances,
    cluster_features,
    cluster_stability,
    extract_features,
    project_2d,
    random_fit_index,
)
from models.registry import BACKBONES  # noqa: E402
from utils import load_config, seed_everything, setup_logging  # noqa: E402
from utils.device import (  # noqa: E402
    DeviceNotAvailableError,
    describe_device,
    resolve_device,
)

logger = None


def build_encoder(cfg, device, work_dir):
    """按 `method` 拿一个特征提取器。

    pretrained: ImageNet 预训练卷积骨架，纯推理，CPU 也能跑。用顶层的 `backbone`。
    simclr:     在本目录数据上自监督训练一个项目自己的卷积骨架。
                用 `simclr.backbone`，与 pretrained 的骨架配置各自独立：两者的字段完全不同
                （arch/pretrained vs widths/depths），共用一个块会在 `_base_` 合并时串味。
    """
    method = str(cfg.get("method", "pretrained"))

    if method == "pretrained":
        backbone = BACKBONES.build(dict(cfg.get("backbone", {"name": "torchvision"})))
        return StageFeatureExtractor(backbone, cfg.get("feature_stages")).to(device).eval()

    if method != "simclr":
        raise ValueError(f"未知 method: {method}，可选 pretrained / simclr")

    from torch.utils.data import DataLoader

    from engine.simclr import SimCLRTrainer

    simclr_cfg = dict(cfg.get("simclr", {}))
    backbone_cfg = dict(simclr_cfg.pop("backbone", {"name": "tiny_conv"}))
    # 自研骨架不是 nn.Sequential，默认取最终输出而不是中间 stage
    stages = simclr_cfg.pop("feature_stages", None)
    backbone = BACKBONES.build(backbone_cfg).to(device)

    resume = simclr_cfg.pop("resume", None)
    checkpoint = simclr_cfg.pop("checkpoint", None)
    if resume and checkpoint:
        raise ValueError(
            "simclr.resume 与 simclr.checkpoint 只能二选一："
            "resume=接着训练，checkpoint=跳过训练直接聚类。"
        )
    if checkpoint and Path(checkpoint).exists():
        state = torch.load(checkpoint, map_location=device, weights_only=False)
        backbone.load_state_dict(state["model"])
        logger.info("复用已有的自监督骨架: %s", checkpoint)
        return StageFeatureExtractor(backbone, stages).to(device).eval()

    view_dataset = build_dataset(
        {**dict(cfg["dataset"]), "views": 2, "view_transforms": cfg.get("view_transforms", [])}
    )
    loader = DataLoader(
        view_dataset,
        batch_size=int(simclr_cfg.pop("batch_size", 128)),
        shuffle=True,
        drop_last=True,
        num_workers=int(cfg.get("loader", {}).get("num_workers", 4)),
        collate_fn=unlabeled_collate,
    )
    full_config = cfg.to_dict() if hasattr(cfg, "to_dict") else dict(cfg)
    SimCLRTrainer(
        backbone, loader, simclr_cfg, device, work_dir, full_config=full_config, resume=resume
    ).fit()
    return StageFeatureExtractor(backbone, stages).to(device).eval()


def write_outputs(cfg, work_dir, metas, result, num_clusters, stability):
    """CSV + JSON + 软链目录 + 预览拼图。人工只需要翻软链目录就能确认分组是否合理。"""
    labels = result.labels.tolist()
    distances = result.distances.tolist()

    with (work_dir / "assignments.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["path", "cluster", "distance_to_centroid", "source_dir", "width", "height"])
        for meta, label, distance in zip(metas, labels, distances, strict=True):
            width, height = meta["size"]
            writer.writerow([meta["path"], label, f"{distance:.6f}", meta["source_dir"], width, height])

    method = str(cfg.get("method", "pretrained"))
    backbone_cfg = (
        dict(cfg.get("simclr", {}).get("backbone", {}))
        if method == "simclr"
        else dict(cfg.get("backbone", {}))
    )
    summary = {
        "num_clusters": num_clusters,
        "num_images": len(metas),
        "method": method,
        # 记录特征来源，供 export_clusters.py 自动命名以及日后追溯
        "feature_source": {
            "method": method,
            "backbone": backbone_cfg,
            "feature_stages": (
                cfg.get("simclr", {}).get("feature_stages")
                if method == "simclr"
                else cfg.get("feature_stages")
            ),
        },
        "image_size": cfg.get("image_size"),
        "cluster_sizes": result.sizes,
        "silhouette": result.silhouette,
        "stability": stability,
        "inertia": result.inertia,
        "mean_distance_per_cluster": [
            float(result.distances[result.labels == index].mean()) if size else None
            for index, size in enumerate(result.sizes)
        ],
        "centroid_distances": centroid_distances(result.centroids).tolist(),
    }
    (work_dir / "summary.json").write_text(
        json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8"
    )

    link_limit = int(cfg.get("link_limit", 300))
    preview_count = int(cfg.get("preview_count", 25))
    clusters_dir = work_dir / "clusters"
    if clusters_dir.exists():
        shutil.rmtree(clusters_dir)

    order = sorted(range(len(metas)), key=lambda i: distances[i])  # 离质心近的排前面
    per_cluster: dict[int, list[int]] = defaultdict(list)
    for index in order:
        per_cluster[labels[index]].append(index)

    for cluster, indices in sorted(per_cluster.items()):
        target = clusters_dir / f"cluster_{cluster:02d}"
        target.mkdir(parents=True, exist_ok=True)
        # 近的取大头，末尾补几张最远的——离群样本最能暴露这个簇是不是混了东西
        chosen = indices[: max(link_limit - 20, 1)] + indices[-20:] if len(indices) > link_limit else indices
        for rank, index in enumerate(chosen):
            source = Path(metas[index]["path"]).resolve()
            link = target / f"{rank:04d}_{source.stem}{source.suffix}"
            if not link.exists():
                link.symlink_to(source)
        save_preview(work_dir / f"cluster_{cluster:02d}_preview.png", metas, indices[:preview_count])

    if cfg.get("scatter", {}).get("enabled", True) and result.points is not None:
        scatter_path = work_dir / "cluster_scatter.png"
        save_scatter(scatter_path, result, cfg)
        logger.info("簇分布图已保存: %s", scatter_path)

    return summary


#: 按优先级尝试的中文字体（macOS / Linux 常见）。都找不到就退回英文标签，避免画出豆腐块。
CJK_FONTS = (
    "PingFang SC",
    "Hiragino Sans GB",
    "Heiti SC",
    "Noto Sans CJK SC",
    "Source Han Sans SC",
    "WenQuanYi Zen Hei",
    "Microsoft YaHei",
    "SimHei",
    "Arial Unicode MS",
)


def setup_cjk_font() -> bool:
    """给 matplotlib 配一个能显示中文的字体，返回是否配上了。"""
    import matplotlib
    from matplotlib import font_manager

    available = {f.name for f in font_manager.fontManager.ttflist}
    for name in CJK_FONTS:
        if name in available:
            matplotlib.rcParams["font.sans-serif"] = [name, *matplotlib.rcParams["font.sans-serif"]]
            matplotlib.rcParams["axes.unicode_minus"] = False  # 中文字体里没有 U+2212
            return True
    logger.warning("未找到中文字体（尝试过 %s），图上改用英文标签", "/".join(CJK_FONTS[:3]))
    return False


def save_scatter(path: Path, result, cfg) -> None:
    """把聚类结果画成二维散点 + 质心距离矩阵，放在 preview 拼图的同级目录。

    投影用 PCA 而不是 t-SNE/UMAP：后两者为了把簇拉开会扭曲簇间距离，
    而这张图的目的就是看簇之间离多远。图上标注了两个主成分的方差占比——
    占比低说明高维结构被压掉得多，图上的远近只能参考，以右边的距离矩阵为准。
    """
    import matplotlib

    matplotlib.use("Agg")  # 无显示环境
    import matplotlib.pyplot as plt

    cjk = setup_cjk_font()

    scatter_cfg = dict(cfg.get("scatter", {}))
    max_points = int(scatter_cfg.get("max_points", 5000))
    points, centroids, ratio = project_2d(result.points, result.centroids)
    labels = result.labels

    # 点太多时随机抽样：既避免糊成一团，也控制文件大小
    if max_points and points.shape[0] > max_points:
        generator = torch.Generator().manual_seed(int(cfg.get("seed", 42)))
        index = torch.randperm(points.shape[0], generator=generator)[:max_points]
        points, labels = points[index], labels[index]
        subtitle = f"随机抽样 {max_points} / {result.labels.shape[0]} 点"
        subtitle_en = f"sampled {max_points} / {result.labels.shape[0]}"
    else:
        subtitle = f"全部 {points.shape[0]} 点"
        subtitle_en = f"all {points.shape[0]} points"

    points = points.cpu().numpy()
    centroids_2d = centroids.cpu().numpy()
    labels = labels.cpu().numpy()
    num_clusters = result.centroids.shape[0]
    colors = plt.get_cmap("tab10")(range(num_clusters))

    figsize = scatter_cfg.get("figsize", [14, 6])
    fig, (ax, ax_matrix) = plt.subplots(
        1, 2, figsize=tuple(figsize), gridspec_kw={"width_ratios": [1.6, 1]}
    )

    for cluster in range(num_clusters):
        mask = labels == cluster
        ax.scatter(
            points[mask, 0],
            points[mask, 1],
            s=float(scatter_cfg.get("point_size", 6)),
            alpha=float(scatter_cfg.get("alpha", 0.35)),
            color=colors[cluster],
            linewidths=0,
            label=f"cluster_{cluster:02d}  n={int(result.sizes[cluster])}",
        )
    ax.scatter(
        centroids_2d[:, 0],
        centroids_2d[:, 1],
        s=220,
        marker="X",
        c=colors,
        edgecolors="black",
        linewidths=1.5,
        zorder=5,
    )
    for cluster in range(num_clusters):
        ax.annotate(
            f"{cluster:02d}",
            centroids_2d[cluster],
            textcoords="offset points",
            xytext=(0, 12),
            fontsize=11,
            fontweight="bold",
            ha="center",
            zorder=6,
            color="black",
            bbox={"boxstyle": "round,pad=0.2", "fc": "white", "ec": "none", "alpha": 0.8},
        )

    if cjk:
        scatter_title = (
            f"PCA 二维投影（{subtitle}）\n"
            f"两个主成分方差占比 {ratio[0]:.1%} + {ratio[1]:.1%} = {sum(ratio):.1%}"
        )
        matrix_title = "质心两两距离（聚类空间，非投影后）"
    else:
        scatter_title = (
            f"PCA 2D projection ({subtitle_en})\n"
            f"explained variance {ratio[0]:.1%} + {ratio[1]:.1%} = {sum(ratio):.1%}"
        )
        matrix_title = "pairwise centroid distance (in cluster space)"
    ax.set_title(scatter_title)
    ax.set_xlabel("PC1")
    ax.set_ylabel("PC2")
    ax.legend(loc="best", fontsize=8, framealpha=0.9)
    ax.grid(alpha=0.2)
    ax.set_aspect("equal", adjustable="datalim")

    distances = centroid_distances(result.centroids).cpu().numpy()
    image = ax_matrix.imshow(distances, cmap="viridis")
    ax_matrix.set_title(matrix_title)
    ax_matrix.set_xticks(range(num_clusters))
    ax_matrix.set_yticks(range(num_clusters))
    ax_matrix.set_xticklabels([f"{i:02d}" for i in range(num_clusters)])
    ax_matrix.set_yticklabels([f"{i:02d}" for i in range(num_clusters)])
    for i in range(num_clusters):
        for j in range(num_clusters):
            ax_matrix.text(
                j, i, f"{distances[i, j]:.2f}",
                ha="center", va="center",
                color="white" if distances[i, j] < distances.max() * 0.6 else "black",
                fontsize=9,
            )
    fig.colorbar(image, ax=ax_matrix, fraction=0.046)

    fig.tight_layout()
    fig.savefig(path, dpi=int(scatter_cfg.get("dpi", 150)))
    plt.close(fig)


def save_preview(path: Path, metas: list[dict], indices: list[int], cell: int = 96, columns: int = 5):
    """把代表样本拼成一张图，方便一眼看出这个簇是什么缺陷。"""
    if not indices:
        return
    from PIL import Image

    rows = (len(indices) + columns - 1) // columns
    canvas = Image.new("RGB", (columns * cell, rows * cell), (32, 32, 32))
    for position, index in enumerate(indices):
        with Image.open(metas[index]["path"]) as raw:
            image = raw.convert("RGB")
        scale = min(cell / image.width, cell / image.height)
        image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale))))
        x = (position % columns) * cell + (cell - image.width) // 2
        y = (position // columns) * cell + (cell - image.height) // 2
        canvas.paste(image, (x, y))
    canvas.save(path)


def sweep_clusters(embeddings, spec: str, cluster_kwargs: dict) -> None:
    """扫一遍簇数，同时看轮廓系数和换种子的稳定性。

    无标签时没有"正确答案"，只能看：簇是否紧致（轮廓），以及换初始化后划分是否还一样
    （稳定性）。稳定性低说明那个 k 是在硬切一团连续的数据。
    """
    start, end = (int(v) for v in spec.split("-"))
    logger.info("%3s %9s %10s %12s  簇大小", "k", "轮廓", "稳定性", "最大簇占比")
    for k in range(start, end + 1):
        result = cluster_features(embeddings, k, seed=42, **cluster_kwargs)
        stability = cluster_stability(embeddings, k, seeds=(0, 7), **cluster_kwargs)
        logger.info(
            "%3d %9.3f %10.3f %12.2f  %s",
            k,
            result.silhouette,
            stability,
            max(result.sizes) / sum(result.sizes),
            result.sizes,
        )


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--config", default="configs/cluster/film.yaml")
    parser.add_argument("--root", default=None, help="覆盖配置里的图片根目录")
    parser.add_argument("--num-clusters", type=int, default=None, help="覆盖配置里的簇数")
    parser.add_argument("--work-dir", default=None)
    parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu")
    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
    parser.add_argument("--save-features", action="store_true", help="额外保存 embedding 供复用")
    parser.add_argument(
        "--sweep-k", default=None, metavar="A-B", help="只扫簇数（如 2-8）打印指标表，不产出分组"
    )
    args = parser.parse_args()

    cfg = load_config(args.config, args.set)
    if args.root:
        cfg["dataset"]["root"] = args.root
    if args.num_clusters:
        cfg["num_clusters"] = args.num_clusters

    work_dir = Path(args.work_dir or cfg.get("work_dir", "runs/cluster"))
    work_dir.mkdir(parents=True, exist_ok=True)
    logger = setup_logging(work_dir / "cluster.log")
    seed_everything(int(cfg.get("seed", 42)))

    num_clusters = int(cfg.get("num_clusters", 4))
    try:
        device = resolve_device(args.device or cfg.get("device", "auto"))
    except DeviceNotAvailableError as exc:
        logger.error("%s", exc)
        return 1

    dataset = build_dataset(dict(cfg["dataset"]))
    logger.info(
        "扫描到 %d 张图，来自 %d 个目录；簇数 %d，method=%s，设备 %s",
        len(dataset),
        len({dataset.source_dir(path) for path in dataset.images}),
        num_clusters,
        cfg.get("method", "pretrained"),
        describe_device(device),
    )

    encoder = build_encoder(cfg, device, work_dir)

    from torch.utils.data import DataLoader

    loader = DataLoader(
        dataset,
        batch_size=int(cfg.get("loader", {}).get("batch_size", 64)),
        shuffle=False,
        num_workers=int(cfg.get("loader", {}).get("num_workers", 4)),
        collate_fn=unlabeled_collate,
    )
    embeddings, metas = extract_features(
        encoder, loader, device, use_shape_features=bool(cfg.get("use_shape_features", False))
    )
    logger.info("embedding: %s", tuple(embeddings.shape))

    cluster_kwargs = {
        "pca_dim": int(cfg.get("pca_dim", 32)),
        "whiten": bool(cfg.get("whiten", False)),
        "n_init": int(cfg.get("n_init", 10)),
    }
    fit_index = random_fit_index(
        len(metas), int(cfg.get("fit_max_samples", 0)), seed=int(cfg.get("seed", 42))
    )
    if fit_index is not None:
        logger.info("随机抽 %d 个样本拟合质心（拟合后仍分配全部样本）", len(fit_index))

    if args.sweep_k:
        sweep_clusters(embeddings, args.sweep_k, cluster_kwargs)
        return 0

    result = cluster_features(
        embeddings, num_clusters, seed=int(cfg.get("seed", 42)), fit_index=fit_index, **cluster_kwargs
    )
    stability = cluster_stability(embeddings, num_clusters, seeds=(0, 7), **cluster_kwargs)
    logger.info(
        "簇大小: %s | 轮廓系数: %.4f | 换种子稳定性: %.4f",
        result.sizes,
        result.silhouette,
        stability,
    )

    if args.save_features:
        torch.save({"embeddings": embeddings, "metas": metas}, work_dir / "features.pt")

    write_outputs(cfg, work_dir, metas, result, num_clusters, stability)
    logger.info("结果已写入 %s（clusters/ 下是按簇建的软链目录，直接翻图即可核对）", work_dir)
    for index, size in enumerate(result.sizes):
        mean_distance = float(result.distances[result.labels == index].mean()) if size else float("nan")
        logger.info("  cluster_%02d  %6d 张  平均距质心 %.3f", index, size, mean_distance)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())