import pandas as pd
import streamlit as st
from sklearn.preprocessing import LabelEncoder
from sklearn.ensemble import RandomForestRegressor

# Leer los datos y seleccionar las columnas necesarias
df = pd.read_csv('./ds_salaries.csv')
df = df[['company_location', 'salary_in_usd']]


# Codificar las ubicaciones de las empresas
le = LabelEncoder()
df['company_location'] = le.fit_transform(df['company_location'])

# Decodificar las ubicaciones de las empresas
decoded_locations = le.inverse_transform(df['company_location'].unique())

# Separar los datos de entrada y salida
X = df.iloc[:, :-1].values
y = df.iloc[:, -1].values

# Entrenar el modelo
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X, y)

# Obtener las ubicaciones de las empresas y sus salarios predichos
locations = df['company_location'].unique()
predicted_salaries = model.predict(locations.reshape(-1, 1))
results_df = pd.DataFrame({'company_location': locations, 'predicted_salary': predicted_salaries})

# Decodificar las ubicaciones de las empresas
results_df['company_location'] = le.inverse_transform(results_df['company_location'])

# Ordenar los resultados por salario predicho
results_df = results_df.sort_values('predicted_salary', ascending=False).reset_index(drop=True)



# Mostrar el título y el top 5 de países mejor pagados
st.markdown("""
    <h1>Top 5 de países mejor pagados</h1>
    """, unsafe_allow_html=True)

# Descripción
st.markdown("""
    <p>Este código utiliza un algoritmo de RandomForest para seleccionar solo algunas empresas de cada país de forma aleatoria y retornar el promedio. Utilizamos este algoritmo para obtener un promedio menos sesgado por outliers</p> 
""", unsafe_allow_html=True)

for i in range(5):
    location = results_df.loc[i, 'company_location']
    salary = results_df.loc[i, 'predicted_salary']
    st.markdown(f'### **{location}**: ${salary:,.2f}', unsafe_allow_html=True)

# Mostrar el menú desplegable para seleccionar un país
st.markdown('---')
st.title('Seleccionar un país')
selected_location = st.selectbox('Ubicación de la empresa', decoded_locations)

# Mostrar el salario predicho para el país seleccionado
predicted_salary = results_df.loc[results_df['company_location'] == selected_location, 'predicted_salary'].iloc[0]
st.markdown(f'### **{selected_location}**: ${predicted_salary:,.2f}', unsafe_allow_html=True)