import joblib
import os
from .features import create_features

BASE_DIR = os.path.dirname(__file__)

# Load models
davis_model = joblib.load(os.path.join(BASE_DIR, "models/davis_model.pkl"))
receptor_model = joblib.load(os.path.join(
    BASE_DIR, "models/receptor_model.pkl"))


def predict_pkd(smiles, protein, protein_type="receptor"):
    if not isinstance(smiles, str) or not isinstance(protein, str):
        raise ValueError("SMILES and protein must be strings")

    protein_type = protein_type.lower()

    features = create_features(smiles, protein)

    if protein_type in ["davis", "kinase"]:
        model = davis_model
    elif protein_type == "receptor":
        model = receptor_model
    else:
        raise ValueError("protein_type must be 'davis'/'kinase' or 'receptor'")

    pred = model.predict(features)[0]
    return float(pred)


def predict_kd(smiles, protein, protein_type="receptor"):
    pkd = predict_pkd(smiles, protein, protein_type)

    # Convert pKd → Kd (nM)
    kd = 10 ** (-pkd) * 1e9

    return float(kd)
