import os
from PIL import Image
import faiss
import streamlit as st
import torch
import torch.nn.functional as F
from torchvision.transforms import Compose, Resize, ToTensor, Normalize

from model import FashionPrediction
transforms = Compose([Resize((232, 232)), ToTensor(),
                          Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

path = './data/final-models/resnet_152_classification.pt'
index_path = os.path.join('./data/index_files/resnet152_unweighted_flat.index')

gender_dict = {0: 'Boys', 1: 'Girls', 2: 'Men', 3: 'Unisex', 4: 'Women'}
master_dict = {0: 'Accessories', 1: 'Apparel', 2: 'Footwear', 3: 'Personal Care'}
subcat_dict = {0: 'Accessories', 1: 'Apparel Set', 2: 'Bags', 3: 'Belts', 4: 'Bottomwear', 5: 'Cufflinks',
               6: 'Dress', 7: 'Eyes', 8: 'Eyewear', 9: 'Flip Flops', 10: 'Fragrance', 11: 'Headwear',
               12: 'Innerwear', 13: 'Jewellery', 14: 'Lips', 15: 'Loungewear and Nightwear', 16: 'Makeup',
               17: 'Mufflers', 18: 'Nails', 19: 'Sandal', 20: 'Saree', 21: 'Scarves', 22: 'Shoe Accessories',
               23: 'Shoes', 24: 'Skin', 25: 'Skin Care', 26: 'Socks', 27: 'Stoles', 28: 'Ties', 29: 'Topwear',
               30: 'Wallets', 31: 'Watches'}
color_dict = {0: 'Beige', 1: 'Black', 2: 'Blue', 3: 'Bronze', 4: 'Brown', 5: 'Burgundy', 6: 'Charcoal',
              7: 'Coffee Brown', 8: 'Copper', 9: 'Cream', 10: 'Gold', 11: 'Green', 12: 'Grey', 13: 'Grey Melange',
              14: 'Khaki', 15: 'Lavender', 16: 'Magenta', 17: 'Maroon', 18: 'Mauve', 19: 'Metallic', 20: 'Multi',
              21: 'Mushroom Brown', 22: 'Mustard', 23: 'Navy Blue', 24: 'Nude', 25: 'Off White', 26: 'Olive',
              27: 'Orange', 28: 'Peach', 29: 'Pink', 30: 'Purple', 31: 'Red', 32: 'Rose', 33: 'Rust',
              34: 'Sea Green', 35: 'Silver', 36: 'Skin', 37: 'Steel', 38: 'Tan', 39: 'Taupe', 40: 'Teal',
              41: 'Turquoise Blue', 42: 'White', 43: 'Yellow'}


@st.experimental_memo
def load_trained_model(model_path):
    model = FashionPrediction()
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    model.eval()
    return model


@st.experimental_memo
def load_index():
    index = faiss.read_index(index_path)
    return index


def get_nearest_k(input_image, top_k=5):
    model = load_trained_model(path)
    transformed_image = transforms(input_image)
    transformed_image = torch.unsqueeze(transformed_image, 0)
    with torch.inference_mode():
        query_embeddings = model(transformed_image, True)
    the_index = load_index()
    dist, indexes = the_index.search(query_embeddings, top_k)
    return dist, indexes


def get_predictions(input_image):
    model = load_trained_model(path)
    transformed_image = transforms(input_image)
    transformed_image = torch.unsqueeze(transformed_image, 0)
    with torch.inference_mode():
        logits = model(transformed_image, False)
    gender_prob = F.softmax(logits[0], dim=1)
    master_prob = F.softmax(logits[1], dim=1)
    subcat_prob = F.softmax(logits[2], dim=1)
    color_prob = F.softmax(logits[3], dim=1)
    top2_gender = torch.topk(gender_prob, 2, dim=1)
    top2_master = torch.topk(master_prob, 2, dim=1)
    top2_subcat = torch.topk(subcat_prob, 2, dim=1)
    top2_color = torch.topk(color_prob, 2, dim=1)

    all_predictions = (list(top2_gender.values.numpy().reshape(-1)), list(top2_gender.indices.numpy().reshape(-1))), \
        (list(top2_master.values.numpy().reshape(-1)), list(top2_master.indices.numpy().reshape(-1))), \
        (list(top2_color.values.numpy().reshape(-1)), list(top2_color.indices.numpy().reshape(-1))), \
        (list(top2_subcat.values.numpy().reshape(-1)), list(top2_subcat.indices.numpy().reshape(-1)))

    pred_dict = {
        'Predicted Master Category': (master_dict[all_predictions[1][1][0]], master_dict[all_predictions[1][1][1]]),
        'Master Category Probability': [round(prob, 3) for prob in all_predictions[1][0]],
        'Predicted Sub Category': (subcat_dict[all_predictions[3][1][0]], subcat_dict[all_predictions[3][1][1]]),
        'Sub Category Probability': [round(prob, 3) for prob in all_predictions[3][0]],
        'Predicted person type': (gender_dict[all_predictions[0][1][0]], gender_dict[all_predictions[0][1][1]]),
        'Person Type Probability': [round(prob, 3) for prob in all_predictions[0][0]],
        'Predicted Color': (color_dict[all_predictions[2][1][0]], color_dict[all_predictions[2][1][1]]),
        'Color Probability': [round(prob, 3) for prob in all_predictions[2][0]]
    }

    return pred_dict


if __name__ == '__main__':
    sample_image = Image.open('./data/small_images_0_9999/0.jpg')
    output = get_predictions(sample_image)
    print(output)