from IPython.display import JSON  # Helpful UI for JSON display
import time
from biolmai import BioLM
import requests  # Will use to make calls to BioLM.ai
import csv  # To read example data
import numpy as np
from numpy.linalg import norm

lines = []
with open('data/protein/data/PLA2.csv', newline='') as csvfile:
    reader = csv.reader(csvfile)
    for row in reader:
        lines.append(row)
print(lines)

SEQ1 = lines[1][1]
SEQ2 = lines[2][1]

print("Sequence length 1: {}".format(len(SEQ1)))
print("Sequence length 2: {}".format(len(SEQ2)))

# Make the request - let's time it!
start = time.time()
result = BioLM(entity="biolmtox2", action="encode", type="sequence", items=[SEQ1, SEQ2])
end = time.time()
print(f"BioLMTox prediction took {end - start:.4f} seconds.")

# If you wish to view the full result, you can expand the tree in the cell below
JSON(result)

# Define similarity measure
def cos_similarity(a, b):
    return np.dot(a,b)/(norm(a)*norm(b))

# convert sequence embeddings to numpy arrays
em_1 = np.asarray(result[0]["mean_representation"])
em_2 = np.asarray(result[1]["mean_representation"])

# compute similarity measures
em_similarity = cos_similarity(em_1, em_2)
print(f'sequence embedding cosine similarity:\n{em_similarity}')