from typing import List

import gradio as gr
from angle_emb import AnglE
from scipy import spatial

def cosine_similarity(vec1: List[int], vec2: List[int]):
    """ Calculate cosine similarity between two vectors.

    :param vec1: a list of integers
    :param vec2: a list of integers
    :return: a float value between 0 and 1, indicating the similarity between the two vectors.
    """
    return 1 - spatial.distance.cosine(vec1, vec2)
    

uae = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1').to('cpu')
uae_code = AnglE.from_pretrained('WhereIsAI/UAE-Code-Large-V1').to('cpu')


def code_similarity(code1, code2, model):
    if model == "UAE-Code-Large-V1":
        v1 = uae_code.encode(code1)
        v2 = uae_code.encode(code2)
    else:
        v1 = uae.encode(code1)
        v2 = uae.encode(code2)
    return cosine_similarity(v1[0], v2[0])


demo = gr.Interface(
    code_similarity,
    [
        gr.Textbox(label='code 1'),
        gr.Textbox(label='code 2'),
        gr.Radio(["UAE-Code-Large-V1", "UAE-Large-V1"], label="model", value="UAE-Code-Large-V1"),
    ],
    "text",
    examples=[
        # sample 1
        ["""# Approach 2: Quicksort using list comprehension

def quicksort(arr):
    if len(arr) <= 1:
        return arr
    else:
        pivot = arr[0]
        left = [x for x in arr[1:] if x < pivot]
        right = [x for x in arr[1:] if x >= pivot]
        return quicksort(left) + [pivot] + quicksort(right)
 
# Example usage
arr = [1, 7, 4, 1, 10, 9, -2]
sorted_arr = quicksort(arr)
print("Sorted Array in Ascending Order:")
print(sorted_arr)""", """def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    else:
        return qsort([x for x in arr[1:] if x < arr[0]])
        + [arr[0]]
        + qsort([x for x in arr[1:] if x >= arr[0]])""", """UAE-Code-Large-V1"""
        ],
        # sample 2
        ["""def bubblesort(elements):
    # Looping from size of array from last index[-1] to index [0]
    for n in range(len(elements)-1, 0, -1):
        swapped = False
        for i in range(n):
            if elements[i] > elements[i + 1]:
                swapped = True
                # swapping data if the element is less than next element in the array
                elements[i], elements[i + 1] = elements[i + 1], elements[i]
        if not swapped:
            # exiting the function if we didn't make a single swap
            # meaning that the array is already sorted.
            return

elements = [39, 12, 18, 85, 72, 10, 2, 18]

print("Unsorted list is,")
print(elements)
bubblesort(elements)
print("Sorted Array is, ")
print(elements)""", """def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    else:
        return qsort([x for x in arr[1:] if x < arr[0]])
        + [arr[0]]
        + qsort([x for x in arr[1:] if x >= arr[0]])""", """UAE-Code-Large-V1"""
        ],
    ]
)

if __name__ == "__main__":
    demo.launch()