from rdkit import Chem
from rdkit.Chem import Crippen, Descriptors, Lipinski, rdMolDescriptors


def local_property_metrics(smiles: str) -> dict[str, float | int]:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError("Cannot predict invalid molecule")
    return {
        "molecular_weight": round(Descriptors.MolWt(mol), 2),
        "logp": round(Crippen.MolLogP(mol), 2),
        "tpsa": round(rdMolDescriptors.CalcTPSA(mol), 2),
        "hbd": int(Lipinski.NumHDonors(mol)),
        "hba": int(Lipinski.NumHAcceptors(mol)),
        "rotatable_bonds": int(Lipinski.NumRotatableBonds(mol)),
    }


def developability_score(metrics: dict[str, float | int]) -> float:
    score = 1.0
    if metrics["molecular_weight"] > 500:
        score -= 0.20
    if metrics["logp"] > 5:
        score -= 0.20
    if metrics["tpsa"] > 140:
        score -= 0.15
    if metrics["hbd"] > 5:
        score -= 0.15
    if metrics["hba"] > 10:
        score -= 0.15
    if metrics["rotatable_bonds"] > 10:
        score -= 0.10
    return max(0.0, round(score, 3))