import json
from pathlib import Path

import pandas as pd
import streamlit as st

st.set_page_config(page_title="BDO.ai Dashboard", layout="wide")


@st.cache_data
def load_traces():
    path = Path("artifacts/training_traces.json")
    if not path.exists():
        return []
    return json.loads(path.read_text())


@st.cache_data
def load_report():
    path = Path("artifacts/train_report.json")
    if not path.exists():
        return {}
    return json.loads(path.read_text())


def main():
    st.title("BDO.ai Dashboard")
    st.markdown("Monitor agent training traces and environment dynamics.")

    traces = load_traces()
    report = load_report()

    if not traces:
        st.warning("No training traces found. Please run `python train.py` first.")
        return

    # Overview
    st.header("Training Overview")
    if report:
        col1, col2, col3 = st.columns(3)
        col1.metric("Scenario", report.get("scenario", "N/A"))
        col2.metric("Best Total Training Reward", report.get("best_total_training_reward", "N/A"))
        col3.metric("Best Belief Accuracy", report.get("best_avg_belief_accuracy", "N/A"))

    # Episode Selection
    st.header("Episode Analysis")
    episodes = [t["episode"] for t in traces]
    selected_ep = st.selectbox("Select Episode to Analyze", episodes)
    
    trace = next((t for t in traces if t["episode"] == selected_ep), None)
    
    if trace:
        st.subheader(f"Episode {selected_ep} Summary")
        col1, col2, col3 = st.columns(3)
        col1.metric("Total Env Reward", trace["total_reward"])
        col2.metric("Total Training Reward", trace["total_training_reward"])
        col3.metric("Avg Consistency Bonus", trace["avg_consistency_bonus"])

        monthly = trace["monthly"]
        if monthly:
            df = pd.DataFrame(monthly)
            
            st.markdown("### Belief Accuracy (Rationality Verifier)")
            st.info("This metric shows how well the BDO models the hidden fraud levels based on noisy signals. A rising curve indicates true world modeling.")
            st.line_chart(df.set_index("month")[["belief_accuracy"]], color=["#17A2B8"])

            st.markdown("### Physical Reward Components")
            st.line_chart(df.set_index("month")[["env_reward", "training_reward"]])

            st.markdown("### Thought Process Log")
            for m in monthly:
                with st.expander(f"Month {m['month']} (Bonus: {m['consistency_bonus']})"):
                    st.write(m.get("thought_process", "No thought process logged."))

if __name__ == "__main__":
    main()