import gradio as gr
import torch
from PIL import Image
import torch.nn.functional as F
import numpy as np
import pickle
import json
import requests
from transformers import CLIPProcessor, AutoModelForSemanticSegmentation, AutoFeatureExtractor, CLIPModel
from torch import nn
import io

# Initialize the models using huggingface

device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the CLIP model from hugging face
clip_hg = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device).eval()
processor_hg = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
seg_hg = AutoModelForSemanticSegmentation.from_pretrained('mattmdjaga/segformer_b2_clothes').to(device).eval()
extractor_hg = AutoFeatureExtractor.from_pretrained('mattmdjaga/segformer_b2_clothes', reduce_labels=False)

# Load the data and normalize the embeddings just in case. 
features = torch.load('features.pt').to(device)
features_main = F.normalize(features)
item_embeddings = torch.load('item_embeds.pt').to(device)
item_embeddings = F.normalize(item_embeddings)
url_list_main = pickle.load(open('new_url_list.pt','rb'))
clothes_tree = json.load(open('clothes_tree_new_data.json'))
rec_dic = json.load(open('top5_mini_new.json'))

# URL for an image if no image is selected
url = 'https://bitsofco.de/content/images/2018/12/Screenshot-2018-12-16-at-21.06.29.png'

# Set up all the variables

label =  ['Background', 'Hat', 'Hair', 'Sunglasses', 'Upper-clothes', 'Skirt', 'Pants', 'Dress', 'Belt',
                  'Left-shoe', 'Right-shoe', 'Face', 'Left-leg', 'Right-leg', 'Left-arm', 'Right-arm', 'Bag', 'Scarf']


clothing_type = ['top', 'bottom', 'dress']
top_type = ['t-shirt', 'tank top', 'blouse', 'sweater', 'hoodie', 'cardigan','turtleneck','blazer','polo','collar shirt','knitwear',
    'tuxedo', 'Compression top','duffle coat', 'peacoat', 'long coat', 'trench coat',
    'biker jacket', 'blazer', 'bomber jacket', 'hooded jacket', 'leather jacket', 'military jacket', 'down jacket', 'shirt jacket',
    'suit jacket', 'dinner jacket', 'gillet', 'track jacket'
    ]
bottom_type = ['skirt', 'leggings', 'sweatpants', 'skinny pants', 'tailored pants', 'track pants', 'wide-leg pants'
    , 'cargo shorts', 'denim shorts', 'track shorts', 'compression shorts', 'cycling shorts','denim pants',
    'cargo pants', 'chino pants', 'chino shorts'
    ]
dress_type = ['casual dress', 'cocktail dress', 'evening dress', 'maxi dress', 'mini dress', 'party dress', 'sundress']
styles = ['plain','polka dot','striped','floral','checkered','zebra print','leopard print','plaid','paisley']
colors = ['blue','red','pink','orange','yellow','purple','gold','white','off white','black','grey','green','brown','beige','cream','navy','maroon']

top_list = [f"{t}, {color}, {style}" for t in top_type for style in styles for color in colors]
bottom_list =  [f"{t}, {color}, {style}" for t in bottom_type for style in styles for color in colors]
dress_list =  [f"{t}, {color}, {style}" for t in dress_type for style in styles for color in colors]
all_items = top_list + bottom_list + dress_list


clothing_type = ['top', 'bottom', 'dress']
all_types = {'top' :top_type,
    'bottom' : bottom_type,
    'dress':dress_type}
patterns_list = styles.copy()
colors_list = colors.copy()

clicks = 0
c_types = []
types = []
colors = []
patterns = []
new_files = []
out = []


clothes_click = 0
global_mask = None
mask_choice = 'Clothes'

# Define all needed functions

def find_closest(target_feature, features):
    '''
    Purpose: Find the closest embedding to the given image embedding
    Inputs: 
        target_feature (tenosr): embedding of our search item
        features (tensor): embedding of all the items in the dataset
    Outputs:
    group_sorted_indices (list): indicies of the closest items in a sorted order
    '''
    cos_similarity = features.to(torch.float32) @ target_feature.to(torch.float32).T
    group_sorted_indices = torch.argsort(cos_similarity, descending=True,dim=0).squeeze(1).cpu().tolist()
    return group_sorted_indices


def filter_function(choices):
    '''
    Purpose: Find a list of items that match the given filters
    Inputs:
        choices (list): list of filters
    Outputs:
        Upating the choices of filters
    '''
    # Import the global variables
    global clicks
    global c_types
    global types
    global colors
    global patterns
    global new_files
    new_choices = []

    # Clicks is just a reference to how far we are in the filter tree
    # We keep going down and saving the selected options until we reach the end
    # Then we add items which had the desired filters to the new_choices list
    # This is then used to filter out the items that don't match the filters in search

    if clicks == 0:
        temp_choices = [choice for choice in choices if choice in clothing_type]
        if len(temp_choices) == 0:
            temp_choices = clothing_type
        for choice in temp_choices:
            c_types.append(choice)
            new_choices.extend(list(clothes_tree[choice].keys()))

    if clicks == 1:
        temp_choices = [choice for c_type in c_types for choice in choices if choice in all_types[c_type]]
        if len(temp_choices) == 0:
            types = []
            for c_type in c_types:
                types.extend([(t,c_type) for t in clothes_tree[c_type].keys()])
        for choice in temp_choices:
            if choice in clothes_tree['top']:
                types.append((choice,'top'))
            elif choice in clothes_tree['bottom']:
                types.append((choice,'bottom'))
            else :
                types.append((choice,'dress'))
        new_choices = list(clothes_tree['top']['t-shirt'].keys())

    if clicks == 2:
        temp_choices = [choice for choice in choices if choice in colors_list]
        if len(temp_choices) == 0:
            colors = colors_list.copy()
        for choice in temp_choices:
            colors.append(choice)
        new_choices = list(clothes_tree['top']['t-shirt']['red'].keys())

    if clicks == 3:
        temp_choices = [choice for choice in choices if choice in patterns_list]
        if len(temp_choices) == 0:
            patterns = patterns_list.copy()
        for choice in temp_choices:
            patterns.append(choice)
        for type_,c_type in types:
            for color in colors:
                for pattern in patterns:
                    new_files.extend(clothes_tree[c_type][type_][color][pattern])
        clicks += 1
        new_choices = ['Press Search to use the set filter. Dont press this button']
        return gr.update(choices=new_choices, label='Press Search to use the filter or press filter to reset the filter')
    if clicks == 4:
        c_types.clear()
        types.clear()
        colors.clear()
        patterns.clear()
        new_files.clear()
        clicks = 0
        new_choices = ['top','bottom','dress']

        return gr.update(choices=new_choices,label='Select the type of clothing you want to search for')
    clicks += 1
    return gr.update(choices=new_choices)

def set_theme(theme):
    '''
    Purpose: Set the theme using filters
    Inputs:
        theme (string): theme to be set
    Outputs:
        Upadting to show the chosen theme
    '''
    global new_files
    new_files.clear()

    # Here we just manually set the filters to the desired theme
    # Then we just find images with the desired filters

    if theme == 'Red carpet':
        types = [('evening dress','dress'), ('tuxedo','top'), ('suit jacket','top'), ('dinner jacket','top'),('maxi dress','dress')]
        colors = ['red','purple','gold','white','off white','black','beige','cream','navy','maroon']
        patterns = ['plain']
    elif theme == 'Sports':
        types = [('running shorts','bottom'), ('track shorts','bottom'), ('track pants','bottom'), ('track jacket','top'),
            ('Compression top','top'), ('cycling top','top'), ('cycling shorts','bottom'),('compression shorts','bottom'),('tank top','top')]
        colors = colors_list.copy()
        patterns = patterns_list.copy()#
    elif theme =='My preference':
        types = [('evening dress','dress'), ('tuxedo','top'), ('suit jacket','top'), ('dinner jacket','top'),('maxi dress','dress')]
        colors = ['red','purple','gold']
        patterns = ['plain','zebra print']
    else:
        return gr.update(label='Chosen theme: None')
    for type_,c_type in types:
        for color in colors:
            for pattern in patterns:
                new_files.extend(clothes_tree[c_type][type_][color][pattern])
    return gr.update(label='Chosen theme: '+theme)
    

def segment(img):
    '''
    Purpose: Segment the image to get the mask
    Inputs:
        img(pil image): image to be segmented
    Outputs:
        img(pil image): original image
        arr(numpy array): array of image
        pred_seg(tensor): mask
    '''

    # Get the segmentation mask then umsample it to the original size

    encoding = extractor_hg(img.convert('RGB'), return_tensors="pt")
    pixel_values = encoding.pixel_values.to(device)
    outputs = seg_hg(pixel_values=pixel_values)
    logits = outputs.logits.cpu()
    upsampled_logits = nn.functional.interpolate(
        logits,
        size=img.size[::-1],
        mode="bilinear",
        align_corners=False,
    )

    pred_seg = upsampled_logits.argmax(dim=1)[0]
    arr_img = np.array(img)
    return img, arr_img, pred_seg

def clean_img(img):
    '''
    Purpose: Clean the image to remove the chosen items
    Inputs:
        img(numpy array): image to be cleaned
    Outputs:
        img(numpy array): cleaned image
    '''

    # Here we remove pixels whihc are not in our desired class

    global global_mask
    global mask_choice
    bad = []
    mask_size = global_mask.shape
    img_size = img.shape[:2]
    if img_size != mask_size:
        return img
    if mask_choice=='Person':
        bad.append(0)
    elif mask_choice=='Clothes':
        bad.extend([0,2,15,14,13,12,11])
    elif mask_choice=='Upper Body/Dress':
        bad.extend([0,5,6,9,10,12,13,16])
    elif mask_choice=='Lower Body':
        bad.extend([0,1,2,3,4,7,8,11,14,15,16])
    elif mask_choice=='Upper Body/Dress, no person':
        bad.extend([0,1,2,15,11,14,5,6,9,10,12,13,16,3])
    for i in bad:
        global_mask[global_mask==i] = 50
    img[global_mask==50] = 255
    return img


def label_to_rec_lables (label):
    '''
    Purpose: Use the label to get the corresponding reccomendation labels
    Inputs:
        label(string): label of the image
    Outputs:
        rec_labels(list): list of reccomendation labels
    '''

    # This function is used to get the reccomendation labels to then
    # filter the reccomendation search to them

    labels = label.split(',')
    new_label = rec_dic[','.join(labels[:2])]
    print('Reccomendation label: ',new_label)
    n = 5 if len(new_label) >= 5 else len(new_label)
    labels = []
    labels = [new_label[i][0].split(',') for i in range(n)]
    chosen = []
    c_types = ['top','bottom','dress']
    for item in labels:

        label_type = item[0]
        label_color = item[1].strip()
        for c_type in c_types:
            if label_type in all_types[c_type]:
                item_type = c_type
                chosen.append([item_type,label_type,label_color])
    print('Chosen: ',chosen)
    return chosen


def filter_features(labels, rec=False, rec_items=None):
    '''
    Purpose: Filter the features to only contain the chosen label
    Inputs:
        labels(str): label string
        rec(bool): if the function is called from the recommendation function
        rec_items(list): list containing the label info
    Outputs:
        url_list(list): list of urls after filtering
        features(tensor): features after filtering
        '''
    global url_list_main
    global features_main

    # Here we filter the features to only contain the desired labels and 
    # also provide the new url list

    labels = labels.split(',')
    label_type = labels[0]
    label_color = labels[1].strip()
    c_types = ['top', 'bottom', 'dress']
    for c_type in c_types:
        if label_type in all_types[c_type]:
            item_type = c_type
    new_list = set()
    if rec:
        item_type = rec_items[0]
        label_type = rec_items[1]
        label_color = rec_items[2]
        for pattern in patterns_list:
            new_list.update(clothes_tree[item_type][label_type][label_color][pattern])
    else:
        for color in colors_list:
            for pattern in patterns_list:
                new_list.update(clothes_tree[item_type][label_type][color][pattern])
    new_files = list(new_list)
    temp_url = []
    temp_features = torch.zeros(len(new_files), 512).to(device)
    for c,i in enumerate(new_files):
        temp_url.append(url_list_main[i])
        temp_features[c] = features_main[i]
    url_list = temp_url
    features = temp_features.to(torch.float32)
    return url_list, features
    
def get_image_from_url(idx,url_list,items=5):
    '''
    Purpose: Get a list of images from the url list using the indecies
    Inputs:
        idx(list): list of indecies
        url_list(list): list of urls
        items(int): number of images to return
    Outputs:
        images(list): list of images
    '''

    # Looping until we have the desired number of images

    res = []
    i = 0
    n = 15 if len(idx) > 15 else len(idx)
    while len(res) != items and i != n:
        try:
            req = requests.get(url_list[idx[i]],stream=True,timeout=5)
            img = Image.open(req.raw).convert('RGB')
            img = np.array(img)
            res.append(img)
            i += 1
        except:
            print('Error with: ' + url_list[i])
            i += 1
            continue
    return res

def get_label(img):
    '''
    Purpose: Get the label of the image
    Inputs:
        img(numpy array or pil image): image to get label of
    Outputs:
        label(string): label of the image
    '''
    img_features = processor_hg(images=img, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        img_features = clip_hg.get_image_features(**img_features)
    idx = find_closest(img_features,item_embeddings)[0]
    label = all_items[idx]
    return label

def resize_img(img,thresh=384):
    '''
    Purpose: Resize the image to have the largest dimension be thresh
    Inputs:
        img(pil image): image to resize
        thresh(int): threshold for the largest dimension
    Outputs:
        img(pil image): resized image
    '''
    size = img.size
    larger_dim = 0 if size[0] > size[1] else 1
    if size[larger_dim] > thresh:
        size = (int(size[0] * thresh / size[larger_dim]), int(size[1] * thresh / size[larger_dim]))
        img = img.resize(size)
    return img

def segment_function(choice):
    '''
    Purpose: Set the mask choice so that it can be called during search
    Inputs:
        choice(string): mask choice
    Outputs:
        None
    '''
    global mask_choice
    mask_choice = choice
    return gr.update(label =f'Selection: {choice}')


def rec_function(option):
    '''
    Purpose: using an image to get a reccomendation return that image and the reccomendations
    Inputs:
        option(int): option to use
    Outputs:
        rec_out(list): list of images
        temp_out(numpy array): choice image
    '''
    global out
    global url_list_main
    global features_main

    # Here we get the items which should be reccomended based on the 
    # chosen image. Then we find the closest items to the chosen image
    # out of the reccomended items. Finally we crop the images so that
    # we only see the reccomended items in the output

    if not out:
        req = requests.get(url,stream=True)
        img = np.array(Image.open(req.raw).convert('RGB'))
        rec_out = [img]*5
        return rec_out
    img = Image.fromarray(out[option])
    choice_img = resize_img(img)
    label = get_label(choice_img)
    target_labels = label_to_rec_lables(label)
    temp_out = []
    img_features = processor_hg(images=choice_img, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        img_features = clip_hg.get_image_features(**img_features)
    n = len(target_labels)
    if n == 1:
        return_items = 5
    elif n == 2:
        return_items = 3
    elif n == 3:
        return_items = 2
    else:
        return_items = 1
    for item in target_labels:
        url_list, features = filter_features(label, rec=True, rec_items=item)
        idx = find_closest(img_features, features)[:5]
        temp_out.extend(get_image_from_url(idx,url_list,items=return_items))
    rec_out = []
    for temp_img in temp_out:
        temp_img = resize_img(Image.fromarray(temp_img))
        img, seg_img, out_mask = segment(temp_img)
     
        label_type = label.split(',')[0].strip()
        bad = []
        if label_type in top_type or label_type in dress_type:
            bad.extend([0,1,2,3,4,7,8,11,14,15,16])
        elif label_type in bottom_type:
            bad.extend([0,5,6,9,10,12,13,16])
        for i in bad:
            out_mask[out_mask==i] = 50
        img = np.array(img)
        img[out_mask==50] = 255
        h, w = img.shape[:2]
        # find the highest and lowest y-coordinates where the pixel is not white
        top = 0
        bottom = h
        for i in range(h):
            if np.all(img[i] == 255):
                top = i
            else:
                break
        for i in range(h-1, 0, -1):
            if np.all(img[i] == 255):
                bottom = i
            else:
                break
        # find the highest and lowest x-coordinates where the pixel is not white
        left = 0
        right = w
        for i in range(w):
            if np.all(img[:, i] == 255):
                left = i
            else:
                break   
        for i in range(w-1, 0, -1):
            if np.all(img[:, i] == 255):
                right = i
            else:
                break
        # crop the image
        # add 10 pixels to the top and bottom if those are not the edges of the image
        if top - 10 > 0:
            top -= 10
        if bottom  + 10 < h:
            bottom += 10
        # add 10 pixels to the left and right if those are not the edges of the image
        if left - 10 > 0:
            left -= 10
        if right + 10 < w:
            right += 10
        
        if top > bottom or right < left:
            rec_out.append(temp_img)
        else:
            temp_img = np.array(temp_img)
            img = temp_img[top:bottom, left:right]
            rec_out.append(img)
    temp_out = [choice_img]
    return rec_out, temp_out
    
def reset_values():
    '''
    Purpose: reset the values of the global variables
    Inputs:
        None
    Outputs:
        None
    '''
    global global_mask
    global out
    global mask_choice
    global clicks
    global c_types
    global types
    global colors
    global patterns
    global new_files
    global_mask = None
    out = None
    mask_choice = None
    clicks = 0
    c_types.clear()
    types.clear()
    colors.clear()
    patterns.clear()
    new_files.clear()
    return [gr.update(choices=['top','bottom','dress'],value=[]),gr.update(choices=['Person','Clothes','Upper Body/Dress','Upper Body/Dress, no person','Lower Body'],value=None)
            ,gr.update(value=None), gr.update(value=[]),gr.update(value=[]),gr.update(value=0)]

def search_function(img, text, use_choice,use_label):
    '''
    Purpose: search for images based on the text input or image input
    Inputs:
        img(pil image): image input
        text(string): text input
        use_choice(boolean): Boolen to know if to use image or text
        use_label(boolean): whether to use the label
    Outputs:
        out(list): list of images
    '''
    global new_files
    global global_mask
    global out
    use_img = False
    use_text = False
    if use_choice == 'Use Image':
        use_img = True
    elif use_choice == 'Use Text':
        use_text = True

    if new_files:
        global url_list_main
        global features_main
        temp_url = []
        new_files = list(set(new_files))
        temp_features = torch.zeros(len(new_files), 512).to(device)
        for c,i in enumerate(new_files):
            temp_url.append(url_list_main[i])
            temp_features[c] = features_main[i]
        url_list = temp_url
        features = temp_features.to(torch.float32)
    else:
        features = features_main.clone()
        url_list = url_list_main.copy()


    if use_text and not use_img:
        text_features = processor_hg(text=text, return_tensors="pt", padding=True).to(device)
        with torch.no_grad():
            text_features = clip_hg.get_text_features(**text_features)
        idx = find_closest(text_features, features)[:15]
        out = get_image_from_url(idx,url_list)
    else :
        if not isinstance(global_mask,type(None)):
            seg_img = clean_img(img)
        else:
            seg_img = img
        img = Image.fromarray(seg_img)
        label = get_label(img)
        if not new_files and use_label:
            url_list, features = filter_features(label)
        img_features = processor_hg(images=img, return_tensors="pt", padding=True).to(device)
        with torch.no_grad():
            img_features = clip_hg.get_image_features(**img_features)
        idx = find_closest(img_features, features)[:15]
        out = get_image_from_url(idx,url_list)
    if use_img:
        out.pop()
        out.insert(0, seg_img)
    return out

def search(img,text, choice,use_label,rotation):
    global global_mask
    try: 
        img = Image.fromarray(img).convert('RGB')
    except:
        img = Image.open(requests.get(url, stream=True).raw).convert('RGB')
    img = img.rotate(rotation)
    img = resize_img(img)
    pil, img, out_mask = segment(img)
    global_mask = out_mask
    res = search_function(img, text, choice,use_label)
    return res 

# Define the app layout

with gr.Blocks() as demo:
    gr.Markdown("Search using image segmentation")
    with gr.Tab("Search"):
        with gr.Row():
            search_image = gr.Image()
            search_input = [search_image,gr.Textbox(lines=2, label="Search Text")]
            with gr.Column():
                search_type = gr.Radio(choices=['Use Image','Use Text'],label='Select the type of search you want to perform',value='Use Image')
                use_label = gr.Checkbox(label="Use Label",value=True)
            image_output = [gr.Gallery(label='Outputs')]
            rec_out = [gr.Gallery(label='Recommendations',interactive=True)]
        with gr.Row():
            rec_selector = gr.Radio(label='Select which item you want a recommendation for',choices = [1,2,3,4],value=1)
            rec_button = gr.Button("Get Recommendation")
        with gr.Row():
            clothes_selector = gr.Radio(label='Choose a segmentation',
                choices=['Person','Clothes','Upper Body/Dress','Upper Body/Dress, no person','Lower Body'],interactive=True)
            theme_radio = gr.Radio(label='Choose a theme',choices=['None','Red carpet','Sports'],interactive=True)
            rotation_radio = gr.Radio(label='Choose a rotation',choices=[0,90,180,270],interactive=True,value=0)
        with gr.Row():
            filter_checkbox = gr.CheckboxGroup(label='Choose the clothing types', choices=['top','bottom','dress'],interactive=True,value=['top'])
            filter_button = gr.Button("Filter Button")
        search_button = gr.Button("Search Button")

    clothes_selector.change(segment_function,inputs=[clothes_selector],outputs=clothes_selector)
    search_image.change(reset_values, inputs=None, outputs=[filter_checkbox,clothes_selector,theme_radio,image_output[0],rec_out[0],rotation_radio])
    theme_radio.change(set_theme, inputs=theme_radio, outputs=theme_radio)
    rec_button.click(rec_function, inputs=rec_selector, outputs=[rec_out[0],image_output[0]])
    filter_button.click(filter_function, inputs=filter_checkbox, outputs=filter_checkbox)
    search_button.click(search, inputs=search_input+[search_type,use_label,rotation_radio], outputs=image_output)

demo.launch(share=False)