# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3.6 (fn_env)
#     language: python
#     name: fn_env
# ---

# %load_ext line_profiler
# %load_ext autoreload
# %autoreload 2

import sys
import gym
import pathlib
import typing as tp
from scipy.special import softmax
import pandas as pd

sys.path.insert(0, '/mnt/ubuntu_data_hdd/school/masters/uwaterloo/1b/cs885/project/codenames/codenames_ai/src')

from default_game import *
from codenames import *
from codenames_env import *

import scann

glove.vectors.shape


# +
# searcher = scann.scann_ops_pybind.builder(glove.vectors, 20, "dot_product").tree(
#     num_leaves=2000, num_leaves_to_search=300, training_sample_size=250000).score_ah(
#     2, anisotropic_quantization_threshold=0.2).reorder(100).build()
# -

def compute_recall(neighbors, true_neighbors):
    total = 0
    for gt_row, row in zip(true_neighbors, neighbors):
        total += np.intersect1d(gt_row, row).shape[0]
    return total / true_neighbors.size


glove.vectorize(["hot dog", "vanilla ice cream"]).shape

# %timeit guesser.generate_word_suggestions_mean_approx(["seal" , "antarctica", "meal"], 20)

# %timeit guesser.generate_word_suggestions_minimax_approx(["seal" , "antarctica", "meal"], 20)

# %timeit guesser.generate_word_suggestions_minimax(["seal" , "antarctica", "meal"], 20)

# %timeit guesser.generate_word_suggestions_mean(["seal" , "antarctica", "meal"], 20)

guesser.generate_word_suggestions_minimax(["seal" , "antarctica", "meal"], 20)

# + tags=[]
guesser.give_hint_candidates(["seal" , "antarctica"], strategy="approx_mean")
# -

codenames_hacked_env.generate_candidates([], 3)

codenames_hacked_env = CodenamesEnvHack(glove, wordlist)

codenames_hacked_env.start_new_game()
codenames_hacked_env.render()
type(codenames_hacked_env.current_observation())


def bench_me():
    action = codenames_hacked_env.action_space.sample()
    step = codenames_hacked_env.step(action)
    print(step[1:3])


bench_me()

qq = np.arange(12).reshape(3, 4)

em = np.arange(100, 244).reshape(12, 12)

em[qq]

glove.vectors[[np.array([0, 1])]] @ glove.vectors[np.array([1, 2, 3])].T

np.min(glove.vectors[[np.array([0, 1])]] @ glove.vectors[np.array([1, 2, 3])].T, axis=1)

# %lprun -f GloveGuesser.generate_word_suggestions_mean_approx bench_me()

# %lprun -f bench_me bench_me()

codenames_hacked_env.render()



StarcoderdataJupyterScriptsDedupFiltered