import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import json
import time
from datetime import datetime

# Import from your fixed procurement agent file
from agentic_sourcing_ppo_sap_colab import (
    suppliers_synthetic, market_signal, rl_recommend_tool, 
    sap_create_po_mock, check_model_tool
)

# Page config
st.set_page_config(
    page_title="🤖 AI Procurement Agent Demo",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        font-weight: bold;
        color: #2E86AB;
        text-align: center;
        margin-bottom: 2rem;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #F24236;
        margin-bottom: 1rem;
    }
    .metric-container {
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        padding: 1rem;
        border-radius: 10px;
        margin: 0.5rem 0;
    }
    .success-box {
        background: #d4edda;
        border: 1px solid #c3e6cb;
        padding: 1rem;
        border-radius: 5px;
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

def create_allocation_pie_chart(allocations):
    """Create pie chart for supplier allocations"""
    df = pd.DataFrame(allocations)
    df = df[df['share'] > 0.01]  # Filter out very small allocations
    
    fig = px.pie(df, values='share', names='supplier', 
                 title="Supplier Allocation Distribution",
                 color_discrete_sequence=px.colors.qualitative.Set3)
    fig.update_traces(textposition='inside', textinfo='percent+label')
    fig.update_layout(height=400)
    return fig

def create_supplier_comparison_chart(suppliers_data):
    """Create radar chart comparing suppliers"""
    df = pd.DataFrame(suppliers_data)
    
    # Select top 5 suppliers by quality score
    df['combined_score'] = df['current_quality'] * 0.4 + df['current_delivery'] * 0.3 + (1-df['financial_risk']) * 0.3
    top_suppliers = df.nlargest(5, 'combined_score')
    
    categories = ['Quality', 'Delivery', 'ESG Score', 'Low Risk', 'Cost Efficiency']
    
    fig = go.Figure()
    
    for _, supplier in top_suppliers.iterrows():
        values = [
            supplier['current_quality'],
            supplier['current_delivery'],
            supplier['esg'],
            1 - supplier['financial_risk'],  # Invert risk for better visualization
            1 - (supplier['base_cost_per_unit'] / 150)  # Normalize cost
        ]
        
        fig.add_trace(go.Scatterpolar(
            r=values,
            theta=categories,
            fill='toself',
            name=supplier['name'],
            opacity=0.7
        ))
    
    fig.update_layout(
        polar=dict(
            radialaxis=dict(visible=True, range=[0, 1])
        ),
        showlegend=True,
        title="Top 5 Suppliers Comparison",
        height=500
    )
    return fig

def main():
    # Header
    st.markdown('<div class="main-header">🤖 AI Procurement Agent Demo</div>', unsafe_allow_html=True)
    st.markdown("### Intelligent Supplier Selection using Reinforcement Learning")
    
    # Create columns for better layout
    col1, col2 = st.columns([1, 2])
    
    with col1:
        st.markdown('<div class="sub-header">🎛️ Control Panel</div>', unsafe_allow_html=True)
        
        # Market Parameters
        st.subheader("Market Conditions")
        volatility = st.selectbox(
            "Market Volatility",
            ["low", "medium", "high"],
            index=1,
            help="Current market volatility level"
        )
        
        demand_mult = st.slider(
            "Demand Multiplier",
            min_value=0.7,
            max_value=1.5,
            value=1.0,
            step=0.05,
            help="Demand change from baseline"
        )
        
        price_mult = st.slider(
            "Price Multiplier",
            min_value=0.8,
            max_value=1.3,
            value=1.0,
            step=0.05,
            help="Price change from baseline"
        )
        
        baseline_demand = st.number_input(
            "Baseline Demand (units)",
            min_value=100,
            max_value=10000,
            value=1000,
            step=100
        )
        
        # Supplier Configuration
        st.subheader("Supplier Configuration")
        num_suppliers = st.slider(
            "Number of Suppliers",
            min_value=3,
            max_value=10,
            value=6,
            help="Number of suppliers to consider"
        )
        
        seed = st.number_input(
            "Random Seed",
            min_value=1,
            max_value=1000,
            value=123,
            help="Seed for reproducible supplier generation"
        )
        
    with col2:
        st.markdown('<div class="sub-header">📊 Real-time Dashboard</div>', unsafe_allow_html=True)
        
        # Action button
        if st.button("🚀 Run Procurement Agent", type="primary", use_container_width=True):
            
            # Progress bar
            progress_bar = st.progress(0)
            status_text = st.empty()
            
            try:
                # Step 1: Generate suppliers
                status_text.text("Step 1/5: Generating supplier data...")
                progress_bar.progress(20)
                
                suppliers_result = suppliers_synthetic(n=num_suppliers, seed=seed)
                suppliers_data = suppliers_result["suppliers"]
                
                # Display suppliers table
                st.subheader("Generated Suppliers")
                df_suppliers = pd.DataFrame(suppliers_data)
                st.dataframe(df_suppliers.round(3), use_container_width=True)
                
                # Step 2: Market signals
                status_text.text("Step 2/5: Analyzing market conditions...")
                progress_bar.progress(40)
                
                market_data = market_signal(volatility, price_mult, demand_mult)
                
                # Display market metrics
                col_m1, col_m2, col_m3 = st.columns(3)
                with col_m1:
                    st.metric("Volatility", volatility.upper(), 
                             delta="High Risk" if volatility == "high" else "Normal")
                with col_m2:
                    st.metric("Demand Change", f"{demand_mult:.1%}", 
                             delta=f"{(demand_mult-1)*100:+.1f}%")
                with col_m3:
                    st.metric("Price Change", f"{price_mult:.1%}", 
                             delta=f"{(price_mult-1)*100:+.1f}%")
                
                # Step 3: Check model
                status_text.text("Step 3/5: Checking AI model availability...")
                progress_bar.progress(60)
                
                model_check = check_model_tool("./supplier_selection_ppo_gymnasium.pkl")
                
                # Step 4: Get recommendations
                status_text.text("Step 4/5: Getting AI recommendations...")
                progress_bar.progress(80)
                
                recommendation_input = {
                    "volatility": market_data["volatility"],
                    "price_multiplier": market_data["price_multiplier"],
                    "demand_multiplier": market_data["demand_multiplier"],
                    "baseline_demand": baseline_demand,
                    "suppliers": suppliers_data,
                    "auto_align_actions": True
                }
                
                recommendations = rl_recommend_tool(recommendation_input)
                
                if recommendations.get("strategy") == "error":
                    st.error(f"AI recommendation failed: {recommendations.get('error', 'Unknown error')}")
                    return
                
                # Step 5: Create PO
                status_text.text("Step 5/5: Creating purchase order...")
                progress_bar.progress(100)
                
                po_data = {
                    "lines": [
                        {
                            "supplier": alloc["supplier"],
                            "quantity": round(recommendations["demand_units"] * alloc["share"], 2)
                        }
                        for alloc in recommendations["allocations"]
                        if alloc["share"] > 0.01
                    ]
                }
                
                po_result = sap_create_po_mock(po_data)
                
                # Clear progress indicators
                status_text.text("✅ Procurement process completed!")
                time.sleep(0.5)
                progress_bar.empty()
                status_text.empty()
                
                # Display results
                st.markdown("---")
                st.subheader("🎯 Procurement Results")
                
                # Key metrics
                col_r1, col_r2, col_r3, col_r4 = st.columns(4)
                with col_r1:
                    st.metric("Strategy", recommendations["strategy"].title())
                with col_r2:
                    active_suppliers = len([a for a in recommendations["allocations"] if a["share"] > 0.01])
                    st.metric("Active Suppliers", active_suppliers)
                with col_r3:
                    st.metric("Total Units", f"{recommendations['demand_units']:,.0f}")
                with col_r4:
                    st.metric("PO Number", po_result["PurchaseOrder"])
                
                # Visualizations
                col_v1, col_v2 = st.columns(2)
                
                with col_v1:
                    # Allocation pie chart
                    fig_pie = create_allocation_pie_chart(recommendations["allocations"])
                    st.plotly_chart(fig_pie, use_container_width=True)
                
                with col_v2:
                    # Supplier comparison radar
                    fig_radar = create_supplier_comparison_chart(suppliers_data)
                    st.plotly_chart(fig_radar, use_container_width=True)
                
                # Detailed allocation table
                st.subheader("📋 Detailed Allocation")
                allocation_df = pd.DataFrame(recommendations["allocations"])
                allocation_df["quantity"] = allocation_df["share"] * recommendations["demand_units"]
                allocation_df["percentage"] = allocation_df["share"] * 100
                
                # Merge with supplier data for additional context
                supplier_df = pd.DataFrame(suppliers_data)
                detailed_df = allocation_df.merge(
                    supplier_df[["name", "base_cost_per_unit", "current_quality", "financial_risk"]], 
                    left_on="supplier", right_on="name"
                )
                
                st.dataframe(
                    detailed_df[["supplier", "percentage", "quantity", "base_cost_per_unit", "current_quality", "financial_risk"]]
                    .round(2), 
                    use_container_width=True
                )
                
                # Purchase Order JSON
                with st.expander("📄 View Purchase Order JSON"):
                    st.json(po_result)
                
                # Success message
                st.markdown(f"""
                <div class="success-box">
                    <strong>✅ Success!</strong> Purchase Order {po_result["PurchaseOrder"]} has been created successfully!
                    <br><em>Note: This is a demonstration. No actual SAP system was contacted.</em>
                </div>
                """, unsafe_allow_html=True)
                
            except Exception as e:
                st.error(f"Error during execution: {str(e)}")
                st.exception(e)

    # Sidebar with information
    with st.sidebar:
        st.markdown("### About This Demo")
        st.info("""
        This demo showcases an AI-powered procurement agent that:
        
        🎯 **Analyzes** market conditions and supplier data
        
        🤖 **Uses** reinforcement learning (PPO) for optimal allocation
        
        📊 **Generates** purchase orders automatically
        
        🔗 **Integrates** with SAP systems (mocked for demo)
        """)
        
        st.markdown("### Key Features")
        st.markdown("""
        - **Real-time Analysis**: Dynamic market condition assessment
        - **Multi-criteria Optimization**: Quality, cost, delivery, ESG factors
        - **Risk Management**: Financial and supply chain risk evaluation
        - **Scalable Architecture**: Handles multiple suppliers efficiently
        """)
        
        st.markdown("### Technology Stack")
        st.markdown("""
        - **RL Framework**: Stable-Baselines3 PPO
        - **Agent Framework**: SmolagentS
        - **Backend**: Python, NumPy, Pandas
        - **Frontend**: Streamlit, Plotly
        """)

if __name__ == "__main__":
    main()