import os, pickle, h5py
import matplotlib.pyplot as plt
import numpy as np
from nsd_visuo_semantics.get_embeddings.word_lists import coco_categories_91
from nsd_visuo_semantics.get_embeddings.nsd_embeddings_utils import sentence_embeddings_sanity_check, get_words_from_multihot
from nsd_visuo_semantics.get_embeddings.embedding_models_zoo import get_embedding_model, get_embeddings


def get_nsd_sentence_embeddings_categories_simple(embedding_model_type, captions_to_embed_path, 
                                                  h5_dataset_path, SAVE_PATH, OVERWRITE):
    '''
    Concatenates the coco categories into a string, and throws that into a sentence embedder.
    There is the option to only keep the coco categories that are also present/absent in the captions.
    embedding_model_type: str, the model to use. See embedding_models_zoo.py for options.
    captions_to_embed_path: str, path to the pickle file containing the captions of nsd.
    h5_dataset_path: str, path to the h5 dataset containing the images and categories of ms-coco/nsd.
    OVERWRITE: bool, if True, overwrite existing embeddings.
    '''

    print(f"GATHERING CATEGORY EMBEDDINGS FOR: {embedding_model_type}\n "
          f"ON: {captions_to_embed_path}") 

    SANITY_CHECK = 1
    GET_EMBEDDINGS = 1
    FINAL_CHECK = 1

    save_embeddings_to = SAVE_PATH
    save_test_imgs_to = f"{save_embeddings_to}/_check_imgs"
    os.makedirs(save_test_imgs_to, exist_ok=1)

    safety_check_metric = 'correlation'

    save_name = f"nsd_{embedding_model_type}_CATEGORY_concatString_mean_embeddings"

    if os.path.exists(f"{save_embeddings_to}/{save_name}_allCats.pkl") and not OVERWRITE:
        print(f"Embeddings already exist at {save_embeddings_to}/{save_name}.pkl. Set OVERWRITE=True to overwrite.")
    else:

        embedding_model = get_embedding_model(embedding_model_type)

        if SANITY_CHECK:
            sentence_embeddings_sanity_check(embedding_model_type, embedding_model, safety_check_metric, save_test_imgs_to)

        if GET_EMBEDDINGS:

            with open(captions_to_embed_path, "rb") as fp:
                loaded_captions = pickle.load(fp)

            with h5py.File(h5_dataset_path,'r') as f:
                loaded_multihot_labels = f['test']['img_multi_hot'][:]

            n_elements = len(loaded_captions)
            dummy_embeddings = get_embeddings(loaded_captions[0], embedding_model, embedding_model_type)

            mean_embeddings_all = np.empty((n_elements, dummy_embeddings.shape[-1]))
            
            cats_per_image = []

            for i in range(n_elements):
                if i % 1000 == 0:
                    print(f"\rRunning... {i/n_elements*100:.2f}%", end="")

                these_captions = loaded_captions[i]

                if not isinstance(these_captions, list):
                    # needed if we are using a single caption per image
                    # in that case, we have a string and convert it to a list
                    # with a single element
                    these_captions = [these_captions]

                img_category_words = get_words_from_multihot(loaded_multihot_labels[i], coco_categories_91)

                # first, we keep all categories
                if len(img_category_words) == 0:
                    all_cat_word_string = "something"
                else:
                    all_cat_word_string = " ".join(img_category_words)
                these_all_cat_embeds = get_embeddings(all_cat_word_string, embedding_model, embedding_model_type)
                mean_embeddings_all[i] = these_all_cat_embeds
                cats_per_image.append(all_cat_word_string)

            with open(f"{save_embeddings_to}/{save_name}_allCats.pkl", "wb") as fp:
                pickle.dump(mean_embeddings_all, fp)

            with open(f"{save_embeddings_to}/{save_name}_categs_per_image.pkl", "wb") as fp:
                pickle.dump(cats_per_image, fp)

        del mean_embeddings_all, cats_per_image
            

    if FINAL_CHECK:
        with h5py.File(h5_dataset_path, "r") as h5_dataset:
            total_n_stims = h5_dataset["test"]["labels"][:].shape[0]
            plot_n_imgs = 10
            step_size = total_n_stims // plot_n_imgs

            with open(captions_to_embed_path, "rb") as fp:
                loaded_captions = pickle.load(fp)
            with open(f"{save_embeddings_to}/{save_name}_allCats.pkl", "rb") as fp:
                loaded_mean_embeddings = pickle.load(fp)
            with open(f"{save_embeddings_to}/{save_name}_categs_per_image.pkl", "rb") as fp:
                loaded_cats_per_image = pickle.load(fp)

            for i in range(0, total_n_stims, step_size):
                plt.imshow(h5_dataset["test"]["data"][i])
                plt.title(
                    f"{loaded_captions[i][0]}\n"
                    f"{loaded_cats_per_image[i]}\n"
                    f"Emb shape, min, max, mean: {loaded_mean_embeddings[i].shape, loaded_mean_embeddings[i].min(), loaded_mean_embeddings[i].max(), loaded_mean_embeddings[i].mean()}"
                )
                plt.savefig(f"{save_test_imgs_to}/{save_name}_check_{i}.png")
                plt.close()
