import time

import tensorflow as tf

PROJECT_ID = "ksalama-cloudml"
BUCKET = "ksalama-cloudml"
INDEX_DIR = f"gs://{BUCKET}/bqml/scann_index"
BQML_MODEL_DIR = f"gs://{BUCKET}/bqml/item_matching_model"
LOOKUP_MODEL_DIR = f"gs://{BUCKET}/bqml/embedding_lookup_model"

songs = {
    "2114406": "Metallica: Nothing Else Matters",
    "2114402": "Metallica: The Unforgiven",
    "2120788": "Limp Bizkit: My Way",
    "2120786": "Limp Bizkit: My Generation",
    "1086322": "Jacques Brel: Ne Me Quitte Pas",
    "3129954": "Édith Piaf: Non, Je Ne Regrette Rien",
    "53448": "France Gall: Ella, Elle l'a",
    "887688": "Enrique Iglesias: Tired Of Being Sorry",
    "562487": "Shakira: Hips Don't Lie",
    "833391": "Ricky Martin: Livin' la Vida Loca",
    "1098069": "Snoop Dogg: Drop It Like It's Hot",
    "910683": "2Pac: California Love",
    "1579481": "Dr. Dre: The Next Episode",
    "2675403": "Eminem: Lose Yourself",
    "2954929": "Black Sabbath: Iron Man",
    "625169": "Black Sabbath: Paranoid",
}

"""
Explanation: Test the Retrieval Latency of Approximate vs Exact Matching
End of explanation
"""


class ExactMatcher(object):
    def __init__(self, model_dir):
        print("Loading exact matchg model...")
        self.model = tf.saved_model.load(model_dir)
        print("Exact matchg model is loaded.")

    def match(self, instances):
        outputs = self.model.signatures["serving_default"](
            tf.constant(instances, tf.dtypes.int64)
        )
        return outputs["predicted_item2_Id"].numpy()

exact_matcher = ExactMatcher(BQML_MODEL_DIR)

exact_matches = {}

start_time = time.time()
for i in range(100):
    for song in songs:
        matches = exact_matcher.match([int(song)])
        exact_matches[song] = matches.tolist()[0]
end_time = time.time()
exact_elapsed_time = end_time - start_time

print(
    f"Elapsed time: {round(exact_elapsed_time, 3)} seconds - average time: {exact_elapsed_time / (100 * len(songs))} seconds"
)

"""
Explanation: Exact Matching
End of explanation
"""


from index_server.matching import ScaNNMatcher

scann_matcher = ScaNNMatcher(INDEX_DIR)
embedding_lookup = tf.saved_model.load(LOOKUP_MODEL_DIR)

approx_matches = dict()

start_time = time.time()
for i in range(100):
    for song in songs:
        vector = embedding_lookup([song]).numpy()[0]
        matches = scann_matcher.match(vector, 50)
        approx_matches[song] = matches
end_time = time.time()
scann_elapsed_time = end_time - start_time

print(
    f"Elapsed time: {round(scann_elapsed_time, 3)} seconds - average time: {scann_elapsed_time / (100 * len(songs))} seconds"
)

speedup_percent = round(exact_elapsed_time / scann_elapsed_time, 1)
print(f"ScaNN speedup: {speedup_percent}x")

"""
Explanation: Approximate Matching (ScaNN)
End of explanation
"""


