import streamlit as st
import pandas as pd


import seaborn as sns
import matplotlib.pyplot as plt

# Apply the default theme and activate color codes
sns.set_theme()
sns.set(color_codes=True)


################### import dataset
penguins = sns.load_dataset("penguins")

########### set the title and subtitle
st.title("Correlation of Penguin Measurements")

########## we add image
st.image("https://raw.githubusercontent.com/allisonhorst/palmerpenguins/main/man/figures/lter_penguins.png")

############## we create filters for our interactive plot
with st.sidebar:
    st.subheader("Filters")

    all_species = sorted(penguins["species"].dropna().unique().tolist())
    selected_species = st.multiselect(
        "Species to show",
        options=all_species,
        default=all_species,
    )

    all_measurements = {
    "Flipper length (mm)": "flipper_length_mm",
    "Bill length (mm)": "bill_length_mm",
    "Bill depth (mm)": "bill_depth_mm",
    "Body mass (g)": "body_mass_g",
    }
    measurement_options = st.multiselect(
    "Measurements to show",
    options=list(all_measurements.keys()),
    default=list(all_measurements.keys())
    )


if not selected_species:
    st.info("Select at least one species to display the plot.")
elif not measurement_options:
    st.info("Select at least one measurement to display the correlation matrix.")
else:
    data = penguins[penguins["species"].isin(selected_species)]
    selected_columns = [all_measurements[m] for m in measurement_options]
    corr_matrix = data[selected_columns].corr()

    labels = [" ".join(selected_columns[i].split("_")[:-1]).title() for i in range(len(selected_columns))]
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(
        corr_matrix,
        annot=True,
        cmap="coolwarm",
        fmt=".2f",
        xticklabels=labels,   # custom column labels
        yticklabels=labels    # custom row labels
    )
    plt.title("Correlation Matrix of Penguin Measurements")
    fig = plt.gcf()
    st.pyplot(fig)
    plt.close()

    # Save image button
    from io import BytesIO
    buf = BytesIO()
    fig.savefig(buf, format="png", dpi=200, bbox_inches="tight") 
    buf.seek(0)
    st.download_button(
        "Save image",
        data=buf,
        file_name="penguins_corr_matrix.png",
        mime="image/png",
    )

########### Footer
st.caption("Developed for SDS M1 course.")