# 安装: pip install rdkit-pypi torch dgl
import torch
from dgl.nn import GINConv
import dgl

# 1. 用 RDKit 将分子转为图
from rdkit import Chem
from rdkit.Chem import AllChem

def mol_to_graph(smiles: str):
    """SMILES → DGL 图"""
    mol = Chem.MolFromSmiles(smiles)
    AllChem.Compute2DCoords(mol)

    # 节点特征：原子类型
    atom_features = []
    for atom in mol.GetAtoms():
        feat = [
            atom.GetAtomicNum(),
            atom.GetDegree(),
            atom.GetFormalCharge(),
            atom.GetNumRadicalElectrons(),
            int(atom.GetIsAromatic()),
        ]
        atom_features.append(feat)

    # 边：化学键
    src, dst = [], []
    for bond in mol.GetBonds():
        src.append(bond.GetBeginAtomIdx())
        dst.append(bond.GetEndAtomIdx())
        src.append(bond.GetEndAtomIdx())  # 无向图
        dst.append(bond.GetBeginAtomIdx())

    g = dgl.graph((src, dst))
    g.ndata["feat"] = torch.tensor(atom_features, dtype=torch.float)

    return g

# 2. 简单 GNN 预测模型
class MoleculePredictor(nn.Module):
    def __init__(self, in_feats=5, hidden=64, out_feats=1):
        super().__init__()
        self.conv1 = GINConv(
            nn.Sequential(
                nn.Linear(in_feats, hidden),
                nn.ReLU(),
                nn.Linear(hidden, hidden)
            )
        )
        self.conv2 = GINConv(
            nn.Sequential(
                nn.Linear(hidden, hidden),
                nn.ReLU(),
                nn.Linear(hidden, hidden)
            )
        )
        self.pred = nn.Linear(hidden, out_feats)

    def forward(self, g):
        h = self.conv1(g, g.ndata["feat"])
        h = self.conv2(g, h)
        g.ndata["h"] = h
        hg = dgl.readout_nodes(g, "h", op="mean")  # 全局池化
        return self.pred(hg).squeeze()

# 3. 预测分子水溶性
model = MoleculePredictor(in_feats=5, out_feats=1)
smiles = "CCO"  # 乙醇
g = mol_to_graph(smiles)
pred = model(g)
print(f"预测水溶性 (logS): {pred.item():.2f}")