import os
os.system('pip install openpyxl')
os.system('pip install sentence-transformers')
import pandas as pd
import gradio as gr
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('all-mpnet-base-v2') #all-MiniLM-L6-v2 #all-mpnet-base-v2

df = pd.read_parquet('df_encoded.parquet')
df['tags'] = df['tags'].apply(lambda x : str(x))
def parse_raised(x):
    if x == 'Undisclosed':
        return 0
    else: 
        quantifier = x[-1]
        x = float(x[1:-1])
        if quantifier == 'K':
            return x/1000
        elif quantifier == 'M':
            return x
df['raised'] = df['raised'].apply(lambda x : parse_raised(x))
df = df.reset_index(drop=True)

from sklearn.neighbors import NearestNeighbors
import pandas as pd
from sentence_transformers import SentenceTransformer

def filter_df(df, column_name, filter_type, filter_value):
    if filter_type == '==':
        df_filtered = df[df[column_name]==filter_value]
    elif filter_type == '>=':
        df_filtered = df[df[column_name]>=filter_value]
    elif filter_type == '<=':
        df_filtered = df[df[column_name]<=filter_value]
    elif filter_type == 'contains':
        df_filtered = df[df['target'].str.contains(filter_value)]
    return df_filtered

def search(df, query):
    product = model.encode(query).tolist()
    # product = df.iloc[0]['text_vector_'] #use one of the products as sample

    #prepare model
    nbrs = NearestNeighbors(n_neighbors=20, algorithm='ball_tree').fit(df['text_vector_'].values.tolist())

    distances, indices = nbrs.kneighbors([product]) #input the vector of the reference object

    #print out the description of every recommended product
    return df.iloc[list(indices)[0]][['name', 'description', 'raised', 'year', 'target', 'size', 'stage', 'tags']]

#the first module becomes text1, the second module file1
def greet(size, target, raised, query): 
    df_size = filter_df(df, 'size', '==', size)
    df_target = filter_df(df_size, 'target', 'contains', target)
    def raised_zero(x):
        if x == 0:
            return 'Undisclosed'
        else:
            return x
    print('a')
    df_raised = df_target[(df_target['raised'] >= raised) | (df_target['raised'] == 0)]
    df_knn = search(df_raised, query)
    #we live the sorting for last
    df_knn = df_knn.sort_values('raised', ascending=False)
    df_knn['raised'] = df_knn['raised'].apply(lambda x : raised_zero(x))

    return df_knn

with gr.Blocks(theme=gr.themes.Soft(primary_hue='amber', secondary_hue='gray', neutral_hue='amber')) as demo:
    gr.Markdown(
    """
    # Startup Search Engine
    """
    )
    size = gr.Radio(['1-10', '11-50', '51-200', '201-500', '500+'], multiselect=False, value='11-50', label='size')
    target = gr.Radio(['B2B', 'B2C', 'B2G', 'B2B2C'], value='B2B', multiselect=False, label='target')
    raised = gr.Slider(0, 20, value=5, step_size=1, label="Minimum raising (in Millions)")
    query = gr.Textbox(label='Describe the Startup you are searching for', value='age reversing')
    btn = gr.Button(value="Search for a Startup")
    output1 = gr.DataFrame(label='value')
    # btn.click(greet, inputs='text', outputs=['dataframe'])
    btn.click(greet, [size, target, raised, query], [output1])
demo.launch(share=False)