# 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")


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()
    st.session_state.deployment_id = deployment_id
    deployment_token = st.text_input("Deeploy API 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 "request_body" not in st.session_state:
    st.session_state.request_body = None

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

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

def form_request_body():
    """Create the request body for the prediction endpoint"""
    marital_status_id = marital_status_dict[st.session_state.marital_status]
    native_country_id = countries_dict[st.session_state.native_country]
    relationship_id = relationship_dict[st.session_state.relationship]
    occupation_id = occupation_dict[st.session_state.occupation]
    education_id = education_dict[st.session_state.education]
    type_of_work_id = type_of_work_dict[st.session_state.type_of_work]
    return {
        "instances": [
            [
                st.session_state.age,
                type_of_work_id,
                education_id,
                marital_status_id,
                occupation_id,
                relationship_id,
                st.session_state.capital_gain,
                st.session_state.capital_loss,
                st.session_state.hours_per_week,
                native_country_id,
            ]
        ]
    }


def predict_callback():
    """Callback function to call the prediction endpoint"""
    request_body = form_request_body() # Make sure we have the latest values after user input
    st.session_state.exp = None
    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=st.session_state.deployment_id
        )
        st.session_state.exp = exp

    st.session_state.predict_button_clicked = True
    st.session_state.evaluation_submitted = False    

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

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

def submit_and_clear(evaluation: str):
    if evaluation == "yes":
        st.session_state.evaluation_input["result"] = 0 # Agree with the prediction
    else:
        desired_output = not predictions[0]
        st.session_state.evaluation_input["result"] = 1
        st.session_state.evaluation_input["value"] = {"predictions": [desired_output]}
    try:
        client.evaluate(
            deployment_id, request_log_id, prediction_log_id, st.session_state.evaluation_input
        )
        st.session_state.evaluation_submitted = True
        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."
        )

# with st.expander("Debug session state", expanded=False):
#     st.write(st.session_state)

# Attributes
with st.expander("**Loan 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=10, max_value=100, value=30, key="age", on_change=predict_callback)
        marital_status = st.selectbox("Marital Status", marital_status_dict.keys(), key="marital_status", on_change=predict_callback,)
        native_country = st.selectbox(
            "Native Country", countries_dict.keys(), index=len(countries_dict) - 1, key="native_country",on_change=predict_callback
        )
        relationship = st.selectbox("Family situation", relationship_dict.keys(), key="relationship", on_change=predict_callback)
        occupation = st.selectbox("Occupation", occupation_dict.keys(), index=1, key="occupation", on_change=predict_callback)

    with col2:
        education = st.selectbox("Highest education level", education_dict.keys(), key="education", index=4, on_change=predict_callback)
        type_of_work = st.selectbox("Type of work", type_of_work_dict.keys(), key="type_of_work", on_change=predict_callback)
        hours_per_week = st.number_input(
            "Working hours per week", min_value=0, max_value=100, value=40, key="hours_per_week", on_change=predict_callback,
        )
        capital_gain = st.number_input(
            "Yearly income [€]", min_value=0, max_value=10000000, value=70000, key="capital_gain", on_change=predict_callback,
        )
        capital_loss = st.number_input(
            "Yearly expenditures [€]", min_value=0, max_value=10000000, value=60000, key="capital_loss", on_change=predict_callback,
        )
    data_df = pd.DataFrame(
        [
            [
                st.session_state.age,
                st.session_state.type_of_work,
                st.session_state.education,
                st.session_state.marital_status,
                st.session_state.occupation,
                st.session_state.relationship,
                st.session_state.capital_gain,
                st.session_state.capital_loss,
                st.session_state.hours_per_week,
                st.session_state.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

# 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=predict_callback,
    )        

if st.session_state.evaluation_submitted:
    st.success("Evaluation submitted successfully!")

# Show prediction and explanation after predict button is clicked
elif 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"
        feat_val_col = "Value"
        exp_df_t.columns = [feat_val_col, weight_feat]
        exp_df_t["Feature"] = exp_df_t.index
        exp_df_t = exp_df_t[["Feature", feat_val_col, weight_feat]]
        exp_df_t[feat_val_col] = exp_df_t[feat_val_col].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, feat_val_col]}"
            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, feat_val_col]}"
            for feat in neg_feats
        ]
        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.error(
                    "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 reasons for loan rejection are:  \n - "
                    + "  \n - ".join(neg_feats)
                )
            with col2:
                st.success(
                    "However, the following factors weigh in favor of the loan applicant:  \n - "
                    + "  \n - ".join(pos_feats)
                )
        try:
            # Show explanation
            if predictions[0]:
                col_pos, col_neg = st.columns(2)
            else:
                col_neg, col_pos = st.columns(2) # Swap columns if prediction is negative
            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,
                        )
                    },
                )
        except Exception as e:
            logging.error(e)
            st.error(
                "Failed to show the explanation."
                + "Refresh the page to reset the application."
                + "Contact Deeploy if the problem persists."
            )

        st.divider()

        if not st.session_state.evaluation_submitted:
            # 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."
            )
            st.session_state.evaluation_input = {}
            comment = st.text_input("Your assessment:", placeholder="For example: 'Income is too low, given applicant's background'")
            if comment:
                st.session_state.evaluation_input["explanation"] = comment

            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",
                    on_click=submit_and_clear,
                    args=["yes"]
                )
                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",
                    on_click=submit_and_clear,
                    args=["no"]
                )
                ChangeButtonColour("No, I disagree", "white", "#DD360C") # Red color for disagree button

    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."
        )