#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
结构感知 + 分布对齐 的 DrugBank 数据集划分脚本
- 读取 CSV（至少包含 SELFIES 或 SMILES 之一，建议还包含 qed）
- 计算 SMILES（若仅有 SELFIES），构建 RDKit Mol
- 计算 Bemis-Murcko Scaffold、分子量等基本属性
- 以 scaffold 为分组单元进行分层划分（train/val/test）
- 对齐 QED/MW 分布（分箱 + 贪心平衡）
- 输出三份 CSV：split_train.csv / split_val.csv / split_test.csv

结构感知 + 分布对齐 的 DrugBank 数据集划分脚本（修正版）
- 关键修复：分配策略改为【容量优先 + 分布对齐】，避免 TRAIN=0 的偏置
"""

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
结构感知 + 分布对齐 的 DrugBank 数据集划分脚本（修正版）
- 关键修复：分配策略改为【容量优先 + 分布对齐】，避免 TRAIN=0 的偏置
"""

import argparse
from pathlib import Path
import math
import hashlib
from collections import defaultdict, Counter

import pandas as pd
import numpy as np

from rdkit import Chem
from rdkit.Chem.Scaffolds import MurckoScaffold
from rdkit.Chem import Descriptors

try:
    import selfies as sf
except Exception:
    sf = None

def to_smiles_from_selfies(s):
    if pd.isna(s) or not isinstance(s, str) or not s.strip():
        return None
    if sf is None:
        return None
    try:
        smi = sf.decoder(s)
        mol = Chem.MolFromSmiles(smi)
        if mol is None:
            return None
        return Chem.MolToSmiles(mol)
    except Exception:
        return None

def canonicalize_smiles(s):
    if pd.isna(s) or not isinstance(s, str) or not s.strip():
        return None
    mol = Chem.MolFromSmiles(s)
    if mol is None:
        return None
    return Chem.MolToSmiles(mol)

def bemis_murcko_scaffold(smiles):
    if smiles is None:
        return None
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    try:
        scaf = MurckoScaffold.GetScaffoldForMol(mol)
        if scaf is None:
            return None
        return Chem.MolToSmiles(scaf)
    except Exception:
        return None

def calc_basic_props(smiles):
    mol = Chem.MolFromSmiles(smiles) if smiles else None
    if mol is None:
        return (np.nan, np.nan, np.nan, np.nan)
    mw = Descriptors.MolWt(mol)
    tpsa = Descriptors.TPSA(mol)
    logp = Descriptors.MolLogP(mol)
    heavy = mol.GetNumHeavyAtoms()
    return (mw, tpsa, logp, heavy)

def hash_str(x):
    return int(hashlib.md5(x.encode('utf-8')).hexdigest(), 16)

def bin_by_edges(x, edges):
    if isinstance(x, float) and math.isnan(x):
        return -1
    for i in range(1, len(edges)):
        if x <= edges[i]:
            return i-1
    return len(edges) - 2

def hist_distance_L2(h1, h2):
    keys = set(h1.keys()) | set(h2.keys())
    return math.sqrt(sum((h1.get(k,0) - h2.get(k,0))**2 for k in keys))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--in-csv", required=True)
    ap.add_argument("--out-dir", default=".")
    ap.add_argument("--seed", type=int, default=2025)
    ap.add_argument("--train-ratio", type=float, default=0.8)
    ap.add_argument("--val-ratio", type=float, default=0.1)
    ap.add_argument("--test-ratio", type=float, default=0.1)
    ap.add_argument("--n_qed_bins", type=int, default=5)
    ap.add_argument("--n_mw_bins", type=int, default=5)
    # 新增：控制是否按 scaffold 大小降序分配（推荐 True）
    ap.add_argument("--largest-first", action="store_true", default=True)
    args = ap.parse_args()

    np.random.seed(args.seed)

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    df = pd.read_csv(args.in_csv)

    def col_like(name_candidates):
        for nc in name_candidates:
            idx = [i for i,c in enumerate(df.columns) if c.lower()==nc.lower()]
            if idx:
                return df.columns[idx[0]]
        return None

    col_smiles = col_like(["smiles", "canonical_smiles"])
    col_selfies = col_like(["selfies"])
    col_id = col_like(["id","molid","drug_id"])
    col_qed = col_like(["qed"])

    if col_smiles is None:
        if col_selfies is None:
            raise SystemExit("需要 SMILES 或 SELFIES 至少一个列。")
        df["SMILES"] = df[col_selfies].map(to_smiles_from_selfies)
        col_smiles = "SMILES"
    else:
        df["SMILES"] = df[col_smiles]

    df["SMILES"] = df["SMILES"].map(canonicalize_smiles)
    df = df[~df["SMILES"].isna()].copy()
    df = df.drop_duplicates(subset=["SMILES"])

    props = df["SMILES"].map(calc_basic_props)
    df["MW"] = [p[0] for p in props]
    df["TPSA"] = [p[1] for p in props]
    df["LogP"] = [p[2] for p in props]
    df["HeavyAtoms"] = [p[3] for p in props]

    df["Scaffold"] = df["SMILES"].map(bemis_murcko_scaffold)
    df["GroupKey"] = np.where(df["Scaffold"].isna(), df["SMILES"], df["Scaffold"])

    if col_qed is None:
        df["qed"] = np.nan
    else:
        df["qed"] = pd.to_numeric(df[col_qed], errors="coerce")

    def compute_edges(series, n_bins):
        ser = series.dropna()
        if len(ser) < n_bins:
            return np.linspace(ser.min() if len(ser)>0 else 0,
                               ser.max() if len(ser)>0 else 1,
                               n_bins+1)
        qs = np.linspace(0,1,n_bins+1)
        return np.quantile(ser, qs)

    qed_edges = compute_edges(df["qed"], args.n_qed_bins)
    mw_edges  = compute_edges(df["MW"],  args.n_mw_bins)

    df["QED_bin"] = df["qed"].map(lambda x: bin_by_edges(x, qed_edges))
    df["MW_bin"]  = df["MW"].map(lambda x: bin_by_edges(x, mw_edges))
    df["Strata"] = list(zip(df["QED_bin"], df["MW_bin"]))

    group2idx = defaultdict(list)
    for i, g in enumerate(df["GroupKey"]):
        group2idx[g].append(i)

    group_hist = {g: Counter(df.loc[idxs, "Strata"].tolist())
                  for g, idxs in group2idx.items()}
    global_hist = Counter(df["Strata"].tolist())

    total = len(df)
    target_counts = {
        "train": int(round(total * args.train_ratio)),
        "val":   int(round(total * args.val_ratio)),
        "test":  int(round(total * args.test_ratio)),
    }
    diff = total - sum(target_counts.values())
    if diff != 0:
        target_counts["train"] += diff

    global_ratio = {
        "train": args.train_ratio,
        "val":   args.val_ratio,
        "test":  args.test_ratio,
    }
    target_hist = {split: {k: v * r for k, v in global_hist.items()}
                   for split, r in global_ratio.items()}

    splits = {"train": [], "val": [], "test": []}
    split_counts = {"train": 0, "val": 0, "test": 0}
    split_hist = {"train": Counter(), "val": Counter(), "test": Counter()}

    scaffolds = list(group2idx.keys())
    rng = np.random.RandomState(args.seed)

    # 关键：按组大小降序（先放大的）
    if args.largest_first:
        scaffolds.sort(key=lambda g: len(group2idx[g]), reverse=True)
    else:
        rng.shuffle(scaffolds)

    def remaining_capacity(split):
        return target_counts[split] - split_counts[split]

    for g in scaffolds:
        idxs = group2idx[g]
        g_size = len(idxs)
        g_hist = group_hist[g]

        # 计算每个 split 的剩余容量
        caps = {s: remaining_capacity(s) for s in ["train","val","test"]}

        # 优先考虑【仍有正容量】的 split
        positive_splits = [s for s,c in caps.items() if c > 0]

        # 如果都有容量<=0，说明目标配额已满；退化为“最小溢出优先”
        candidate_splits = positive_splits if positive_splits else ["train","val","test"]

        # 在候选中，综合“容量优先 + 分布对齐（L2更小）”
        best = None
        for s in candidate_splits:
            cap = caps[s]
            proj_hist = split_hist[s] + g_hist
            l2 = hist_distance_L2(proj_hist, target_hist[s])

            # 排序键：
            #   1) -max(cap, 0): 正容量越大越好；都负时此项=0
            #   2) abs(min(cap, 0)): 负容量越小（溢出越少）越好
            #   3) l2: 分布差越小越好
            key = (-max(cap, 0), abs(min(cap, 0)), l2)

            if (best is None) or (key < best[0]):
                best = (key, s)

        chosen = best[1]

        splits[chosen].append(g)
        split_counts[chosen] += g_size
        split_hist[chosen] += g_hist

    # 收集索引
    idx_train, idx_val, idx_test = [], [], []
    for g in splits["train"]:
        idx_train.extend(group2idx[g])
    for g in splits["val"]:
        idx_val.extend(group2idx[g])
    for g in splits["test"]:
        idx_test.extend(group2idx[g])

    df_train = df.iloc[sorted(idx_train)].copy()
    df_val   = df.iloc[sorted(idx_val)].copy()
    df_test  = df.iloc[sorted(idx_test)].copy()

    def report(name, sub):
        print(f"\n== {name} ==")
        print(f"size = {len(sub)}")
        h = Counter(sub["Strata"].tolist())
        for k, v in h.most_common(8):
            print(f"  {k}: {v}")
        print("  QED mean/std:", np.nanmean(sub["qed"]), np.nanstd(sub["qed"]))
        print("  MW  mean/std:", np.nanmean(sub["MW"]),  np.nanstd(sub["MW"]))

    print(f"Total: {len(df)} (targets: {target_counts})")
    report("TRAIN", df_train)
    report("VAL",   df_val)
    report("TEST",  df_test)

    out_train = out_dir / "split_train.csv"
    out_val   = out_dir / "split_val.csv"
    out_test  = out_dir / "split_test.csv"
    df_train.to_csv(out_train, index=False)
    df_val.to_csv(out_val, index=False)
    df_test.to_csv(out_test, index=False)

    edges = pd.DataFrame({
        "qed_edges": list(qed_edges) + [np.nan]*(max(len(mw_edges),len(qed_edges))-len(qed_edges)),
        "mw_edges":  list(mw_edges)  + [np.nan]*(max(len(mw_edges),len(qed_edges))-len(mw_edges)),
    })
    edges.to_csv(out_dir / "bin_edges_qed_mw.csv", index=False)

    # 额外提示：打印容量偏差
    print("\nFinal counts vs targets:")
    for s in ["train","val","test"]:
        print(f"  {s}: {split_counts[s]} / {target_counts[s]} (delta={split_counts[s]-target_counts[s]})")

    print(f"\nSaved:\n  {out_train}\n  {out_val}\n  {out_test}\n  {out_dir/'bin_edges_qed_mw.csv'}")
    print("\n验证/测试阶段，仅使用 split_val / split_test 作为 feed-chemical。")

if __name__ == "__main__":
    main()