"""
predict() entry point for the Tox21 leaderboard (ml-jku/tox21_leaderboard).
Input: list of SMILES. Output: nested dict {smiles: {target: probability}}.
Unparseable molecules return 0.5 for every target.

Model: VIDRAFT rich-feature XGBoost (Morgan2048 + MACCS167 + RDKit 2D descriptors),
one XGBoost per Tox21 task, trained on the official ml-jku/tox21 train split.
Feature order is loaded from checkpoints/feat_spec.json so it is identical to training
regardless of the local RDKit version.
"""
import os
import json
import numpy as np
import xgboost as xgb
from rdkit import Chem, RDLogger, DataStructs
from rdkit.Chem import AllChem, MACCSkeys
from rdkit.ML.Descriptors import MoleculeDescriptors

RDLogger.DisableLog("rdApp.*")

TASKS = ["NR-AhR", "NR-AR", "NR-AR-LBD", "NR-Aromatase", "NR-ER", "NR-ER-LBD",
         "NR-PPAR-gamma", "SR-ARE", "SR-ATAD5", "SR-HSE", "SR-MMP", "SR-p53"]

_HERE = os.path.dirname(os.path.abspath(__file__))
_CK = os.path.join(_HERE, "checkpoints")
_spec = json.load(open(os.path.join(_CK, "feat_spec.json")))
_DESC = _spec["desc"]
_calc = MoleculeDescriptors.MolecularDescriptorCalculator(_DESC)
_med = np.load(os.path.join(_CK, "feat_median.npy"))
_FPN, _MACCSN = 2048, 167
_DIM = _FPN + _MACCSN + len(_DESC)

_models = {}
for _t in TASKS:
    _m = xgb.XGBClassifier()
    _m.load_model(os.path.join(_CK, "xgb_%s.json" % _t))
    _models[_t] = _m


def _feat_one(smi):
    mol = Chem.MolFromSmiles(str(smi))
    if mol is None:
        return None
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=_FPN)
    a = np.zeros(_FPN, np.int8); DataStructs.ConvertToNumpyArray(fp, a)
    mc = MACCSkeys.GenMACCSKeys(mol)
    b = np.zeros(_MACCSN, np.int8); DataStructs.ConvertToNumpyArray(mc, b)
    d = np.array(_calc.CalcDescriptors(mol), dtype=np.float32)
    return np.concatenate([a.astype(np.float32), b.astype(np.float32), d])


def predict(smiles_list):
    n = len(smiles_list)
    X = np.zeros((n, _DIM), np.float32)
    ok = np.zeros(n, bool)
    for i, s in enumerate(smiles_list):
        f = _feat_one(s)
        if f is not None:
            X[i] = f; ok[i] = True
    X[~np.isfinite(X)] = np.nan
    bad = np.where(np.isnan(X))
    if bad[0].size:
        X[bad] = np.take(_med, bad[1])
    probs = {t: _models[t].predict_proba(X)[:, 1] for t in TASKS}
    out = {}
    for i, s in enumerate(smiles_list):
        out[s] = {t: (float(probs[t][i]) if ok[i] else 0.5) for t in TASKS}
    return out