from __future__ import annotations

import argparse
import json
from pathlib import Path

import joblib
import pandas as pd

from src.features.rdkit_features import featurize_smiles, FeatConfig


def load_artifacts(artifact_dir: Path):
    meta = json.loads((artifact_dir / "metadata.json").read_text(encoding="utf-8"))
    pre = joblib.load(artifact_dir / "preprocess.joblib")
    models = {t: joblib.load(artifact_dir / "models" / f"{t}.joblib") for t in meta["targets"]}
    return meta, pre, models


def main():
    ap = argparse.ArgumentParser(description="Predict ADMET from SMILES using saved artifacts.")
    ap.add_argument("--artifact_dir",type=str,default="artifacts/admet_smiles_merge_didb_seed202")
    ap.add_argument("--smiles", type=str, required=True)
    ap.add_argument("--smiles_col", type=str, default="smiles")
    ap.add_argument("--threshold", type=float, default=0.5)
    args = ap.parse_args()

    meta, pre, models = load_artifacts(Path(args.artifact_dir))

    cfg = FeatConfig(fp_radius=meta["features"]["fp"]["radius"], fp_nbits=meta["features"]["fp"]["nbits"])
    df = pd.DataFrame({args.smiles_col: [args.smiles]})

    X, valid = featurize_smiles(df, smiles_col=args.smiles_col, config=cfg, add_physchem=True, drop_invalid=False)
    if not valid.all():
        raise ValueError("Invalid SMILES.")

    feature_names = X.columns.tolist()
    Xt = pd.DataFrame(pre.transform(X), columns=feature_names, index=X.index)


    out = {"smiles": args.smiles, "pred": {}, "meta": {"model_version": meta["model_version"]}}
    for t, m in models.items():
        prob = float(m.predict_proba(Xt)[:, 1][0])
        pred = int(prob >= args.threshold)

        inv = {v: k for k, v in meta["label_maps"][t]["mapping"].items()}
        out["pred"][t] = {"prob": prob, "pred": pred, "label": inv[pred]}

    print(json.dumps(out, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()