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

logging.basicConfig(level=logging.INFO)

st.set_page_config(layout="wide")

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

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

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

def disabled():
    st.session_state.disabled = True

def rerun():
    st.session_state.predict_button_clicked = True
    st.session_state.submitted_disabled = False

def submitted_disabled():
    st.session_state.submitted_disabled = True

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 Model Token", "my-secret-token") # my-secret-token
    if deployment_token == "my-secret-token":
        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
    ) #on_click=lambda: st.experimental_rerun() 
        ChangeButtonColour("Get suspicious transaction", '#FFFFFF', "#00052D")#'#FFFFFF', "#00052D"


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

if 'predict_button' not in st.session_state:
    st.session_state.predict_button = False

if st.session_state.predict_button: #  and not st.session_state.predict_button_clicked
    st.session_state.predict_button_clicked = True

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

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)
            # request_log_id = exp["requestLogId"]
            # prediction_log_id = exp["predictionLogIds"][0]
            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."
        )

if not st.session_state.got_explanation:
        st.info(
        "Fill in left hand side and click on button to observe a potential fraudulent transaction"
        )
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

    col1, col2 = st.columns(2)

    with col1:
        create_data_input_table(datapoint_pd, COL_NAMES)

    with col2:
        st.subheader('AML Model Hit')
        # st.success(f'{certainty}')
        # st.metric(label='Model Certainty', value=certainty)
        # style_metric_cards(border_left_color='#00052D', box_shadow=False)
        # # st.markdown('#### Model Certainty')
        st.metric(label='Model Certainty', value=certainty)
        

        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_pd)
        create_table(explainability_texts, explainability_values, weights, 'Important Suspicious Factors')

    st.subheader("")
    # st.markdown("<h2 style='text-align: center; white: red;'>Evaluation</h2>", unsafe_allow_html=True)

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


    col3, col4 = st.columns(2)
    with col3:
        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') #'#FFFFFF', "#DD360C"
        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
        }
        
    with col4:
        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') # '#FFFFFF', "#46B071",   '#FFFFFF', "#666666"
        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
        desired_output = 1
        st.session_state.evaluation_input = {
            "result": 1, # Disagree with the prediction
            "value": {"predictions": [desired_output]},
        }

    success = False
    if st.session_state.eval_selected:
        # st.write('after eval')
        # st.write(st.session_state)
        if st.session_state.yes_button:
            explanation = get_comment_explanation(certainty, explainability_texts, explainability_values)
            comment = st.text_area("Reason for evaluation:", explanation)
            st.session_state.evaluation_input["explanation"] = comment
        if st.session_state.no_button:
            comment = st.text_area("Reason for evaluation:", "I don't think this transaction is money laundering because...")
            st.session_state.evaluation_input["explanation"] = comment
        logging.debug("Selected feedback:" + str(st.session_state.evaluation_input))
        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 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 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") #'#FFFFFF', #F9B917" "#DD360C  #457EA4