# type: ignore -- ignores linting import issues when using multiple virtual environments
import streamlit.components.v1 as components
import streamlit as st
import pandas as pd
import logging
from deeploy import Client
from constants import (
    relationship_dict,
    occupation_dict,
    education_dict,
    type_of_work_dict,
    countries_dict,
    marital_status_dict,
)

# reset Plotly theme after streamlit import
import plotly.io as pio


pio.templates.default = "plotly"

logging.basicConfig(level=logging.INFO)

st.set_page_config(layout="wide")

st.title("Loan application model example")

st.markdown(
    """
    <style>
        section[data-testid="stSidebar"] {
            width: 300px !important; # Set the width to your desired value
        }
    </style>
    """,
    unsafe_allow_html=True,
)


def get_model_url():
    model_url = st.text_area(
        "Model URL (without the /explain endpoint, default is the demo deployment)",
        "https://api.app.deeploy.ml/workspaces/708b5808-27af-461a-8ee5-80add68384c7/deployments/dc8c359d-5f61-4107-8b0f-de97ec120289/",
        height=125,
    )
    elems = model_url.split("/")
    try:
        workspace_id = elems[4]
        deployment_id = elems[6]
    except IndexError:
        workspace_id = ""
        deployment_id = ""
    return model_url, workspace_id, deployment_id


def ChangeButtonColour(widget_label, font_color, background_color="transparent"):
    # func to change button colors
    htmlstr = f"""
        <script>
            var elements = window.parent.document.querySelectorAll('button');
            for (var i = 0; i < elements.length; ++i) {{ 
                if (elements[i].innerText == '{widget_label}') {{ 
                    elements[i].style.color ='{font_color}';
                    elements[i].style.background = '{background_color}'
                }}
            }}
        </script>
        """
    components.html(f"{htmlstr}", height=0, width=0)


with st.sidebar:
    st.image("deeploy_logo_wide.png", width=250)

    # Ask for model URL and token
    host = st.text_input("Host (Changing is optional)", "app.deeploy.ml")
    model_url, workspace_id, deployment_id = get_model_url()
    deployment_token = st.text_input("Deeploy Model Token", "my-secret-token")
    if deployment_token == "my-secret-token":
        st.warning("Please enter Deeploy API token.")
    # Split model URL into workspace and deployment ID

    # st.write("Values below are for debug only:")
    # st.write("Workspace ID: ", workspace_id)
    # st.write("Deployment ID: ", deployment_id)

    client_options = {
        "host": host,
        "deployment_token": deployment_token,
        "workspace_id": workspace_id,
    }
    client = Client(**client_options)


if "expander_toggle" not in st.session_state:
    st.session_state.expander_toggle = True

if "evaluation_submitted" not in st.session_state:
    st.session_state.evaluation_submitted = False

if "predict_button_clicked" not in st.session_state:
    st.session_state.predict_button_clicked = False

if "eval_selected" not in st.session_state:
    st.session_state["eval_selected"] = False

if "exp" not in st.session_state:
    st.session_state.exp = None

def set_prediction_view():
    st.session_state.predict_button_clicked = True
    st.session_state.evaluation_submitted = False    
    hide_expander()

def hide_expander():
    st.session_state.expander_toggle = False

def show_expander():
    st.session_state.expander_toggle = True

def submit_and_clear():
    try:
        # Call the explain endpoint as it also includes the prediction
        client.evaluate(
            deployment_id, request_log_id, prediction_log_id, st.session_state.evaluation_input
        )
        st.session_state.eval_selected = False
        st.session_state.evaluation_submitted = True
        st.session_state.eval_selected = False
        st.session_state.predict_button_clicked = False
        st.session_state.exp = None
        show_expander()
    except Exception as e:
        logging.error(e)
        st.error(
            "Failed to submit feedback."
            + "Check whether you are using the right model URL and token for evaluations. "
            + "Contact Deeploy if the problem persists."
        )

# Attributes
st.subheader("Loan Application")
with st.expander("Application form", expanded=st.session_state.expander_toggle):
    # Split view in 2 columns
    col1, col2 = st.columns(2)
    with col1:
        # Create input fields for attributes from constant dicts
        age = st.number_input("Age", min_value=0, max_value=100, value=30)
        marital_status = st.selectbox("Marital Status", marital_status_dict.keys())
        marital_status_id = marital_status_dict[marital_status]
        native_country = st.selectbox(
            "Native Country", countries_dict.keys(), index=len(countries_dict) - 1
        )
        native_country_id = countries_dict[native_country]
        relationship = st.selectbox("Family situation", relationship_dict.keys())
        relationship_id = relationship_dict[relationship]
        occupation = st.selectbox("Occupation", occupation_dict.keys(), index=1)
        occupation_id = occupation_dict[occupation]

    with col2:
        education = st.selectbox("Highest education level", education_dict.keys(), index=4)
        education_id = education_dict[education]
        type_of_work = st.selectbox("Type of work", type_of_work_dict.keys())
        type_of_work_id = type_of_work_dict[type_of_work]
        hours_per_week = st.number_input(
            "Working hours per week", min_value=0, max_value=100, value=40
        )
        capital_gain = st.number_input(
            "Yearly income [€]", min_value=0, max_value=1000000, value=70000
        )
        capital_loss = st.number_input(
            "Yearly expenditures [€]", min_value=0, max_value=1000000, value=60000
        )
    data_df = pd.DataFrame(
        [
            [
                age,
                type_of_work,
                education,
                marital_status,
                occupation,
                relationship,
                capital_gain,
                capital_loss,
                hours_per_week,
                native_country,
            ]
        ],
        columns=[
            "Age",
            "Type of work",
            "Highest education level",
            "Marital Status",
            "Occupation",
            "Family situation",
            "Yearly Income [€]",
            "Yearly expenditures [€]",
            "Working hours per week",
            "Native Country",
        ],
    )
    data_df_t = data_df.T
    request_body = {
        "instances": [
            [
                age,
                type_of_work_id,
                education_id,
                marital_status_id,
                occupation_id,
                relationship_id,
                capital_gain,
                capital_loss,
                hours_per_week,
                native_country_id,
            ]
        ]
    }

# Show predict button if token is set
if deployment_token != "my-secret-token" and st.session_state.exp is None:
    predict_button = st.button(
        "Send loan application", key="predict_button", help="Click to get the AI prediction.", on_click=set_prediction_view
    )
    if predict_button:
        with st.spinner("Loading prediction and explanation..."):
            # Call the explain endpoint as it also includes the prediction
            exp = client.explain(
                request_body=request_body, deployment_id=deployment_id
            )
            st.session_state.exp = exp

# Show prediction and explanation after predict button is clicked
if st.session_state.predict_button_clicked and st.session_state.exp is not None:
    try:
        exp = st.session_state.exp
        # Read explanation to dataframe from json
        predictions = exp["predictions"]
        request_log_id = exp["requestLogId"]
        prediction_log_id = exp["predictionLogIds"][0]

        exp_df = pd.DataFrame(
            [exp["explanations"][0]["shap_values"]], columns=exp["featureLabels"]
        )

        exp_df.columns = data_df.columns

        exp_df_t = exp_df.T

        # Merge data and explanation
        exp_df_t = data_df_t.merge(exp_df_t, left_index=True, right_index=True)

        weight_feat = "Weight"
        exp_df_t.columns = ["Feature value", weight_feat]
        exp_df_t["Feature"] = exp_df_t.index
        exp_df_t = exp_df_t[["Feature", "Feature value", weight_feat]]
        exp_df_t["Feature value"] = exp_df_t["Feature value"].astype(str)
        
        # Filter values below 0.01
        exp_df_t = exp_df_t[
            (exp_df_t[weight_feat] > 0.01) | (exp_df_t[weight_feat] < -0.01)
        ]
        exp_df_t[weight_feat] = exp_df_t[weight_feat].astype(float).round(2)

        pos_exp_df_t = exp_df_t[exp_df_t[weight_feat] > 0]
        pos_exp_df_t = pos_exp_df_t.sort_values(by=weight_feat, ascending=False)

        neg_exp_df_t = exp_df_t[exp_df_t[weight_feat] < 0]
        neg_exp_df_t = neg_exp_df_t.sort_values(by=weight_feat, ascending=True)
        neg_exp_df_t[weight_feat] = neg_exp_df_t[weight_feat].abs()

        # Get 3 features with highest positive relevance score
        pos_feats = pos_exp_df_t[weight_feat].nlargest(3).index.tolist()
        # For feature, get feature value and concatenate into a single string
        pos_feats = [
            f"{feat}: {pos_exp_df_t.loc[feat, 'Feature value']}"
            for feat in pos_feats
        ]
        # Get 3 features with highest negative relevance score
        neg_feats = neg_exp_df_t[weight_feat].nlargest(3).index.tolist()
        # For feature, get feature value and concatenate into a single string
        neg_feats = [
            f"{feat}: {neg_exp_df_t.loc[feat, 'Feature value']}"
            for feat in neg_feats
        ]
        if not st.session_state.evaluation_submitted:  
            if predictions[0]:
                # Show prediction
                st.subheader("Loan Decision: :green[Approve]", divider="green")
                # Format subheader to green
                st.markdown(
                    "<style>.css-1v3fvcr{color: green;}</style>", unsafe_allow_html=True
                )
                col1, col2 = st.columns(2)
                with col1:
                    # If prediction is positive, first show positive features, then negative features
                    st.success(
                        "The most important characteristics in favor of loan approval are:  \n - "
                        + "  \n- ".join(pos_feats)
                    )
                with col2:
                    st.warning(
                        "However, the following features weight against the loan applicant:  \n - "
                        + "  \n- ".join(neg_feats)
                        # + "  \n For more details, see full explanation of the credit assessment below.",
                    )
            else:
                st.subheader("Loan Decision: :red[Reject]", divider="red")
                col1, col2 = st.columns(2)
                with col1:
                    # If prediction is negative, first show negative features, then positive features
                    st.error(
                        "The most important characteristics in favor of loan rejection are:  \n - "
                        + "  \n - ".join(neg_feats)
                    )
                with col2:
                    st.warning(
                        "However, the following factors weigh in favor of the loan applicant:  \n - "
                        + "  \n - ".join(pos_feats)
                        # + "  \n For more details, see full explanation of the credit assessment below.",
                    )
            explanation_expander = st.expander("Show explanation")
            with explanation_expander:
                # Show explanation
                col_pos, col_neg = st.columns(2)

                with col_pos:
                    st.subheader("Factors :green[in favor] of loan approval")
                    # st.success("**Factors in favor of loan approval**")
                    st.dataframe(
                        pos_exp_df_t,
                        hide_index=True,
                        width=600,
                        column_config={
                            "Weight": st.column_config.ProgressColumn(
                                "Weight",
                                width="small",
                                format=" ",
                                min_value=0,
                                max_value=1,
                            )
                        },
                    )

                with col_neg:
                    st.subheader("Factors :red[against] loan approval")
                    # st.error("**Factors against loan approval**")
                    st.dataframe(
                        neg_exp_df_t,
                        hide_index=True,
                        width=600,
                        column_config={
                            "Weight": st.column_config.ProgressColumn(
                                "Weight",
                                width="small",
                                format=" ",
                                min_value=0,
                                max_value=1,
                            )
                        },
                    )

            st.divider()
          
            # Add prediction evaluation
            st.subheader("Evaluation: Do you agree with the loan assessment?")
            st.write(
                "AI model predictions always come with a certain level of uncertainty. Evaluate the correctness of the assessment based on your expertise and experience."
            )
            cols = st.columns(4)
            col_yes, col_no = cols[:2]
            with col_yes:
                yes_button = st.button(
                    "Yes, I agree",
                    key="yes_button",
                    use_container_width=True,
                    help="Click if you agree with the prediction",
                )
                ChangeButtonColour("Yes, I agree", "white", "green")
            with col_no:
                no_button = st.button(
                    "No, I disagree",
                    key="no_button",
                    use_container_width=True,
                    help="Click if you disagree with the prediction",
                    type="primary",
                )
                ChangeButtonColour("No, I disagree", "white", "#DD360C")
                # ChangeButtonColour("No, I disagree", "#DD360C", "#F0F0F0")
            
            if yes_button:
                st.session_state.eval_selected = True
                st.session_state.evaluation_input = {
                    "result": 0  # Agree with the prediction
                }
                st.session_state.placeholder = "Income is sufficient, given applicant's background"
            if no_button:
                st.session_state.eval_selected = True
                desired_output = not predictions[0]
                st.session_state.evaluation_input = {
                    "result": 1,  # Disagree with the prediction
                    "value": {"predictions": [desired_output]},
                }
                st.session_state.placeholder = "Income is too low, given applicant's background"

            if st.session_state.eval_selected:
                comment = st.text_input("Would you like to add a comment?", placeholder="For example: '" + st.session_state.placeholder + "'")
                if comment:
                    st.session_state.evaluation_input["explanation"] = comment

                st.button("Submit", key="submit_button", on_click=submit_and_clear)
        else:
            st.success("Evaluation submitted successfully!")                    

    except Exception as e:
        logging.error(e)
        st.error(
            "Failed to retrieve the prediction or explanation."
            + "Check whether you are using the right model URL and Token. "
            + "Contact Deeploy if the problem persists."
        )