from fastapi import FastAPI
from pydantic import BaseModel

from model import *

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

model = Model.load_from_checkpoint("methylsight2.ckpt")
model.eval()

class Protein(BaseModel):
    sequence: str

app = FastAPI()

@app.post("/predict")
def predict(protein: Protein):
    scores = []
    print("Embedding the sequences on the CPU...")
    site_embeddings = embed_sequences([("sequence", protein.sequence)], gpu=False)
    print("Making predictions on the CPU...")

    for protein, sites in tqdm.tqdm(site_embeddings.items()):
        try:
            for position, representation in sites.items():
                emb = representation
                with torch.no_grad():
                    logit = model(emb, 'methylation').squeeze(-2)[0]
                    scores.append({
                        'protein': protein,
                        'position': position,
                        #'logit': float(logit),
                        #'uncorrected_score': float(sigmoid(logit)),
                        'score': float(sigmoid(logit - MAX_LOGIT))
                    })
        except Exception as e:
            print(e)
            print(f"Could not do {protein}... skipping.")
            continue

    return {"results": scores }