import torch
from torch import optim, nn
from torchvision import models, transforms
from torchvision.models.vgg import VGG16_Weights


class FeatureExtractor(nn.Module):
  def __init__(self, model):
    super(FeatureExtractor, self).__init__()
		# Extract VGG-16 Feature Layers
    self.features = list(model.features)
    self.features = nn.Sequential(*self.features)
		# Extract VGG-16 Average Pooling Layer
    self.pooling = model.avgpool
		# Convert the image into one-dimensional vector
    self.flatten = nn.Flatten()
		# Extract the first part of fully-connected layer from VGG16
    self.fc = model.classifier[0]
  
  def forward(self, x):
		# It will take the input 'x' until it returns the feature vector called 'out'
    out = self.features(x)
    out = self.pooling(out)
    out = self.flatten(out)
    out = self.fc(out) 
    return out 

# Initialize the model
model = models.vgg16(weights=VGG16_Weights.DEFAULT)
new_model = FeatureExtractor(model)

# Change the device to GPU
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
new_model = new_model.to(device)

IMG_RESIZE_SIZE = 224
IMG_PATH = "data"

import cv2
from tqdm import tqdm
import numpy as np

# Transform the image, so it becomes readable with the model
transform = transforms.Compose([
  transforms.ToPILImage(),
  transforms.Resize((IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)),
  transforms.ToTensor()                              
])

# Will contain the feature
features = []

mappings = {}

import glob
# files = glob.glob("/content/drive/Shareddrives/ndl/kao/dataset 3/*.jpg")
files = glob.glob(f"{IMG_PATH}/*.jpg")

files.sort()

for index in tqdm(range(len(files))):

  path = files[index]

  img = cv2.imread(path)
  # Transform the image
  img = transform(img)
  # Reshape the image. PyTorch model reads 4-dimensional tensor
  # [batch_size, channels, width, height]
  # img = img.reshape(1, 3, 448, 448)
  img = img.reshape(1, 3, IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)
  img = img.to(device)
  # We only extract features, so we don't need gradient
  with torch.no_grad():
    # Extract the feature from the image
    feature = new_model(img)
  # Convert to NumPy Array, Reshape it, and save it to features variable
  features.append(feature.cpu().detach().numpy().reshape(-1))

  mappings[index] = {
      "nconst": path.split("/")[-1].split(".")[0],
      "name": "",
      "url": ""
  }

# Convert to NumPy Array
features = np.array(features)

import json

with open('mappings.json', mode='wt', encoding='utf-8') as file:
  json.dump(mappings, file, ensure_ascii=False, indent=2)

N_TREES = 1000

from annoy import AnnoyIndex

annoy_index = AnnoyIndex(features.shape[1], metric='euclidean')

for i in range(len(features)):

    feature = features[i]

    annoy_index.add_item(i, feature)

# k-d tree をビルドする
annoy_index.build(n_trees=N_TREES)

annoy_index.save("../models/index.ann")