"""Inference helper for models uploaded by hf_upload_universal.py.

Usage:
    from infer import predict_smiles
    predict_smiles(model, ["CCO", "CC(=O)O"])

The model was trained on the features listed in config.json
(e.g. Morgan fingerprint 1024 + physicochemical descriptors).
This model is a sklearn Pipeline (ColumnTransformer +
ExtraTreesRegressor). LazyPredict pipelines select columns by STRING
name, so features must be passed as a pandas DataFrame with columns
feature_0..feature_1029.
Feature order: Morgan fingerprint 1024 (radius 2) first, then 6 physchem
descriptors in this order: MolWt, MolLogP, TPSA, NumHDonors,
NumHAcceptors, NumRotatableBonds.
Unpickling model.pkl requires lazypredict (see requirements.txt).
"""
import pickle

import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit.Chem import AllChem, Descriptors


def featurize(smiles_list, features):
    rows = []
    for smi in smiles_list:
        mol = Chem.MolFromSmiles(smi)
        if mol is None:
            raise ValueError(f"Invalid SMILES: {smi}")
        parts = []
        if "morgan1024" in features:
            fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024)
            parts.append(np.asarray(fp, dtype=float))
        if "physchem" in features:
            parts.append(np.array([
                Descriptors.MolWt(mol), Descriptors.MolLogP(mol),
                Descriptors.TPSA(mol), Descriptors.NumHDonors(mol),
                Descriptors.NumHAcceptors(mol),
                Descriptors.NumRotatableBonds(mol),
            ], dtype=float))
        rows.append(np.concatenate(parts) if parts else np.array([]))
    X = np.vstack(rows)
    return pd.DataFrame(X, columns=[f"feature_{i}" for i in range(X.shape[1])])


def load_model(path="model.pkl"):
    with open(path, "rb") as f:
        return pickle.load(f)


def predict_smiles(model, smiles_list, features=("morgan1024", "physchem")):
    X = featurize(smiles_list, features)
    return model.predict(X)