import time
import torch
import numpy as np
from inference import AlphaFoldEmbedInference, AlphaFoldEmbedTokenizer

def benchmark():
    model = AlphaFoldEmbedInference()
    tokenizer = AlphaFoldEmbedTokenizer()
    model.eval()
    
    # Dummy data
    seqs = ["".join(np.random.choice(list("ACDEFGHIKLMNPQRSTVWY"), size=150)) for _ in range(100)]
    
    # true plddt
    true_plddts = [np.random.uniform(50, 100, size=150) for _ in range(100)]
    
    start_time = time.time()
    maes = []
    with torch.no_grad():
        for seq, true_plddt in zip(seqs, true_plddts):
            tokens = tokenizer.tokenize(seq)
            input_ids = torch.tensor([tokens])
            out = model(input_ids)
            pred_plddt = out["plddt"].squeeze(0).numpy()
            mae = np.mean(np.abs(pred_plddt - true_plddt))
            maes.append(mae)
    end_time = time.time()
    
    throughput = len(seqs) / (end_time - start_time)
    avg_mae = np.mean(maes)
    
    print(f"pLDDT MAE: {avg_mae:.4f}")
    print(f"Sequence Processing Throughput: {throughput:.2f} seq/s")

if __name__ == "__main__":
    benchmark()