# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_features.ipynb (unless otherwise specified).

__all__ = ['Extractor', 'SIFTExtractor', 'CNNExtractor', 'imagenet_normalize_transform', 'get_transforms',
           'SimCLRExtractor', 'gen_vcodebook', 'gen_codebooks', 'get_df', 'get_bovw', 'extract_features',
           'new_get_bovw', 'calc_tf_idf']

# Cell
import cv2
import pickle
# import ffmpeg
import random
import time
import torch

import numpy as np

from abc import ABC, abstractmethod
from pathlib import Path
from PIL import Image
from .prep import *
from tqdm.auto import tqdm
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity

# Cell
class Extractor(ABC):

    def __init__(self, extractor):
        self.extractor = extractor
        super().__init__()

    @abstractmethod
    def extract(self, img):
        pass

# Cell
class SIFTExtractor(Extractor):

    '''Exposed SIFTExtractor class used for retrieving features.'''

    def extract(self, img):
        '''Given an image, extract features using SIFT. Returns the feature vector.'''
        img = np.array(img)
        _, features = self.extractor.detectAndCompute(img, None)
        return features

class CNNExtractor(Extractor):

    '''Exposed CNNExtractor class used for retrieving features.'''

    def extract(self, img):
        '''Given an image, extract features from the layers of a CNN. Returns the feature vector.'''

        return self.extractor.getFeatures(img)

# Cell
from torchvision import transforms

def imagenet_normalize_transform():
    return transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

def get_transforms(size=224):
    return transforms.Compose([
        transforms.Resize(size=(size, size)),
        transforms.ToTensor(),
        imagenet_normalize_transform()
    ])

class SimCLRExtractor(Extractor):

    '''Exposed CNNExtractor class used for retrieving features.'''
    def __init__(self, extractor):
        super().__init__(extractor)
        self.transforms = get_transforms()

    def extract(self, img):
        '''Given an image, extract features from the layers of a CNN. Returns the feature vector.'''
        img = self.transforms(img).float()
        img = img.unsqueeze(0)
        return self.extractor(img).detach().numpy()

# Cell
def gen_vcodebook(path, img_paths, model_name, extractor, vwords):
    """
        Constructs a visual codebook based on the given images.
        You can change vwords to increase the vocabulary of the codebook.
    """
    fname = path/f'models/features/{model_name}/{len(img_paths)}n_features.pkl'
    features_list = []
    feature_time = 0
    if not fname.is_file():
        feature_start = time.time()
        for img in tqdm(img_paths):
            features = extractor.extract(Image.open(img))
            if features is None: continue
            features_list.extend(features)

        features_list = np.asarray(features_list)
        pickle.dump(features_list, open(fname, 'wb'), protocol=4)
        feature_end = time.time()
        feature_time = feature_end - feature_start
        fname = path/f'models/features/{model_name}/{len(img_paths)}n_features_elapsed_time.txt'
        with open(fname, 'w') as f:
            f.write(f'{feature_time}')
    else:
        features_list = pickle.load(open(fname, 'rb'))
        with open(fname, 'r') as f:
            feature_time = [float(x) for x in f][0]
    codebooks = []
    for vw in vwords:
        cb_start = time.time()
        codebook = KMeans(n_clusters = vw)
        codebook.fit(features_list)
        cb_end = time.time()
        codebooks.append((cb_end - cb_start + feature_time, codebook))

    return codebooks

# Cell
def gen_codebooks(path, models, vwords, samples = 15_000):
    rico_path = path/'rico-images/data'
    img_paths = random.sample(list(rico_path.glob('*.jpg')), samples)
    for model_name, model in models:
        codebooks = gen_vcodebook(path, img_paths, model_name, model, vwords)
        for (cb_time, codebook), vw in zip(codebooks, vwords):
            fname = path/f'models/codebooks/{model_name}/cookbook_{model_name}_{vw}vw.model'
            pickle.dump(codebook, open(fname, 'wb'))
            fname = path/f'models/codebooks/{model_name}/cookbook_{model_name}_{vw}vw_elapsed_time.txt'
            with open(fname, 'w') as f:
                f.write(f'{cb_time}')

# Cell
def get_df(imgs, extractor, codebook, vwords):
    """Generates the document frequency for the visual words"""
    arr = []
    for img in imgs:
        features = extractor.extract(img)
        vw = codebook.predict(features)
        arr.extend(vw)
    arr = np.asarray(arr)

    return np.histogram(arr, bins = range(vwords + 1))

# Cell
def get_bovw(vid_path, extractor, codebook, vwords, n = None):
    """Generates the bag of visual words (bovw) for an entire video."""
    vid = cv2.VideoCapture(str(vid_path))
    if n is None: n = vid.get(cv2.CAP_PROP_FRAME_COUNT)

    # checks whether frames were extracted
    success = 1
    bovw = np.array([])
    for i in progress_bar(range(n)):
        # vid object calls read
        # function extract frames
        success, img = vid.read()
        if success:
            features = extractor.extract(img)
            vw = codebook.predict(features)
            bovw = np.concatenate((bovw, vw))

    hist = np.histogram(bovw, bins = range(vwords + 1))[0]
    return hist, bovw

# Cell
def extract_features(vid, extractor, fps = 30, frames_to_keep = 1):
    extracted_features = []
    for i in range(0, len(vid), int(fps / frames_to_keep)):
        img = vid[i]
        if not img: continue
        extracted_features.append(extractor.extract(img))

    return extracted_features

# Cell
def new_get_bovw(features, codebook, vwords):
    bovw = []
    for f in features:
        vw = codebook.predict(f)
        bovw.extend(vw)

    bovw = np.array(bovw)
    bovw = np.histogram(bovw, bins = range(vwords + 1))[0]
    return bovw

# Cell
def calc_tf_idf(tfs, dfs):
    tf_idf = np.array([])
    for tf, df in zip(tfs, dfs):
        tf = tf / np.sum(tfs)
        idf = np.log(len(tfs) / (df + 1))
        tf_idf = np.append(tf_idf, tf * idf)

    return tf_idf