import streamlit as st
import pandas as pd
import logging
from deeploy import Client
from utils import (
    get_request_body,
    get_fake_certainty,
    get_model_url,
    get_random_suspicious_transaction,
)
from utils import (
    get_explainability_texts,
    get_explainability_values,
    send_evaluation,
    get_comment_explanation,
)
from utils import COL_NAMES, feature_texts
from utils import (
    create_data_input_table,
    create_table,
    ChangeButtonColour,
    get_weights,
    modify_datapoint,
)

logging.basicConfig(level=logging.INFO)

st.set_page_config(layout="wide")

st.title("Smart AML:tm:")
st.divider()

# Import data
data = pd.read_pickle("data/preprocessed_data.pkl")

# instantiate important vars in session state
if "predict_button_clicked" not in st.session_state:
    st.session_state.predict_button_clicked = False

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

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

if "no_button_text" not in st.session_state:
    st.session_state.no_button_text = (
        "I don't think this transaction is money laundering because..."
    )

if "yes_button_text" not in st.session_state:
    st.session_state.yes_button_text = ""

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


# define functions to be run when buttons are clicked
# func to be run when input changes in no button text area
def get_input_no_button():
    st.session_state.no_button_text = comment.replace(
        st.session_state.no_button_text, st.session_state.no_comment
    )
    st.session_state.evaluation_input["explanation"] = st.session_state.no_button_text


# func to be run when input changes in yes button text area
def get_input_yes_button():
    st.session_state.yes_button_text = comment.replace(
        st.session_state.yes_button_text, st.session_state.yes_comment
    )
    st.session_state.evaluation_input["explanation"] = st.session_state.yes_button_text


# func to disable click again for button "Get suspicious transactions"
def disabled():
    st.session_state.disabled = True


# func for Next button to rerun and get new prediction
def rerun():
    st.session_state.predict_button_clicked = True
    st.session_state.submitted_disabled = False
    st.session_state.no_button_text = (
        "I don't think this transaction is money laundering because..."
    )


# func for submit button to disable resubmit
def submitted_disabled():
    st.session_state.submitted_disabled = True


# color specs for sidebar
st.markdown(
    """
    <style>
        [data-testid=stSidebar] {
            background-color: #E0E0E0; ##E5E6EA
        }
    </style>
    """,
    unsafe_allow_html=True,
)

with st.sidebar:
    # Add deeploy logo
    st.image("deeploy_logo.png", width=270)
    # 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 API token", "my-secret-token")
    if deployment_token == "my-secret-token":
        # show warning until token has been filled in
        st.warning("Please enter Deeploy API token.")
    else:
        st.button(
            "Get suspicious transaction",
            key="predict_button",
            help="Click to get a suspicious transaction",
            use_container_width=True,
            on_click=disabled,
            disabled=st.session_state.disabled,
        )
        ChangeButtonColour("Get suspicious transaction", "#FFFFFF", "#00052D")

    # define client options and instantiate client
    client_options = {
        "host": host,
        "deployment_token": deployment_token,
        "workspace_id": workspace_id,
    }
    client = Client(**client_options)

# instantiate session state vars to define whether predict button has been clicked
# and explanation was retrieved
if "predict_button" not in st.session_state:
    st.session_state.predict_button = False

if st.session_state.predict_button:
    st.session_state.predict_button_clicked = True

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

# make prediction and explanation calls and store important vars
if st.session_state.predict_button_clicked:
    try:
        with st.spinner("Loading..."):
            datapoint_pd = get_random_suspicious_transaction(data)
            request_body = get_request_body(datapoint_pd)
            # Call the explain endpoint as it also includes the prediction
            exp = client.explain(request_body=request_body, deployment_id=deployment_id)
            st.session_state.shap_values = exp["explanations"][0]["shap_values"]
            st.session_state.request_log_id = exp["requestLogId"]
            st.session_state.prediction_log_id = exp["predictionLogIds"][0]
            st.session_state.datapoint_pd = datapoint_pd
            certainty = get_fake_certainty()
            st.session_state.certainty = certainty
            st.session_state.got_explanation = True
            st.session_state.predict_button_clicked = False
    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."
        )

# create warning or info to be shown until prediction has been retrieved
if not st.session_state.got_explanation:
    st.info(
        "Fill in left hand side and click on button to observe a potential fraudulent transaction"
    )

# store important vars from result of prediction and explanation call
if st.session_state.got_explanation:
    shap_values = st.session_state.shap_values
    request_log_id = st.session_state.request_log_id
    prediction_log_id = st.session_state.prediction_log_id
    datapoint_pd = st.session_state.datapoint_pd
    certainty = st.session_state.certainty
    datapoint = modify_datapoint(datapoint_pd)

    # create two columns to show data input used and explanation
    col1, col2 = st.columns(2)

    # col1 contains input data table
    with col1:
        create_data_input_table(datapoint, COL_NAMES)

    # col 2 contains model certainty and explanation table of top 5 features
    with col2:
        st.subheader("AML Model Hit")
        st.metric(label="Model Certainty", value=certainty, delta="threshold: 75%")
        explainability_texts, sorted_indices = get_explainability_texts(
            shap_values, feature_texts
        )
        weights = get_weights(shap_values, sorted_indices)
        explainability_values = get_explainability_values(sorted_indices, datapoint)
        create_table(
            explainability_texts,
            explainability_values,
            weights,
            "Important Suspicious Factors",
        )

    st.subheader("")

    # add var to session state to discern if user has started an evaluation
    if "eval_selected" not in st.session_state:
        st.session_state["eval_selected"] = False

    # define two columns for agree and disagree button + text area for evaluation input
    col3, col4 = st.columns(2)

    # col 3 contains yes button
    with col3:
        # create empty state so that button disappears when st.empty is cleared
        eval1 = st.empty()
        eval1.button(
            "Send to FIU",
            key="yes_button",
            use_container_width=True,
            disabled=st.session_state.submitted_disabled,
        )
        ChangeButtonColour("Send to FIU", "#FFFFFF", "#4C506C")
        st.session_state.yes_button_clicked = False

    if st.session_state.yes_button:
        st.session_state.eval_selected = True
        st.session_state.evaluation_input = {"result": 0}  # Agree with the prediction

    # col 4 contains no button
    with col4:
        # create empty state so that button disappears when st.empty is cleared
        eval2 = st.empty()
        eval2.button(
            "Not money laundering",
            key="no_button",
            use_container_width=True,
            disabled=st.session_state.submitted_disabled,
        )
        ChangeButtonColour("Not money laundering", "#FFFFFF", "#4C506C")
        st.session_state.no_button_clicked = False
    if st.session_state.no_button:
        st.session_state.no_button_clicked = True
    if st.session_state.no_button_clicked:
        st.session_state.eval_selected = True
        st.session_state.evaluation_input = {
            "result": 1,  # Disagree with the prediction
            "value": {"predictions": [1]},
        }

    # define process for evaluation
    success = False
    if st.session_state.eval_selected:
        # if agree button clicked ("Send to FIU"), prefill explanation as comment for evaluation
        # change evaluation is user decides to fill in own text
        if st.session_state.yes_button:
            st.session_state.yes_button_clicked = True
            yes_button = True
            explanation = get_comment_explanation(
                certainty, explainability_texts, explainability_values
            )
            st.session_state.yes_button_text = explanation
            comment = st.text_area(
                "Reason for evaluation:",
                st.session_state.yes_button_text,
                key="yes_comment",
                on_change=get_input_yes_button,
            )
            st.session_state.evaluation_input[
                "explanation"
            ] = st.session_state.yes_button_text

        # if disagree button clicked ("Not money laundering") prefill with text that user
        # has to finish as a reason for evaluation
        if st.session_state.no_button:
            comment = st.text_area(
                "Reason for evaluation:",
                st.session_state.no_button_text,
                key="no_comment",
                on_change=get_input_no_button,
            )
            st.session_state.evaluation_input[
                "explanation"
            ] = st.session_state.no_button_text
        # create empty state so that button submit disappears when st.empty is cleared
        eval3 = st.empty()
        eval3.button(
            "Submit",
            key="submit_button",
            use_container_width=True,
            on_click=submitted_disabled,
            disabled=st.session_state.submitted_disabled,
        )
        ChangeButtonColour("Submit", "#FFFFFF", "#00052D")

        # if submit button is clicked, send evaluation to Deeploy
        if st.session_state.submit_button:
            st.session_state.eval_selected = False
            success = send_evaluation(
                client,
                deployment_id,
                request_log_id,
                prediction_log_id,
                st.session_state.evaluation_input,
            )
    # if the sending of evaluation was successful, remove buttons and enable Next button
    # to be clicked for next prediction and explanation to appear
    if success:
        st.session_state.eval_selected = False
        st.session_state.submitted = True
        eval1.empty()
        eval2.empty()
        eval3.empty()
        st.warning("Feedback submitted successfully")
        st.button("Next", key="next", use_container_width=True, on_click=rerun)
        ChangeButtonColour("Next", "#FFFFFF", "#00052D")