import streamlit as st

st.set_page_config(layout="wide")

margin1, center, margin2 = st.columns([1,9,1])

with center:

    st.title("AI Model Fairness Auditing")

# --------------------------------------------------
# Strategy Definitions
# --------------------------------------------------

model_owner_actions = {
    "Overfit to One Audit Dataset": "overfit",
    "Overfit to Audit Dataset-like Distributions": "overfitdist",
}

auditor_actions = {
    "Audit with One Dataset": "auditone",
    "Audit with Multiple Datasets": "auditmany",
    "Use Private Auditing Datasets": "private",
    "Audit Model Internals": "internals",
}

# --------------------------------------------------
# Session State
# --------------------------------------------------

if "owner_selected" not in st.session_state:
    st.session_state.owner_selected = []

if "auditor_selected" not in st.session_state:
    st.session_state.auditor_selected = []

# --------------------------------------------------
# Toggle Functions
# --------------------------------------------------

def toggle_owner(action):
    if action in st.session_state.owner_selected:
        st.session_state.owner_selected.remove(action)
    else:
        st.session_state.owner_selected.append(action)

def toggle_auditor(action):
    if action in st.session_state.auditor_selected:
        st.session_state.auditor_selected.remove(action)
    else:
        st.session_state.auditor_selected.append(action)

# --------------------------------------------------
# Layout
# --------------------------------------------------

# left, right = st.columns(2)
margin1, left, gap, right, margin2 = st.columns([1,4,1,4,1])

# --------------------------------------------------
# Model Owner
# --------------------------------------------------

with left:

    st.subheader("Model Owner")

    st.caption("Chosen")

    with st.container(border=True):
        if not st.session_state.owner_selected:
            st.write("None")

        for action in st.session_state.owner_selected:
            st.button(
                action,
                key=f"owner_selected_{action}",
                on_click=toggle_owner,
                args=(action,),
                use_container_width=True,
                type="primary"
            )

    st.caption("Available")

    with st.container(border=True):
        for action in model_owner_actions:
            if action not in st.session_state.owner_selected:
                st.button(
                    action,
                    key=f"owner_available_{action}",
                    on_click=toggle_owner,
                    args=(action,),
                    use_container_width=True,
                    type="tertiary"
                )

# --------------------------------------------------
# Auditor
# --------------------------------------------------

with right:

    st.subheader("Auditor")

    st.caption("Chosen")

    with st.container(border=True):
        if not st.session_state.auditor_selected:
            st.write("None")

        for action in st.session_state.auditor_selected:
            st.button(
                action,
                key=f"auditor_selected_{action}",
                on_click=toggle_auditor,
                args=(action,),
                use_container_width=True,
                type="primary"
            )

    st.caption("Available")

    with st.container(border=True):
        for action in auditor_actions:
            if action not in st.session_state.auditor_selected:
                st.button(
                    action,
                    key=f"auditor_available_{action}",
                    on_click=toggle_auditor,
                    args=(action,),
                    use_container_width=True,
                    type="tertiary"
                )

# --------------------------------------------------
# Score Calculation
# --------------------------------------------------

auditor_active_list = []
for action in st.session_state.auditor_selected:
    auditor_active_list.append(auditor_actions[action])
owner_active_list = []
for action in st.session_state.owner_selected:
    owner_active_list.append(model_owner_actions[action])



auditor_cost = "🟢 Low"
if "auditmany" in auditor_active_list:
    auditor_cost = "🟡 Medium"
if "internals" in auditor_active_list:
    auditor_cost = "🔴 Very High"

owner_cost = "None"
if "overfit" in owner_active_list:
    owner_cost = "🟢 Low"
if "overfitdist" in owner_active_list:
    owner_cost = "🟡 Medium"

reliability = "❌ Low"
if "auditone" in auditor_active_list and len(owner_active_list)==0:
    reliability = "⚠️ Medium"
if "auditmany" in auditor_active_list and len(owner_active_list)==0:
    reliability = "✅ High"
if "auditmany" in auditor_active_list and "overfit" in owner_active_list:
    reliability = "⚠️ Medium"
if "private" in auditor_active_list and "overfitdist" not in owner_active_list:
    reliability = "⚠️ Medium"
if "private" in auditor_active_list and "auditmany" in auditor_active_list and "overfitdist" not in owner_active_list:
    reliability = "✅ High"
if "private" in auditor_active_list and "auditmany" in auditor_active_list and "overfitdist" in owner_active_list:
    reliability = "⚠️ Medium"
if "internals" in auditor_active_list:
    reliability = "✅ High"

openness = "None"
if "internals" in auditor_active_list:
    openness = "🔍 High (Model Weight Sharing)"

if len(auditor_active_list)==0 and len(owner_active_list)==0:
    auditor_cost, owner_cost, reliability, openness = "N/A", "N/A", "N/A", "N/A"

# --------------------------------------------------
# Dashboard
# --------------------------------------------------

st.divider()

margin1, c1, c2, c3, c4, margin2 = st.columns([1,2.25,2.25,2.25,2.25,1])

with c1:
    st.write("**Auditing Cost**")
    st.write(auditor_cost)

with c2:
    st.write("**Cost to Cheat for Model Owner**")
    st.write(owner_cost)

with c3:
    st.write("**Audit Reliability**")
    st.write(reliability)

with c4:
    st.write("**Openness Requirements**")
    st.write(openness)