# -*- coding: utf-8 -*-
from __future__ import annotations
import math
import random
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any, Tuple

import numpy as np
import pandas as pd

from rdkit import Chem, RDLogger
from rdkit.Chem import rdMolDescriptors
from rdkit import DataStructs
from rdkit.ML.Cluster import Butina

from sklearn.cluster import DBSCAN, AgglomerativeClustering

try:
    from scipy.cluster.hierarchy import linkage, fcluster
except Exception:
    linkage = None
    fcluster = None

# 静音 RDKit
RDLogger.DisableLog('rdApp.*')


# ---------- 指纹 & 相似度工具 ----------
@dataclass
class FPConfig:
    radius: int = 2          # ECFP 半径（2=ECFP4）
    n_bits: int = 2048       # 位数（1024 或 2048 常用）
    use_count: bool = False  # True: 计数 FP；False: bit 向量（推荐用于 Tanimoto）


def smiles_to_mols(smiles: List[str]) -> Tuple[List[Chem.Mol], List[int]]:
    mols, idx = [], []
    for i, smi in enumerate(smiles):
        m = Chem.MolFromSmiles(str(smi))
        if m is not None:
            mols.append(m)
            idx.append(i)
    return mols, idx


def mols_to_ecfp(mols: List[Chem.Mol], cfg: FPConfig):
    fps = []
    if cfg.use_count:
        gen = rdMolDescriptors.GetMorganFingerprint
        for m in mols:
            fps.append(gen(m, cfg.radius))
    else:
        for m in mols:
            fps.append(rdMolDescriptors.GetMorganFingerprintAsBitVect(m, cfg.radius, nBits=cfg.n_bits))
    return fps


def tanimoto(a, b) -> float:
    return DataStructs.TanimotoSimilarity(a, b)


def bulk_tanimoto_to_many(fp, fps) -> List[float]:
    # 对大量 pair 更快
    return DataStructs.BulkTanimotoSimilarity(fp, fps)


# ---------- 稳定可扩展的聚类（首选） ----------
def cluster_butina(fps, sim_cutoff: float = 0.6) -> np.ndarray:
    """
    RDKit Butina 聚类
    fps: RDKit Fingerprint list
    sim_cutoff: 相似度阈值（如 0.6~0.8）
    返回：簇标签 (0..K-1)
    """
    n = len(fps)
    # 构建完整的压缩距离矩阵
    dists = []
    for i in range(1, n):
        sims = DataStructs.BulkTanimotoSimilarity(fps[i], fps[:i])
        dists.extend([1.0 - s for s in sims])  # 全部保存

    # 调用 Butina 聚类
    cs = Butina.ClusterData(dists, n, 1.0 - sim_cutoff, isDistData=True)

    # 转换成标签数组
    labels = np.full(n, -1, dtype=int)
    for cid, members in enumerate(cs):
        for m in members:
            labels[m] = cid

    # 把孤立点分配为独立簇（便于后续处理）
    if np.any(labels == -1):
        max_id = labels.max() if labels.max() >= 0 else -1
        for i in np.where(labels == -1)[0]:
            max_id += 1
            labels[i] = max_id

    return labels


def cluster_dbscan_threshold(fps, sim_cutoff: float = 0.6, eps: float = 0.5, min_samples: int = 5) -> np.ndarray:
    """
    用相似度阈值构稀疏邻接，再交给 DBSCAN(metric='precomputed')。
    """
    n = len(fps)
    dm = np.full((n, n), 1.0, dtype=np.float32)
    np.fill_diagonal(dm, 0.0)
    for i in range(n):
        sims = bulk_tanimoto_to_many(fps[i], fps)
        d = 1.0 - np.array(sims, dtype=np.float32)
        near = d <= (1.0 - sim_cutoff)
        dm[i, near] = d[near]
    clt = DBSCAN(metric='precomputed', eps=eps, min_samples=min_samples, n_jobs=-1)
    labels = clt.fit_predict(dm)
    # 保持 DBSCAN 的 -1 为噪声；如需改为独立簇可在外面处理
    return labels


# ---------- 小数据可选：完整距离矩阵 ----------
def distance_matrix_from_fps(fps) -> np.ndarray:
    n = len(fps)
    dm = np.zeros((n, n), dtype=np.float32)
    for i in range(n):
        sims = bulk_tanimoto_to_many(fps[i], fps[i+1:])
        if len(sims):
            d = 1.0 - np.array(sims, dtype=np.float32)
            dm[i, i+1:] = d
            dm[i+1:, i] = d
    return dm


def cluster_agglomerative_precomputed(fps, n_clusters=5, linkage_method='average') -> np.ndarray:
    dm = distance_matrix_from_fps(fps)
    clt = AgglomerativeClustering(n_clusters=n_clusters, metric='precomputed', linkage=linkage_method)
    return clt.fit_predict(dm)


def cluster_scipy_linkage_cut(fps, method='average', t=0.7, criterion='distance') -> np.ndarray:
    if linkage is None:
        raise RuntimeError("scipy 未安装，无法使用 linkage。")
    dm = distance_matrix_from_fps(fps)
    n = dm.shape[0]
    iu = np.triu_indices(n, k=1)
    condensed = dm[iu]
    Z = linkage(condensed, method=method)
    labels = fcluster(Z, t=t, criterion=criterion) - 1
    return labels


# ---------- 代表分子 ----------
def pick_medoid_indices(idx: List[int], fps) -> int:
    """返回簇内 medoid 的下标（全对全，适合簇不大时）"""
    if len(idx) == 1:
        return idx[0]
    sub = idx
    n = len(sub)
    sums = []
    for i in range(n):
        sims = [tanimoto(fps[sub[i]], fps[sub[j]]) for j in range(n)]
        dsum = float(np.sum(1.0 - np.array(sims)))
        sums.append(dsum)
    return sub[int(np.argmin(sums))]


def pick_maxmin_indices(idx: List[int], fps, k=1) -> List[int]:
    """在指定索引集合里挑 k 个多样性最大的（MaxMin）"""
    if not idx:
        return []
    chosen = [idx[0]]
    while len(chosen) < min(k, len(idx)):
        best_j, best_min_d = None, -1.0
        for j in idx:
            if j in chosen:
                continue
            mind = 1.0
            for i in chosen:
                d = 1.0 - tanimoto(fps[i], fps[j])
                if d < mind:
                    mind = d
            if mind > best_min_d:
                best_min_d, best_j = mind, j
        chosen.append(best_j)
    return chosen


def select_representatives(labels: np.ndarray, fps, per_cluster: int = 1,
                           strategy: str = "medoid") -> List[int]:
    reps = []
    for c in np.unique(labels):
        members = np.where(labels == c)[0].tolist()
        if not members:
            continue
        if strategy == "medoid":
            reps.append(pick_medoid_indices(members, fps))
        elif strategy == "maxmin":
            reps.extend(pick_maxmin_indices(members, fps, k=per_cluster))
        else:
            raise ValueError("strategy ∈ {'medoid','maxmin'}")
    return sorted(set(reps))


# ---------- 评估指标 ----------
@dataclass
class ClusterStats:
    n_samples: int
    n_clusters: int
    largest_cluster_size: int
    largest_cluster_ratio: float
    sizes: Dict[int, int] = field(default_factory=dict)


def cluster_stats(labels: np.ndarray) -> ClusterStats:
    n = len(labels)
    vals, counts = np.unique(labels, return_counts=True)
    largest = int(counts.max())
    return ClusterStats(
        n_samples=int(n),
        n_clusters=int(len(vals)),
        largest_cluster_size=largest,
        largest_cluster_ratio=float(largest / n),
        sizes={int(v): int(c) for v, c in zip(vals, counts)}
    )


# ---------- 对外 API ----------
@dataclass
class TanimotoClusteringAPI:
    fp_cfg: FPConfig = field(default_factory=FPConfig)
    params: dict = field(default_factory=dict)

    def fit_from_smiles(self, smiles: List[str], method: str = "butina",
                        method_kwargs: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        返回：{ 'labels': np.ndarray, 'keep_idx': List[int], 'fps': List[FP], 'stats': ClusterStats }
        """
        method_kwargs = method_kwargs or {}
        mols, keep_idx = smiles_to_mols(smiles)
        if len(mols) == 0:
            raise ValueError("没有有效的 SMILES。")
        fps = mols_to_ecfp(mols, self.fp_cfg)

        if method == "butina":
            labels = cluster_butina(fps, **method_kwargs)
        elif method == "dbscan_threshold":
            labels = cluster_dbscan_threshold(fps, **method_kwargs)
        elif method == "agglomerative":
            labels = cluster_agglomerative_precomputed(fps, **method_kwargs)
        elif method == "scipy_linkage":
            labels = cluster_scipy_linkage_cut(fps, **method_kwargs)
        else:
            raise ValueError("未知 method")

        stats = cluster_stats(labels)
        return {"labels": labels, "keep_idx": keep_idx, "fps": fps, "stats": stats}

    # ChemPlot 可视化（可选）
    def chemplot_visualize(self, smiles_valid: List[str], labels: np.ndarray,
                           target: Optional[pd.Series] = None,
                           target_type: Optional[str] = None,
                           random_state: int = 42):
        from chemplot import Plotter
        import matplotlib.pyplot as plt
        cp = Plotter.from_smiles(smiles_valid, sim_type="structural",
                                 target=target, target_type=target_type)
        cp.umap(random_state=random_state)
        cp._Plotter__df_2_components["clusters"] = labels
        fig = cp.visualize_plot(clusters=True)
        return cp, fig

    # 网格探索（大样本优先用 Butina / 稀疏 DBSCAN）
    def explore(self, smiles: List[str],
                butina_cuts=(0.5, 0.6, 0.7, 0.75, 0.8),
                dbscan_params=((0.6, 0.5, 5), (0.6, 0.4, 5), (0.7, 0.5, 5))
                ) -> List[Dict[str, Any]]:
        """
        返回每个配置的 stats 列表，便于挑选“最大簇占比最小”的方案。
        """
        results = []
        mols, keep_idx = smiles_to_mols(smiles)
        fps = mols_to_ecfp(mols, self.fp_cfg)

        for cut in butina_cuts:
            labels = cluster_butina(fps, sim_cutoff=cut)
            stats = cluster_stats(labels)
            results.append({"method": "butina", "params": {"sim_cutoff": cut}, "stats": stats})

        for sim_cut, eps, min_samples in dbscan_params:
            labels = cluster_dbscan_threshold(fps, sim_cutoff=sim_cut, eps=eps, min_samples=min_samples)
            stats = cluster_stats(labels)
            results.append({"method": "dbscan_threshold",
                            "params": {"sim_cutoff": sim_cut, "eps": eps, "min_samples": min_samples},
                            "stats": stats})
        return results


# ---------- 自动搜索封装（对外 API 保持一致） ----------
@dataclass
class SearchConfig:
    # 指纹配置
    radii: List[int] = field(default_factory=lambda: [2, 3])
    n_bits_list: List[int] = field(default_factory=lambda: [1024, 2048])
    # 方法搜索空间
    butina_cuts: List[float] = field(default_factory=lambda: [0.5, 0.6, 0.7, 0.75, 0.8])
    dbscan_params: List[Tuple[float, float, int]] = field(default_factory=lambda: [(0.6, 0.5, 5), (0.6, 0.4, 5), (0.7, 0.5, 5)])
    # 小样本时可选
    try_agglomerative: bool = True
    agglo_k_list: List[int] = field(default_factory=lambda: [5, 8, 10])
    try_scipy: bool = False
    scipy_t_list: List[float] = field(default_factory=lambda: [0.6, 0.7])
    linkage: str = "average"

def _score(stats: ClusterStats) -> float:
    # 越大越好：更均衡（最大簇比例小）+ 簇数量>1的奖励
    balance = 1.0 - stats.largest_cluster_ratio
    bonus = 0.1 if stats.n_clusters > 1 else 0.0
    return balance + bonus

def search_best_config(smiles: List[str], cfg: Optional[SearchConfig] = None) -> Tuple[TanimotoClusteringAPI, ClusterStats, pd.DataFrame]:
    """
    自动尝试不同 FP 半径/位数 + (Butina/DBSCAN/可选Agglo/SciPy)，
    以 1-最大簇比例 为主的评分，返回：最佳 API（已设置好 fp_cfg 与 method/参数）、其 stats、以及历史对比表。
    """
    cfg = cfg or SearchConfig()
    trials: List[Dict[str, Any]] = []

    for r in cfg.radii:
        for nb in cfg.n_bits_list:
            api = TanimotoClusteringAPI(fp_cfg=FPConfig(radius=r, n_bits=nb))
            # 1) Butina
            for cut in cfg.butina_cuts:
                res = api.fit_from_smiles(smiles, method="butina", method_kwargs={"sim_cutoff": cut})
                sc = _score(res["stats"])
                trials.append({"fp_radius": r, "fp_n_bits": nb, "method": "butina",
                               "params": {"sim_cutoff": cut}, "stats": res["stats"], "score": sc})
            # 2) 稀疏 DBSCAN
            for (sim_cut, eps, ms) in cfg.dbscan_params:
                res = api.fit_from_smiles(smiles, method="dbscan_threshold",
                                          method_kwargs={"sim_cutoff": sim_cut, "eps": eps, "min_samples": ms})
                sc = _score(res["stats"])
                trials.append({"fp_radius": r, "fp_n_bits": nb, "method": "dbscan_threshold",
                               "params": {"sim_cutoff": sim_cut, "eps": eps, "min_samples": ms},
                               "stats": res["stats"], "score": sc})
            # 3) （可选）Agglomerative（小样本或探索时）
            if cfg.try_agglomerative:
                for k in cfg.agglo_k_list:
                    res = api.fit_from_smiles(smiles, method="agglomerative",
                                              method_kwargs={"n_clusters": k, "linkage_method": cfg.linkage})
                    sc = _score(res["stats"])
                    trials.append({"fp_radius": r, "fp_n_bits": nb, "method": "agglomerative",
                                   "params": {"n_clusters": k, "linkage_method": cfg.linkage},
                                   "stats": res["stats"], "score": sc})
            # 4) （可选）SciPy linkage（需小样本 & 安装 scipy）
            if cfg.try_scipy and linkage is not None:
                for t in cfg.scipy_t_list:
                    res = api.fit_from_smiles(smiles, method="scipy_linkage",
                                              method_kwargs={"method": cfg.linkage, "t": t, "criterion": "distance"})
                    sc = _score(res["stats"])
                    trials.append({"fp_radius": r, "fp_n_bits": nb, "method": "scipy_linkage",
                                   "params": {"t": t, "method": cfg.linkage, "criterion": "distance"},
                                   "stats": res["stats"], "score": sc})

    # 排序选最优
    rows = []
    for t in trials:
        s: ClusterStats = t["stats"]
        rows.append({
            "fp_radius": t["fp_radius"],
            "fp_n_bits": t["fp_n_bits"],
            "method": t["method"],
            "params": t["params"],
            "n_samples": s.n_samples,
            "n_clusters": s.n_clusters,
            "largest_cluster_ratio": s.largest_cluster_ratio,
            "sizes": s.sizes,
            "score": t["score"],
        })
    hist = pd.DataFrame(rows).sort_values("score", ascending=False).reset_index(drop=True)

    best = hist.iloc[0]
    # 组装一个“已配置好的” API 返回
    api_best = TanimotoClusteringAPI(fp_cfg=FPConfig(radius=int(best["fp_radius"]), n_bits=int(best["fp_n_bits"])))
    # 触发一次拟合以便调用方马上拿 labels/keep_idx
    res = api_best.fit_from_smiles(smiles, method=best["method"], method_kwargs=best["params"])
    return api_best, res["stats"], hist