import numpy as np
import os
import torch
from rdkit import Chem
from rdkit.Chem import AllChem, DataStructs

def smiles_to_ecfp(smiles, radius=2, n_bits=1024):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return np.zeros(n_bits)
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)
    arr = np.zeros(n_bits, dtype=int)
    DataStructs.ConvertToNumpyArray(fp, arr)
    return arr

class ModelWrapper:

    def __init__(self, model_name: str = None):
        self.model = None
    
        model_name = os.path.join(os.environ.get("MODELS_DIR"), model_name)

        print(model_name)

        if model_name and os.path.exists(model_name):
            try: 
                self.model = torch.load(model_name, map_location="cpu", weights_only=False)
            except Exception as e:
                print(e)
                self.model = None
        print(self.model)
        self.featurizer = smiles_to_ecfp

    def predict(self, X):

        X = self.featurizer(X)
        
        X = np.asarray(X, dtype=float)

        

        # self.model.eval()
        with torch.no_grad():
            t = torch.tensor(X, dtype=torch.float32)
            out = self.model(t)
            # print(out.cpu().numpy().item())
            return out.cpu().numpy().item()