import streamlit as st
import pandas as pd
import json
import os
from st_aggrid import AgGrid, GridOptionsBuilder, GridUpdateMode
import plotly.graph_objects as go

# -----------------------------
# Page Configuration
# -----------------------------
st.set_page_config(
    page_title="Golden Bank Kenya - Risk Dashboard",
    layout="wide",
    initial_sidebar_state="expanded"
)

with open("assets/styles.css", encoding="utf-8") as f:
    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

col_title, col_logo = st.columns([4, 1])
with col_title: 
    st.markdown("""
    <style>
        .stImage > img {
            width: 200px; /* Increase size */
            float: right; /* Align to the right */
            border-radius: 10px;
            box-shadow: 5px 5px 10px rgba(0, 0, 0, 0.2);
            transition: transform 0.3s ease-in-out;
        }
    </style>
    """, unsafe_allow_html=True)
    st.image("assets/bank-logo.png", width=90)

with col_logo:
    st.markdown("""
    <style>
        .stImage > img {
            width: 200px; /* Increase size */
            float: right; /* Align to the right */
            border-radius: 10px;
            box-shadow: 5px 5px 10px rgba(0, 0, 0, 0.2);
            transition: transform 0.3s ease-in-out;
        }
    </style>
    """, unsafe_allow_html=True)
    st.image("assets/ailabs_logo.png", width=250)

col_title, col_loan_line = st.columns([4, 1])
with col_title:
    st.markdown(
        "<h1 style='text-align: left; color: #ffffff;'>Early Detection Dashboard for Credit Risk</h1>",
        unsafe_allow_html=True
    )

with col_loan_line:
    st.markdown("""
    <style>
        .stImage > img {
            width: 200px; /* Increase size */
            float: right; /* Align to the right */
            border-radius: 10px;
            box-shadow: 5px 5px 10px rgba(0, 0, 0, 0.2);
            transition: transform 0.3s ease-in-out;
        }
    </style>
    """, unsafe_allow_html=True)
    st.image("assets/loan_line.png", width=350)

st.markdown("---")

# -----------------------------
# Load Data
# -----------------------------
def load_data():
    df = pd.read_csv("data/kenya_personal_loan_data_with_shap.csv")
    df["Top Contributors"] = df["Top Contributors"].apply(json.loads)
    return df

df = load_data()

# -----------------------------
# Sidebar Filters
# -----------------------------
st.sidebar.header("Filter Customers")

target_date = "2026-01-31"
last_month_df = df[df['Date'] == target_date]
filtered_df = last_month_df.copy()

# Filter options
risk_options = st.sidebar.multiselect(
    "Risk Category", options=df["Risk Category"].unique(), default=list(df["Risk Category"].unique())
)
filtered_df = filtered_df[filtered_df["Risk Category"].isin(risk_options)]

min_score, max_score = st.sidebar.slider(
    "Risk Score", 0, 999, (df["Risk Score"].min(), df["Risk Score"].max())
)
filtered_df = filtered_df[
    (filtered_df["Risk Score"] >= min_score) & (filtered_df["Risk Score"] <= max_score)
]

min_loan, max_loan = st.sidebar.slider(
    "Initial Loan Amount (USD)", 0, int(df["Initial Loan Amount"].max()), 
    (int(df["Initial Loan Amount"].min()), int(df["Initial Loan Amount"].max()))
)
filtered_df = filtered_df[
    (filtered_df["Initial Loan Amount"] >= min_loan) & (filtered_df["Initial Loan Amount"] <= max_loan)
]

min_balance, max_balance = st.sidebar.slider(
    "Outstanding Balance (USD)", 0, int(df["Outstanding Balance"].max()), 
    (int(df["Outstanding Balance"].min()), int(df["Outstanding Balance"].max()))
)
filtered_df = filtered_df[
    (filtered_df["Outstanding Balance"] >= min_balance) & (filtered_df["Outstanding Balance"] <= max_balance)
]

delinquency_bucket_options = st.sidebar.multiselect(
    "Delinquency Bucket", options=df["Delinquency Bucket"].unique(), default=list(df["Delinquency Bucket"].unique())
)
filtered_df = filtered_df[filtered_df["Delinquency Bucket"].isin(delinquency_bucket_options)]

delinquency_status_options = st.sidebar.multiselect(
    "Delinquency Status", options=df["Delinquency Status"].unique(), default=list(df['Delinquency Status'].unique())
)
filtered_df = filtered_df[filtered_df["Delinquency Status"].isin(delinquency_status_options)]


# -----------------------------
# Pagination Controls
# -----------------------------
PAGE_SIZE = 20
total_records = len(filtered_df)
total_pages = (total_records - 1) // PAGE_SIZE + 1 if total_records > 0 else 1

current_page = st.sidebar.number_input(
    "Page Number",
    min_value=1,
    max_value=total_pages,
    value=1,
    step=1
)

start_idx = (current_page - 1) * PAGE_SIZE
end_idx = start_idx + PAGE_SIZE
paginated_df = filtered_df.iloc[start_idx:end_idx]

st.sidebar.write(f"Showing records {start_idx + 1} to {min(end_idx, total_records)} of {total_records}")

# -----------------------------
# Main Panel Layout
# -----------------------------
col1, col2 = st.columns([4,1])

with col1:
    st.markdown("### Customer Risk Overview")

    col3, col4 = st.columns([11, 1])
    with col3:
        st.markdown(f'##### Prediction Month: :violet[2026-02-28]')

        display_df = paginated_df[[
            "Account ID", "Risk Score", "Risk Category", "Initial Loan Amount", 
            "Installment Amount", "Outstanding Balance", "Monthly Payments", 
            "Payment Behaviour", "Delinquency Bucket", "Delinquency Status", 
            "Recommended Risk Action"
        ]]

        st.caption(f"Page {current_page} of {total_pages} — Showing {len(paginated_df)} accounts")

    with col4:
        os.makedirs("outputs", exist_ok=True)
        download_path = f"outputs/ABC_Bank_Kenya_Customer_Portfolio_{target_date}_filtered_data.csv"
        filtered_df.to_csv(download_path, index=False)

        st.download_button(
            label="**Download**",
            data=open(download_path, "rb"),
            file_name=download_path,
            mime="text/csv",
            use_container_width=False,
            type='primary'
        )

    gb = GridOptionsBuilder.from_dataframe(display_df)
    gb.configure_selection(selection_mode="single", use_checkbox=False)
    grid_options = gb.build()

    # Defining the columns to format with two decimal places
    decimal_columns = [
        "Initial Loan Amount", 
        "Installment Amount", 
        "Outstanding Balance", 
        "Monthly Payments"
    ]

    for col in grid_options["columnDefs"]:
        col["cellStyle"] = {"textAlign": "center", "color": "white", "backgroundColor": "black"}

        # Apply decimal formatting
        if col["field"] in decimal_columns:
            col["valueFormatter"] = "x.toFixed(2)"

        # Apply color rules for Risk Category
        if col["field"] == "Risk Category":
            col["cellClassRules"] = {
                "high-risk": "x == 'High Risk'",
                "medium-risk": "x == 'Medium Risk'",
                "low-risk": "x == 'Low Risk'"
            }

    GRID_HEIGHT = 700  # Adjust this to ensure ~20 rows are visible without scroll
    response = AgGrid(
        display_df,
        gridOptions=grid_options,
        update_mode=GridUpdateMode.MODEL_CHANGED,
        height=GRID_HEIGHT,
        theme="material",
        custom_css={
            ".ag-theme-material": {
                "background-color": "material !important",
                "color": "white !important",

            },
            ".ag-cell": {
                "background-color": "#191414 !important",
                "color": "white !important",
            },
            ".ag-header": {
                "background-color": "#1e1e1e !important",
                "color": "white !important",
            },
            ".ag-header-cell-label": {
                "font-weight": "bold",
                "color": "white !important",
                "justify-content": "center",
            },
            ".high-risk": {
                "background-color": "#BD100D !important",
                "color": "white !important"
            },
            ".medium-risk": {
                "background-color": "#B88F12 !important",
                "color": "white !important"
            },
            ".low-risk": {
                "background-color": "#3A800F !important",
                "color": "white !important"
            },
            ".ag-cell[col-id='Risk Category']": {
                "font-weight": "bold !important"
            },
        }
    )

with col2:
    st.markdown("### Customer Details")
    selected_row = response["selected_rows"]

    if isinstance(selected_row, pd.DataFrame):
        selected_row = selected_row.to_dict("records")

    if isinstance(selected_row, list) and len(selected_row) > 0:
        account_id = selected_row[0]["Account ID"]
        record = filtered_df[filtered_df["Account ID"] == account_id].iloc[0]

        tab1, tab2, tab3 = st.tabs(["General Info", "Financial Info", "Top Risk Contributors"])

        with tab1:
            st.markdown("#### General Info")
            st.write(f"**Account ID:** {record['Account ID']}")
            st.write(f"**Risk Score:** {record['Risk Score']}")
            st.write(f"**Risk Category:** {record['Risk Category']}")
            st.write(f"**Recommended Risk Action:** {record['Recommended Risk Action']}")

            account_history = df[df['Account ID'] == account_id].sort_values(by="Date")

            if not account_history.empty:
                fig = go.Figure()
                fig.add_trace(go.Scatter(
                    x=account_history['Date'],
                    y=account_history['Risk Score'],
                    mode='lines+markers',
                    line=dict(color='#8D1DB9', width=2),
                    marker=dict(color="#B808FD", size=5),
                    name='Risk Score'
                ))
                fig.update_layout(
                    title='Risk Score Trend',
                    xaxis_title='Date',
                    yaxis_title='Risk Score',
                    template='plotly_white',
                    title_x=0.4
                )
                st.plotly_chart(fig, use_container_width=True)
            else:
                st.write("No historical risk score data available for this account.")

        with tab2:
            st.markdown("#### Financial Info")
            st.write(f"**Initial Loan Amount:** USD {record['Initial Loan Amount']:,.2f}")
            st.write(f"**Installment Amount:** USD {record['Installment Amount']:,.2f}")
            st.write(f"**Outstanding Balance:** USD {record['Outstanding Balance']:,.2f}")
            st.write(f"**Payments for the Month:** USD {record['Monthly Payments']:,.2f}")
            st.write(f"**Interest Rate:** {record['Interest rate (%)']:.2f}%")
            st.write(f"**Payment Behaviour:** {record['Payment Behaviour']}")
            st.write(f"**Delinquency Bucket:** {record['Delinquency Bucket']}")
            st.write(f"**Delinquency Status:** {record['Delinquency Status']}")

        with tab3:
            st.markdown("#### Key Features Driving the Risk Score")
            shap_plot_path = f"SHAP_plots/SHAP_Plot_{account_id}.jpg"
            if os.path.exists(shap_plot_path):
                st.image(shap_plot_path, caption=f"SHAP Plot for Account ID: {account_id}")
            else:
                st.write("Top Contributors plot not available for this account.")

st.caption("ⓘ The recommended risk action is derived based on the risk score of the customer in agreement with the bank's risk management policies.")

st.header("Behavioural Trend Analysis")

unique_accounts = df['Account ID'].unique()
selected_account = st.selectbox("Select Account ID", sorted(unique_accounts)) # dropdown to select account ID

# Filter the data for the selected account
account_df = df[df['Account ID'] == selected_account].sort_values(by="Date")

# Plot the trend of the risk score over time for a selected account
risk_score_trend = go.Figure()

risk_score_trend.add_trace(go.Scatter(
    x=account_df["Date"],
    y=account_df["Risk Score"],
    mode='lines+markers',
    name='Risk Score',
    line=dict(color="#0DABBD", width=3),
    marker=dict(size=10, color='cyan', line=dict(width=1, color='white')),
    hovertemplate='<b>Date:</b> %{x|%b %Y}<br><b>Risk Score:</b> %{y}<extra></extra>'
))

risk_score_trend.update_layout(
    title={
        'text': f"Risk Score Trend for Account ID: {selected_account}",
        'x': 0.5,
        'xanchor': 'center',
        'font': dict(size=22, family='Helvetica Neue, sans-serif')
    },
    xaxis_title='Date',
    yaxis_title='Risk Score',
    template='plotly_white',
    margin=dict(t=80, b=50, l=50, r=50),
    height=450
)
st.plotly_chart(risk_score_trend, use_container_width=True)

# Plot the trend of payments over time for the selected account
payments_trend = go.Figure()

payments_trend.add_trace(go.Scatter(
    x=account_df["Date"],
    y=account_df['Monthly Payments'],
    mode='lines+markers',
    line=dict(color='purple', width=3),
    marker=dict(size=10, color='magenta', line=dict(width=1, color='white')),
    name='Payments',
    hovertemplate='<b>Date:</b> %{x|%b %Y}<br><b>Payments:</b> %{y:,.2f}<extra></extra>'
))

payments_trend.update_layout(
    title={
        'text': f'Payment Trend for Account ID: {selected_account}',
        'x': 0.5,
        'xanchor': 'center',
        'font': dict(size=22, family='Helvetica Neue, sans-serif')
    },
    xaxis_title='Date',
    yaxis_title='Payments (USD)',
    template='plotly_white',
    margin=dict(t=60, b=60, l=50, r=50),
    height=450
)
st.plotly_chart(payments_trend, use_container_width=True)

st.markdown("---")
# st.caption("Powered by :violet[**IronOne AI Labs**]")