import numpy as np
import pandas as pd
import joblib
import argparse
import os
import sys
import torch
from rdkit import Chem
from rdkit.Chem import AllChem, MACCSkeys

BASE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, BASE)

from build_vocab import WordVocab
from pretrain_trfm import TrfmSeq2seq
from utils import split

cache_data   = np.load(os.path.join(BASE, "prott5_cache.npz"), allow_pickle=True)
prott5_cache = cache_data["cache"].item()
scaler       = joblib.load(os.path.join(BASE, "scaler.pkl"))
smiles_cache = joblib.load(os.path.join(BASE, "smiles_cache.pkl"))

N_FOLDS = 5
et_kcat_models = [joblib.load(os.path.join(BASE, f"et_kcat_fold{i}.pkl")) for i in range(N_FOLDS)]
et_km_models   = [joblib.load(os.path.join(BASE, f"et_km_fold{i}.pkl"))   for i in range(N_FOLDS)]

vocab = WordVocab.load_vocab(os.path.join(BASE, "vocab.pkl"))
pad_index = 0; unk_index = 1; eos_index = 2; sos_index = 3

trfm = TrfmSeq2seq(len(vocab), 256, len(vocab), 4)
trfm.load_state_dict(torch.load(os.path.join(BASE, "trfm_12_23000.pkl"), map_location="cpu"))
trfm.eval()

def smiles_to_vec(smiles):
    x_split = [split(smiles)]
    xid = []
    for tokens in x_split:
        seq_len = 220
        if len(tokens) > 218:
            tokens = tokens[:109] + tokens[-109:]
        ids = [vocab.stoi.get(t, unk_index) for t in tokens]
        ids = [sos_index] + ids + [eos_index]
        pad_len = seq_len - len(ids)
        ids += [pad_index] * pad_len
        xid.append(ids)
    xid = torch.tensor(xid).T
    with torch.no_grad():
        out = trfm.encode(xid)
    return np.array(out, dtype=np.float32).squeeze(0)

def get_maccs(smiles):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None: return np.zeros(167, dtype=np.float32)
    return np.array(MACCSkeys.GenMACCSKeys(mol), dtype=np.float32)

def get_morgan_fp(smiles, radius=2, nbits=2048):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None: return np.zeros(nbits, dtype=np.float32)
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=nbits)
    return np.array(fp, dtype=np.float32)

def predict(input_path, output_path):
    df = pd.read_parquet(input_path)
    rows = []
    for _, row in df.iterrows():
        prot_emb   = prott5_cache.get(row["sequence"], np.zeros(1024, dtype=np.float32))
        smiles_emb = smiles_cache.get(row["smiles"])
        if smiles_emb is None:
            smiles_emb = smiles_to_vec(row["smiles"])
        maccs  = get_maccs(row["smiles"])
        morgan = get_morgan_fp(row["smiles"])
        rows.append(np.concatenate([prot_emb, smiles_emb, maccs, morgan]))
    X = np.array(rows, dtype=np.float32)
    X = scaler.transform(X)
    pred_kcat = np.mean([m.predict(X) for m in et_kcat_models], axis=0)
    pred_km   = np.mean([m.predict(X) for m in et_km_models],   axis=0)
    out = pd.DataFrame({
        "id"         : df["id"].values,
        "log10_kcat" : pred_kcat,
        "log10_km"   : pred_km,
    })
    out.to_csv(output_path, index=False)
    print(f"Saved {len(out)} predictions to {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input",  required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()
    predict(args.input, args.output)