"""
Streamlit App for AuditRepairEnv++ Demo
Interactive demonstration of the RL environment for ledger repair
"""

import streamlit as st
import numpy as np
from chronostasis import LedgerRepairEnv
import plotly.graph_objects as go

# Page config
st.set_page_config(
    page_title="AuditRepairEnv++",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Styling
st.markdown("""
<style>
    .stTabs [data-baseurlpath] {color: #667eea;}
    h1 {color: #667eea;}
    h2 {color: #764ba2;}
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'env' not in st.session_state:
    st.session_state.env = None
if 'episode_history' not in st.session_state:
    st.session_state.episode_history = []
if 'current_obs' not in st.session_state:
    st.session_state.current_obs = None
if 'current_info' not in st.session_state:
    st.session_state.current_info = None

# Header
col1, col2 = st.columns([4, 1])
with col1:
    st.title("🤖 AuditRepairEnv++")
    st.markdown("**RL Environment for Cost-Constrained Iterative Ledger Repair**")

with col2:
    st.metric("Version", "1.0.0")

st.markdown("""
Fix financial ledger errors while managing costs and avoiding cascading problems.
An interactive Reinforcement Learning environment for multi-step decision making.
""")

st.divider()

# Sidebar - Configuration
with st.sidebar:
    st.header("⚙️ Configuration")
    
    scenario = st.selectbox(
        "Choose Scenario",
        ["Easy", "Medium", "Hard"],
        help="Difficulty level affects complexity"
    )
    
    # Scenario presets
    scenarios = {
        "Easy": {
            "num_transactions": 15,
            "error_probability": 0.25,
            "budget": 250.0,
            "max_steps": 40
        },
        "Medium": {
            "num_transactions": 25,
            "error_probability": 0.35,
            "budget": 200.0,
            "max_steps": 50
        },
        "Hard": {
            "num_transactions": 40,
            "error_probability": 0.45,
            "budget": 150.0,
            "max_steps": 60
        }
    }
    
    config = scenarios[scenario]
    
    # Advanced options
    with st.expander("🔧 Advanced Settings"):
        config["num_transactions"] = st.slider(
            "Transactions", 5, 100, config["num_transactions"]
        )
        config["error_probability"] = st.slider(
            "Error Probability", 0.0, 1.0, config["error_probability"], 0.05
        )
        config["budget"] = st.slider(
            "Budget ($)", 50.0, 500.0, config["budget"], 10.0
        )
        config["max_steps"] = st.slider(
            "Max Steps", 10, 200, config["max_steps"], 10
        )
    
    st.divider()
    st.subheader("📖 Help")
    st.markdown("""
    ### Actions
    - **Fix (0)**: Repair an error • Cost: $10
    - **Revert (1)**: Undo last action • Cost: $5
    - **Skip (2)**: Do nothing • Cost: $0
    
    ### Goal
    Achieve 100% consistency while staying under budget.
    """)

# Main content - Tabs
tab1, tab2, tab3, tab4 = st.tabs(
    ["🎮 Play", "📊 Metrics", "📋 Details", "ℹ️ About"]
)

with tab1:
    st.header("Play the Game")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        if st.button("🔄 Reset Environment", key="reset_btn", use_container_width=True):
            st.session_state.env = LedgerRepairEnv(**config)
            obs, info = st.session_state.env.reset()
            st.session_state.current_obs = obs
            st.session_state.current_info = info
            st.session_state.episode_history = [{
                "step": 0,
                "action": "RESET",
                "reward": 0.0,
                "cost": 0.0,
                "errors": info['num_errors'],
                "consistency": 0.0
            }]
            st.success("✅ Environment reset!")
            st.rerun()
    
    with col2:
        st.write("")  # Spacer
    
    with col3:
        st.write("")  # Spacer
    
    if st.session_state.env is None:
        st.info("👈 Click 'Reset Environment' to start")
    else:
        env = st.session_state.env
        obs = st.session_state.current_obs
        info = st.session_state.current_info
        
        # Current state display
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric(
                "Errors Remaining",
                info['num_errors'],
                f"-{env.initial_error_count - info['num_errors']} {env.initial_error_count - info['num_errors'] != 1 and 's' or ''}",
                delta_color="inverse"
            )
        
        with col2:
            st.metric(
                "Budget Remaining",
                f"${info['budget_remaining']:.2f}",
                f"spent: ${info['total_cost']:.2f}",
            )
        
        with col3:
            consistency = (env.initial_error_count - info['num_errors']) / max(env.initial_error_count, 1) * 100
            st.metric("Consistency", f"{consistency:.1f}%")
        
        with col4:
            st.metric("Step", info['step'])
        
        st.divider()
        
        # Action buttons
        st.subheader("Choose Action:")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            if st.button("🔧 Fix Entry (Cost: $10)", use_container_width=True, key="fix"):
                obs, reward, terminated, truncated, info = env.step(0)
                st.session_state.current_obs = obs
                st.session_state.current_info = info
                
                st.session_state.episode_history.append({
                    "step": info['step'],
                    "action": "FIX",
                    "reward": reward,
                    "cost": info['total_cost'],
                    "errors": info['num_errors'],
                    "consistency": env.ledger.consistency_ratio() * 100
                })
                
                if terminated:
                    st.balloons()
                    st.success(f"🎉 Episode Complete! Final Score: {sum([h['reward'] for h in st.session_state.episode_history]):.2f}")
                
                if truncated:
                    st.warning("⏱️ Max steps reached!")
                
                st.rerun()
        
        with col2:
            if st.button("↩️ Revert (Cost: $5)", use_container_width=True, key="revert"):
                obs, reward, terminated, truncated, info = env.step(1)
                st.session_state.current_obs = obs
                st.session_state.current_info = info
                
                st.session_state.episode_history.append({
                    "step": info['step'],
                    "action": "REVERT",
                    "reward": reward,
                    "cost": info['total_cost'],
                    "errors": info['num_errors'],
                    "consistency": env.ledger.consistency_ratio() * 100
                })
                
                st.rerun()
        
        with col3:
            if st.button("⏯️ Skip (Cost: $0)", use_container_width=True, key="skip"):
                obs, reward, terminated, truncated, info = env.step(2)
                st.session_state.current_obs = obs
                st.session_state.current_info = info
                
                st.session_state.episode_history.append({
                    "step": info['step'],
                    "action": "SKIP",
                    "reward": reward,
                    "cost": info['total_cost'],
                    "errors": info['num_errors'],
                    "consistency": env.ledger.consistency_ratio() * 100
                })
                
                st.rerun()
        
        st.divider()
        
        # Remaining errors display
        if info['num_errors'] > 0:
            st.subheader("⚠️ Remaining Errors:")
            error_list = list(env.ledger.errors.items())[:5]
            for entry_id, error_desc in error_list:
                st.warning(f"**Entry {entry_id}:** {error_desc}")
            if len(env.ledger.errors) > 5:
                st.info(f"... and {len(env.ledger.errors) - 5} more errors")
        else:
            st.success("✅ All errors fixed!")

with tab2:
    st.header("📊 Episode Metrics")
    
    if not st.session_state.episode_history or len(st.session_state.episode_history) <= 1:
        st.info("👈 Play the game to see metrics")
    else:
        history = st.session_state.episode_history[1:]  # Skip reset
        
        # Charts
        col1, col2 = st.columns(2)
        
        with col1:
            # Cumulative reward
            steps = [h['step'] for h in history]
            cumulative_rewards = np.cumsum([h['reward'] for h in history])
            
            fig = go.Figure()
            fig.add_trace(go.Scatter(
                x=steps, y=cumulative_rewards,
                mode='lines+markers',
                name='Cumulative Reward',
                line=dict(color='#667eea', width=2),
                fill='tozeroy'
            ))
            fig.update_layout(
                title="Cumulative Reward",
                xaxis_title="Step",
                yaxis_title="Reward",
                height=400,
                template="plotly_white"
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            # Cost and consistency
            costs = [h['cost'] for h in history]
            consistency = [h['consistency'] for h in history]
            
            fig = go.Figure()
            fig.add_trace(go.Scatter(
                x=steps, y=costs,
                mode='lines+markers',
                name='Total Cost',
                line=dict(color='#ef4444', width=2),
                yaxis='y'
            ))
            fig.add_trace(go.Scatter(
                x=steps, y=consistency,
                mode='lines+markers',
                name='Consistency %',
                line=dict(color='#10b981', width=2),
                yaxis='y2'
            ))
            fig.update_layout(
                title="Cost vs Consistency",
                xaxis_title="Step",
                yaxis=dict(title="Cost ($)", side='left'),
                yaxis2=dict(title="Consistency (%)", side='right', overlaying='y'),
                height=400,
                template="plotly_white",
                hovermode='x unified'
            )
            st.plotly_chart(fig, use_container_width=True)
        
        # Statistics
        st.divider()
        st.subheader("📈 Statistics")
        
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric("Total Steps", len(history))
        
        with col2:
            total_reward = sum([h['reward'] for h in history])
            st.metric("Total Reward", f"{total_reward:.2f}")
        
        with col3:
            final_cost = history[-1]['cost']
            st.metric("Final Cost", f"${final_cost:.2f}")
        
        with col4:
            final_consistency = history[-1]['consistency']
            st.metric("Final Consistency", f"{final_consistency:.1f}%")

with tab3:
    st.header("📋 Episode History")
    
    if not st.session_state.episode_history or len(st.session_state.episode_history) <= 1:
        st.info("👈 Play the game to see history")
    else:
        import pandas as pd
        
        history_df = pd.DataFrame(st.session_state.episode_history[1:])
        st.dataframe(
            history_df,
            use_container_width=True,
            hide_index=True,
            column_config={
                "step": st.column_config.NumberColumn("Step", format="%d"),
                "action": st.column_config.TextColumn("Action"),
                "reward": st.column_config.NumberColumn("Reward", format="%.2f"),
                "cost": st.column_config.NumberColumn("Cost", format="$%.2f"),
                "errors": st.column_config.NumberColumn("Errors", format="%d"),
                "consistency": st.column_config.NumberColumn("Consistency", format="%.1f%%"),
            }
        )

with tab4:
    st.header("ℹ️ About AuditRepairEnv++")
    
    st.markdown("""
    ### 🎯 What is This?
    
    AuditRepairEnv++ is an OpenAI Gymnasium-compatible RL environment where agents must 
    iteratively repair inconsistencies in a financial ledger while:
    
    - **Managing Costs**: Each action has a monetary cost
    - **Avoiding Cascade Errors**: Fixing one error can introduce new errors
    - **Meeting Constraints**: Stay within a budget while maximizing consistency
    
    ### 🤖 Real-World Applications
    
    - Financial reconciliation systems
    - Audit ledger repair
    - Transaction correction in payment systems
    - Data cleaning and consistency checking
    
    ### 📊 Environment Metrics
    
    Your performance is evaluated on:
    
    1. **Consistency (40%)**: How many errors you fix
    2. **Cost Efficiency (35%)**: How well you stay under budget
    3. **Action Efficiency (15%)**: How few actions you take
    4. **Stability (10%)**: How few overcorrections you make
    
    ### 🚀 Try Different Scenarios
    
    - **Easy**: Simple ledgers with fewer errors
    - **Medium**: Complex patterns with cascading effects
    - **Hard**: Large-scale problems with hidden dependencies
    
    ### 📚 Learn More
    
    - [GitHub Repository](https://github.com/your-repo/auditrepairenv-plus)
    - [OpenAPI Docs](/docs)
    - [Gymnasium Framework](https://gymnasium.farama.org/)
    
    ### 💡 Tips for Success
    
    1. Start with **Easy** difficulty
    2. Watch for **cascading errors** (fixing one can break another)
    3. Balance **speed** with **cost**
    4. Use **Revert** strategically when mistakes happen
    """)
    
    st.divider()
    st.markdown("**Built with ❤️ for the AI community** | v1.0.0")