"""GraphMVP 离线 embedding (300d)

依赖：torch-geometric, ogb
权重：https://github.com/chao1224/GraphMVP → Drive 文件夹内
      GraphMVP_complate_features_for_regression/GraphMVP/pretraining_model.pth

注意：新版 OGB 的 chirality 类别数为 5，而 checkpoint 为 4，
      加载时按形状逐层对齐（前 4 行用预训练，新增类别保持随机初始化）。

用法：
    python encode_graphmvp.py --ckpt pretraining_model.pth \
        --csv data/interim/internal.csv --out data/interim/graphmvp_emb.npy
"""
import argparse
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
from ogb.utils.mol import smiles2graph
from torch_geometric.data import Batch, Data
from torch_geometric.nn import MessagePassing, global_mean_pool


class GINConv(MessagePassing):
    def __init__(self, emb_dim: int) -> None:
        super().__init__(aggr="add")
        self.mlp = nn.Sequential(
            nn.Linear(emb_dim, 2 * emb_dim), nn.BatchNorm1d(2 * emb_dim),
            nn.ReLU(), nn.Linear(2 * emb_dim, emb_dim),
        )
        self.eps = nn.Parameter(torch.Tensor([0]))
        self.bond_encoder = BondEncoder(emb_dim=emb_dim)

    def forward(self, x, edge_index, edge_attr):
        e = self.bond_encoder(edge_attr)
        return self.mlp((1 + self.eps) * x + self.propagate(edge_index, x=x, edge_attr=e))

    def message(self, x_j, edge_attr):
        return F.relu(x_j + edge_attr)

    def update(self, aggr_out):
        return aggr_out


class GNNComplete(nn.Module):
    def __init__(self, num_layer: int = 5, emb_dim: int = 300, drop_ratio: float = 0.0) -> None:
        super().__init__()
        self.num_layer, self.drop_ratio = num_layer, drop_ratio
        self.atom_encoder = AtomEncoder(emb_dim)
        self.gnns = nn.ModuleList([GINConv(emb_dim) for _ in range(num_layer)])
        self.batch_norms = nn.ModuleList([nn.BatchNorm1d(emb_dim) for _ in range(num_layer)])

    def forward(self, x, edge_index, edge_attr):
        h = self.atom_encoder(x)
        for layer in range(self.num_layer):
            h = self.gnns[layer](h, edge_index, edge_attr)
            h = self.batch_norms[layer](h)
            h = F.dropout(h if layer == self.num_layer - 1 else F.relu(h),
                          self.drop_ratio, training=self.training)
        return h


def load_weights(model: nn.Module, ckpt_path: str) -> None:
    sd = torch.load(ckpt_path, map_location="cpu")
    msd = model.state_dict()
    loaded, skipped, partial = 0, [], []
    for k, v in sd.items():
        if k not in msd:
            continue
        if msd[k].shape == v.shape:
            msd[k] = v
            loaded += 1
        elif v.dim() == 2 and msd[k].shape[1] == v.shape[1] and msd[k].shape[0] > v.shape[0]:
            msd[k][: v.shape[0]] = v          # OGB 类别数变化，前 N 行用预训练
            partial.append(k)
            loaded += 1
        else:
            skipped.append((k, tuple(v.shape), tuple(msd[k].shape)))
    model.load_state_dict(msd)
    print(f"已加载 {loaded}/{len(sd)} 个张量 | 部分加载: {partial} | 跳过: {skipped}")
    assert not skipped, "有张量形状不兼容"


def main(ckpt: str, csv_path: str, out_npy: str, col: str = "smiles", batch: int = 64) -> None:
    model = GNNComplete().cuda().eval()
    load_weights(model, ckpt)

    smi = pd.read_csv(csv_path, low_memory=False)[col].astype(str).tolist()
    uniq = list(dict.fromkeys(smi))
    print(f"{len(smi)} 行, {len(uniq)} 唯一分子")

    lut = {}
    with torch.no_grad():
        for i in range(0, len(uniq), batch):
            datas = []
            for s in uniq[i:i + batch]:
                g = smiles2graph(s)
                datas.append(Data(
                    x=torch.tensor(g["node_feat"], dtype=torch.long),
                    edge_index=torch.tensor(g["edge_index"], dtype=torch.long),
                    edge_attr=torch.tensor(g["edge_feat"], dtype=torch.long),
                ))
            b = Batch.from_data_list(datas).cuda()
            hg = global_mean_pool(model(b.x, b.edge_index, b.edge_attr), b.batch).cpu().numpy()
            for j, s in enumerate(uniq[i:i + batch]):
                lut[s] = hg[j]

    E = np.stack([lut[s] for s in smi])
    np.save(out_npy, E)
    print(f"{out_npy} {E.shape}")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", required=True)
    p.add_argument("--csv", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--col", default="smiles")
    a = p.parse_args()
    main(a.ckpt, a.csv, a.out, a.col)