import numpy as np
import torch
from tqdm import tqdm
import clip
from glob import glob
import gradio as gr
import os
import torchvision
import pickle
from collections import Counter

from SimSearch import FaissCosineNeighbors

# HELPERS
to_np = lambda x: x.data.to('cpu').numpy()

# DOWNLOAD THE DATASET and Files

torchvision.datasets.utils.download_file_from_google_drive('1kB1vNdVaNS1OGZ3K8BspBUKkPACCsnrG', '.', 'GTAV-Videos.zip')
torchvision.datasets.utils.download_file_from_google_drive('1pgvIBTs_6h23wIU28EdqO5y2T1wUfOak', '.', 'GTAV-embedding-vit32.zip')

# EXTRACT
torchvision.datasets.utils.extract_archive(from_path='GTAV-embedding-vit32.zip', to_path='Embeddings/VIT32/', remove_finished=False)
torchvision.datasets.utils.extract_archive(from_path='GTAV-Videos.zip', to_path='Videos/', remove_finished=False)

# Initialize CLIP model
clip.available_models()

# # Searcher

class GamePhysicsSearcher:
  def __init__(self, CLIP_MODEL, GAME_NAME, EMBEDDING_PATH='./Embeddings/VIT32/'):
    self.CLIP_MODEL = CLIP_MODEL
    self.GAME_NAME = GAME_NAME
    self.simsearcher = FaissCosineNeighbors()

    self.all_embeddings = glob(f'{EMBEDDING_PATH}{self.GAME_NAME}/*.npy')
    
    self.filenames = [os.path.basename(x) for x in self.all_embeddings]
    self.file_to_class_id = {x:i for i, x in enumerate(self.filenames)}
    self.class_id_to_file = {i:x for i, x in enumerate(self.filenames)}
    self.build_index()
    
  def read_features(self, file_path):
    with open(file_path, 'rb') as f:
      video_features = pickle.load(f)
    return video_features
  
  def read_all_features(self):
    features = {}
    filenames_extended = []

    X_train = []
    y_train = []

    for i, vfile in enumerate(tqdm(self.all_embeddings)):
      vfeatures = to_np(self.read_features(vfile))
      features[vfile.split('/')[-1]] = vfeatures
      X_train.extend(vfeatures)
      y_train.extend([i]*vfeatures.shape[0])
      filenames_extended.extend(vfeatures.shape[0]*[vfile.split('/')[-1]])

    X_train = np.asarray(X_train)
    y_train = np.asarray(y_train)
    
    return X_train, y_train
  
  def build_index(self):
    X_train, y_train = self.read_all_features()
    self.simsearcher.fit(X_train, y_train)
    
  def text_to_vector(self, query):
    text_tokens = clip.tokenize(query).cuda()
    with torch.no_grad():
      text_features = self.CLIP_MODEL.encode_text(text_tokens).float()
    text_features /= text_features.norm(dim=-1, keepdim=True)
    return to_np(text_features)

  # Source: https://stackoverflow.com/a/480227
  def f7(self, seq):
    seen = set()
    seen_add = seen.add # This is for performance improvement, don't remove
    return [x for x in seq if not (x in seen or seen_add(x))]
  
  def search_top_k(self, q, k=5, pool_size=1000, search_mod='Majority'):
    q = self.text_to_vector(q)
    nearest_data_points = self.simsearcher.get_nearest_labels(q, pool_size)
    
    if search_mod == 'Majority':
      topKs = [x[0] for x in Counter(nearest_data_points[0]).most_common(k)]
    elif  search_mod == 'Top-K':
      topKs = list(self.f7(nearest_data_points[0]))[:k]
    
    video_filename = [f'./Videos/{self.GAME_NAME}/' + self.class_id_to_file[x].replace('npy', 'mp4') for x in topKs]
    
    return video_filename



################ SEARCH CORE ################
# CRAETE CLIP MODEL
vit_model, vit_preprocess = clip.load("ViT-B/32")
vit_model.cuda().eval()

saved_searchers = {}
def gradio_search(query, game_name, selected_model, aggregator, pool_size, k=6):
  # print(query, game_name, selected_model, aggregator, pool_size)
  if f'{game_name}_{selected_model}' in saved_searchers.keys():
    searcher = saved_searchers[f'{game_name}_{selected_model}']
  else:
    if selected_model == 'ViT-B/32':
      model = vit_model
      searcher = GamePhysicsSearcher(CLIP_MODEL=model, GAME_NAME=game_name)
    else:
      raise
    
    saved_searchers[f'{game_name}_{selected_model}'] = searcher
  
  results = []
  relevant_videos = searcher.search_top_k(query, k=k, pool_size=pool_size, search_mod=aggregator)
  params = ', '.join(map(str, [query, game_name, selected_model, aggregator, pool_size]))
  results.append(params)
  results.extend(relevant_videos)
  print(results)
  return results

list_of_games = ['Grand Theft Auto V']





# GRADIO APP
iface = gr.Interface(fn=gradio_search, 
                     inputs =[ gr.inputs.Textbox(lines=1, placeholder='Search Query', default="A man in the air", label=None),
                               gr.inputs.Radio(list_of_games, label="Game To Search"),
                               gr.inputs.Radio(['ViT-B/32'], label="MODEL"),
                               gr.inputs.Radio(['Majority', 'Top-K'], label="Aggregator"),
                                gr.inputs.Slider(300, 2000, label="Pool Size"),
                             ], 
                     outputs=[
                              gr.outputs.Textbox(type="auto", label='Search Params'),
                              gr.outputs.Video(type='mp4', label='Result 1'),
                              gr.outputs.Video(type='mp4', label='Result 2'),
                              gr.outputs.Video(type='mp4', label='Result 3'),
                              gr.outputs.Video(type='mp4', label='Result 4'),
                              gr.outputs.Video(type='mp4', label='Result 5')],
                     server_port=7878,
                     server_name="0.0.0.0",
                     # examples=[],
                     title='CLIP Meets Game Physics Demo'
                    )
iface.launch()