#!/usr/bin/env python
"""Inference for the DPP-4 Kd LTN (Delower/dpp4-kd-ltn).

Predicts pKd for human DPP-4 (UniProt P27487 / CHEMBL284) from a SMILES string.
Features are exactly those used in training: Morgan fingerprint (r=2, 2048 bits)
concatenated with the fixed DPP-4 ESM-2 650M embedding (1280 dims) = 3328, no scaling.

Usage:
  python predict_kd.py "Cn1c(=O)cc(N2CCC[C@@H](N)C2)n(Cc2ccccc2C#N)c1=O"
  python predict_kd.py --csv compounds.csv --smiles-col canonical_smiles --out preds.csv

Outputs per compound: predicted pKd, Kd in nM (= 10^(9-pKd)),
deltaG in kcal/mol (= -1.3642 * pKd), and potency class
(potent pKd>=7 / moderate 6-7 / weak <6).

Requires: torch, rdkit, numpy. Files ltn_stage2.pt, esm2_dpp4.npz and config.json
must sit next to this script (or pass --model-dir).
"""
import argparse, csv, json, os, sys
import numpy as np
import torch
import torch.nn as nn

HERE = os.path.dirname(os.path.abspath(__file__))


class LtnNet(nn.Module):
    """Same architecture as training (train_kd_ltn.py)."""
    def __init__(self, in_dim=3328):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Linear(in_dim, 768), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(768, 384), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(384, 192), nn.GELU())
        self.reg = nn.Linear(192, 1)
        self.clf = nn.Linear(192, 3)

    def forward(self, x):
        h = self.trunk(x)
        return self.reg(h).squeeze(-1), self.clf(h)


def band_name(pkd):
    return "potent" if pkd >= 7.0 else ("moderate" if pkd >= 6.0 else "weak")


def fingerprints(smiles_list):
    from rdkit import Chem
    from rdkit.Chem import rdFingerprintGenerator
    gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
    out = np.zeros((len(smiles_list), 2048), dtype=np.float32)
    for i, s in enumerate(smiles_list):
        m = Chem.MolFromSmiles(s)
        if m is not None:
            out[i] = np.asarray(gen.GetFingerprint(m), dtype=np.float32)
    return out


class KdLtnPredictor:
    def __init__(self, model_dir=HERE):
        with open(os.path.join(model_dir, "config.json")) as f:
            self.config = json.load(f)
        z = np.load(os.path.join(model_dir, "esm2_dpp4.npz"), allow_pickle=True)
        keys = z["target_keys"].tolist()
        key = self.config["target"]["esm2_target_key"]
        self.protein_vec = z["embeddings"][keys.index(key)].astype(np.float32)
        self.net = LtnNet(self.config["features"]["total_dim"])
        sd = torch.load(os.path.join(model_dir, self.config["weights_file"]),
                        map_location="cpu", weights_only=True)
        self.net.load_state_dict(sd)
        self.net.eval()

    def predict(self, smiles_list):
        fp = fingerprints(smiles_list)
        X = np.concatenate([fp, np.tile(self.protein_vec, (len(smiles_list), 1))], axis=1)
        with torch.no_grad():
            pkd, _ = self.net(torch.tensor(X, dtype=torch.float32))
        out = []
        for s, p in zip(smiles_list, pkd.numpy()):
            p = float(p)
            out.append({
                "smiles": s,
                "predicted_pKd": round(p, 4),
                "predicted_Kd_nM": 10 ** (9 - p),
                "predicted_deltaG_kcal_mol": round(-1.3642 * p, 4),
                "potency_class": band_name(p),
            })
        return out


def main():
    ap = argparse.ArgumentParser(description="DPP-4 Kd LTN inference (human DPP-4, P27487)")
    ap.add_argument("smiles", nargs="*", help="one or more SMILES strings")
    ap.add_argument("--csv", help="input CSV with a SMILES column")
    ap.add_argument("--smiles-col", default="smiles", help="SMILES column name for --csv")
    ap.add_argument("--out", help="write results to this CSV instead of stdout")
    ap.add_argument("--model-dir", default=HERE)
    args = ap.parse_args()

    smiles = list(args.smiles)
    if args.csv:
        with open(args.csv, newline="") as f:
            smiles += [r[args.smiles_col] for r in csv.DictReader(f)]
    if not smiles:
        ap.error("give at least one SMILES or --csv")

    rows = KdLtnPredictor(args.model_dir).predict(smiles)
    fields = ["smiles", "predicted_pKd", "predicted_Kd_nM", "predicted_deltaG_kcal_mol", "potency_class"]
    if args.out:
        with open(args.out, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=fields)
            w.writeheader()
            for r in rows:
                w.writerow({**r, "predicted_Kd_nM": f"{r['predicted_Kd_nM']:.4g}"})
        print(f"wrote {len(rows)} predictions to {args.out}")
    else:
        for r in rows:
            print(f"pKd {r['predicted_pKd']:7.4f} | Kd {r['predicted_Kd_nM']:12.4g} nM | "
                  f"dG {r['predicted_deltaG_kcal_mol']:8.4f} kcal/mol | {r['potency_class']:8s} | {r['smiles']}")


if __name__ == "__main__":
    main()