import json
import streamlit as st
import pandas as pd
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from st_social_media_links import SocialMediaIcons


AVERAGE_COLUMN_NAME = "Average"
SENTIMENT_COLUMN_NAME = "Sentiment"
RESULTS_COLUMN_NAME = "Results"
UNDERSTANDING_COLUMN_NAME = "Language understanding"
PHRASEOLOGY_COLUMN_NAME = "Phraseology"

# Function to load data from JSON file
def load_data(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    return pd.DataFrame(data)

# Function to style the DataFrame
def style_dataframe(df: pd.DataFrame):
    df[RESULTS_COLUMN_NAME] = df.apply(lambda row: [row[SENTIMENT_COLUMN_NAME], row[UNDERSTANDING_COLUMN_NAME], row[PHRASEOLOGY_COLUMN_NAME]], axis=1)

    # Insert the new column after the 'Average' column
    cols = list(df.columns)
    cols.insert(cols.index(AVERAGE_COLUMN_NAME) + 1, cols.pop(cols.index(RESULTS_COLUMN_NAME)))
    df = df[cols]

    # Create a color ramp using Seaborn
    return df

def styler(df: pd.DataFrame):
    palette = sns.color_palette("RdYlGn", as_cmap=True)
    styled_df = df.style.background_gradient(cmap=palette, subset=[AVERAGE_COLUMN_NAME, SENTIMENT_COLUMN_NAME, PHRASEOLOGY_COLUMN_NAME, UNDERSTANDING_COLUMN_NAME]).format(precision=2)
    return styled_df


### Streamlit app
st.set_page_config(layout="wide")

st.markdown("""
        <style>
               .block-container {
                    padding-top: 0%;
                    padding-bottom: 0%;
                    padding-left: 3%;
                    padding-right: 3%;
                    scrollbar-width: thin;
                }
        </style>
        """, unsafe_allow_html=True)

### Prepare layout
st.subheader("")

st.markdown("""
<style>
    .center {
        display: block;
        margin-left: auto;
        margin-right: auto;
        width: 50%;
    }
    .center-text {
        text-align: center;
    }
    .table-responsive {
        text-align: center;
        font-size: 0.9em;
        margin-left: 0%;
        margin-right: 0%;
        overflow-x: auto;
        -ms-overflow-style: 3px;  /* Internet Explorer 10+ */
        scrollbar-width: thin;  /* Firefox */
    }
    .table-responsive::-webkit-scrollbar { 
        /*display: none;*/  /* Safari and Chrome */
        width: 6px;
    }

    #table_id {
      display: block;
    }

    #table_id th {
      display: inline-block;
    }

    #table_id td {
        padding-left: 0.7rem; 
        padding-right: 0.7rem;
        display: inline-block;
    }
    #table_id td:hover {
        color:#FDA428;
    }

    a:link {color:#A85E00;}      /* unvisited link */
    a:hover {color:#FDA428;}   /* Mouse over link */
    a:visited {color:#A85E00;}  /* visited link */
    a:active {color:#A85E00;}  /* selected link */

    .image-container {
      position: relative;
      display: inline-block;
      transition: transform 0.3s ease;
    }

    .image-container img {
      vertical-align: middle;
    }

    .image-container::after {
      content: "";
      position: absolute;
      left: 0;
      bottom: 0;
      width: 100%;
      height: 2px;
      background-color: #FDA428; /* Change this to your desired color */
      transform: scaleX(0);
      transition: transform 0.3s ease;
    }

    .image-container:hover {
      transform: translateY(-3px); /* Change the value to adjust the upward movement */
    }

    .image-container:hover::after {
      transform: scaleX(1);
    }

/* ---------------------------------------------------------------- */
</style>
""", unsafe_allow_html=True)

# --- Colors info ---
# Primary Color: #FDA428
# Secondary Color: #A85E00
# Grey Color: #7B7B7B
# Background Color: #1C1C1C
# {'LOW': '#7B7B7B', 'MEDIUM': '#A85E00', 'HIGH': '#FDA428'}
# ----------------------------------------------------------

### Row: 1 --> Title + links to SpeakLeash.org website / GitHub / X (Twitter)
social_media_links = [
    "https://discord.com/invite/ZJwCMrxwT7",
    "https://github.com/speakleash",
    "https://x.com/Speak_Leash",
    "https://www.linkedin.com/company/speakleash/",
    "https://www.facebook.com/Speakleash/"
]

light_orange = "#FDA428"
dark_orange = "#A85E00"
white_color = "#FFFFFF"
black_color = "#000000"
links_color = white_color

# if theme.get('background_color') == white_color:
#     links_color = black_color

social_media_links_colors = [
    dark_orange,
    dark_orange,
    dark_orange,
    dark_orange,
    dark_orange
]

social_media_icons = SocialMediaIcons(social_media_links, social_media_links_colors)
social_media_icons.render(justify_content='right')

# Add logo, title, and subheader in a flexible container with equal spacing
st.markdown("""
    <div class="header-container">
        <br>
        <img src="https://speakleash.org/wp-content/uploads/2023/09/SpeakLeash_logo.svg" alt="SpeakLeash Logo">
        <hr>
        <div class="title-container">
            <h1 style='color: #FDA428; margin-top: -1rem; font-size: 3.1em;'>Phrase-Bench</h1>
            <h3 style="margin-top: 0;">Understanding of Polish text, sentiment and phraseological compounds</h2>
        </div>
    </div>
    """, unsafe_allow_html=True)

# Create tabs
tab1, tab2 = st.tabs([RESULTS_COLUMN_NAME, "Opis"])

with tab1:
    st.write("This benchmark evaluates the ability of language models to correctly interpret Polish texts with complex implicatures, such as sarcasm and idiomatic expressions. Models are assessed on sentiment analysis, understanding of true intentions, and identification of idiomatic phrases.")

    # Display the styled DataFrame
    data = load_data('data.json')
    data['Params'] = data['Params'].str.replace('B', '')
    data = data.sort_values(by=AVERAGE_COLUMN_NAME, ascending=False)
    styled_df_show = style_dataframe(data)
    styled_df_show = styler(styled_df_show)

    st.data_editor(styled_df_show, column_config={
        "Model": st.column_config.TextColumn("Model", help="Model name", width="large"),
        "Params": st.column_config.NumberColumn("Params [B]", format="%.1f"),
        AVERAGE_COLUMN_NAME: st.column_config.NumberColumn(AVERAGE_COLUMN_NAME),
        RESULTS_COLUMN_NAME: st.column_config.BarChartColumn(
            RESULTS_COLUMN_NAME, help="Summary of the results of each task",
            y_min=0,y_max=5,),
        SENTIMENT_COLUMN_NAME: st.column_config.NumberColumn(SENTIMENT_COLUMN_NAME, help='Ability to analyze sentiment'),
        PHRASEOLOGY_COLUMN_NAME: st.column_config.NumberColumn(PHRASEOLOGY_COLUMN_NAME, help='Ability to understand phraseological compounds'),
        UNDERSTANDING_COLUMN_NAME: st.column_config.NumberColumn(UNDERSTANDING_COLUMN_NAME, help='Ability to understand language'),
        }, hide_index=True, disabled=True, height=500)

    # st.divider()

    # Add selection for models and create a bar chart for selected models using the AVERAGE_COLUMN_NAME, SENTIMENT_COLUMN_NAME, PHRASEOLOGY_COLUMN_NAME, UNDERSTANDING_COLUMN_NAME
    # Add default selection of 3 best models from AVERAGE_COLUMN_NAME and 1 best model with "Bielik" in Model column
    default_models = list(data.sort_values(AVERAGE_COLUMN_NAME, ascending=False)['Model'].head(3))
    bielik_model = data[data['Model'].str.contains('Bielik')].sort_values(AVERAGE_COLUMN_NAME, ascending=False)['Model'].iloc[0]
    if bielik_model not in default_models:
        default_models.append(bielik_model)
    selected_models = st.multiselect("Select models to compare", data["Model"].unique(), default=default_models)
    # selected_models = st.multiselect("Select models to compare", data["Model"].unique())
    selected_data = data[data["Model"].isin(selected_models)]
    categories = [AVERAGE_COLUMN_NAME, SENTIMENT_COLUMN_NAME, PHRASEOLOGY_COLUMN_NAME, UNDERSTANDING_COLUMN_NAME]

    if selected_models:
        # Kolorki do wyboru:
        # colors = px.colors.sample_colorscale("viridis", len(selected_models)+1)
        colors = px.colors.qualitative.D3[:len(selected_models)]
        # Create a chart with bars for each model for each category
        fig_bars = go.Figure()
        for model, color in zip(selected_models, colors):
            values = selected_data[selected_data['Model'] == model][categories].values.flatten().tolist()
            fig_bars.add_trace(go.Bar(
                x=categories,
                y=values,
                name=model,
                marker_color=color
            ))

        # Update layout to use a custom color scale
        fig_bars.update_layout(
            showlegend=True,
            legend=dict(orientation="h", yanchor="top", y=-0.3, xanchor="center", x=0.5),
            title="Comparison of Selected Models",
            yaxis_title="Score",
            template="plotly_dark"
        )
        st.plotly_chart(fig_bars)


### Zakładka 2 --> Opis
with tab2:
    st.header("Opis")
    st.write("Tutaj znajduje się trochę tekstu jako wypełniacz.")
    st.write("To jest przykładowy tekst, który może zawierać dodatkowe informacje o benchmarku, metodologii, itp.")


# Ending :)
st.divider()
st.markdown("""
### Authors:
- [Jan Sowa](https://www.linkedin.com/in/janpiotrsowa) - leadership, writing texts, benchmark code
- [Agnieszka Kosiak](https://www.linkedin.com/in/agn-kosiak/) - writing texts
- [Magdalena Krawczyk](https://www.linkedin.com/in/magdalena-krawczyk-7810942ab/) - writing texts, labeling
- [Remigiusz Kinas](https://www.linkedin.com/in/remigiusz-kinas/) - methodological support
- [Krzysztof Wróbel](https://www.linkedin.com/in/wrobelkrzysztof/) - engineering, methodological support
- [Szymon Baczyński](https://www.linkedin.com/in/szymon-baczynski/) - front-end / streamlit assistant
""")

st.divider()

# Run the app with `streamlit run your_script.py`