import pandas as pd
import plotly.express as px
import streamlit as st


def _viz_wealth(results):
    wealth = results["wealth"]
    concepts = results["concepts"]
    significance_level = results["significance_level"]

    wealth_mu = wealth.mean(axis=0)

    wealth_df = []
    for concept_idx, concept in enumerate(concepts):
        for t in range(wealth.shape[1]):
            wealth_df.append(
                {"time": t, "concept": concept, "wealth": wealth_mu[t, concept_idx]}
            )
    wealth_df = pd.DataFrame(wealth_df)

    fig = px.line(wealth_df, x="time", y="wealth", color="concept")
    fig.add_hline(
        y=1 / significance_level,
        line_dash="dash",
        line_color="black",
        annotation_text="Rejection threshold (1 / α)",
        annotation_position="bottom right",
    )
    fig.update_yaxes(range=[0, 1.5 * 1 / significance_level])
    st.plotly_chart(fig, use_container_width=True)


def viz_results():
    results = st.session_state.results

    if results is None:
        st.info("Test concepts to show results", icon="ℹ️")
    else:
        rank_tab, wealth_tab = st.tabs(["Rank of importance", "Wealth process"])

        with rank_tab:
            st.subheader("Rank of Semantic Importance")

            with st.expander("Details"):
                st.write(
                    "This tab shows the rank of semantic importance of the concepts for the predictions of the model on the image. Concepts are sorted by increasing rejection time, where a shorter rejection time indicates higher importance."
                )
        with wealth_tab:
            st.subheader("Wealth Process of Testing Procedures")

            with st.expander("Details"):
                st.write(
                    "This tab shows the average wealth process of the testing procedures for random draws of conditioning subsets with the same cardinality. The black dashed line represents the rejection threshold."
                )

            if results is not None:
                _viz_wealth(results)