# 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, CreateEvaluation
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.")

    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..."):
        try:
            # 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
        except Exception as e:
            logging.error(e)
            st.error(
                "Failed to get a prediction and explanation."
                + "Check whether you are using the right model URL and token for predictions. "
                + "Contact Deeploy if the problem persists."
            )

    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(agree: str, comment: str = None):
    if agree == "yes":
        evaluation_input: CreateEvaluation = {
            "agree": True,
            "comment": comment,
        }
    else:
        desired_output = not predictions[0]
        evaluation_input: CreateEvaluation = {
            "agree": False,
            "desired_output": { "predictions": [desired_output] },
            "comment": comment,
        }

    try:
        client.evaluate(st.session_state.deployment_id, prediction_log_id, 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'")

            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", comment]
                )
                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", comment]
                )
                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."
        )