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

from utils import generate_recommendation
from config import DEFAULT_BUDGET,ALL_STOCKS

# Page Config
st.set_page_config(
    page_title="Dynamic Allocation System",
    page_icon="📈",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom Styling
st.markdown("""
<style>
.main {
    background-color:#0e1117;
}
.block-container {
    padding-top:2rem;
}
div[data-testid="metric-container"] {
    background:#1b1f2a;
    border-radius:15px;
    padding:15px;
    border:1px solid #2d3748;
}
.stButton>button {
    width:100%;
    border-radius:10px;
    height:55px;
    font-size:18px;
    font-weight:bold;
}
</style>
""", unsafe_allow_html=True)

# Main Title & Status Badges
st.title("📈 Dynamic Allocation System")

badge_col1, badge_col2 = st.columns([1, 5])
with badge_col1:
    st.success("🟢 Model Status: Loaded")
with badge_col2:
    st.info("🤖 Engine: PPO Reinforcement Learning")

st.caption("Machine Learning Based Dynamic Allocation System using Proximal Policy Optimization (PPO)")


# Clean fallback safety line

# Sidebar Setup
st.sidebar.image("Assets/logo.png", use_container_width=True)
st.sidebar.header("Investment Settings")

budget = st.sidebar.number_input(
    "Investment Budget (₹)",
    min_value=10000,
    max_value=10000000,
    value=DEFAULT_BUDGET,
    step=10000
)

risk = st.sidebar.selectbox(
    "Risk Profile",
    ["Conservative", "Moderate", "Aggressive"]
)

horizon = st.sidebar.radio(
    "Investment Horizon (Days)",
    [5, 10]
)

st.sidebar.divider()

# Updated Model Sidebar (Pipelines and Architecture Overview)
st.sidebar.info(f"""
### Model Pipeline

📊 Feature Engineering
   
⬇️
   
🤖 Gradient Boosting Return Prediction
   
⬇️
   
🧠 PPO Reinforcement Learning
   
⬇️
   
💼 Dynamic Portfolio Allocation

---

**Universe:** 10 Indian Stocks

**Data Window:** Historical Static Dataset

**Future Enhancement:** Live Market Data Integration
""")

generate = st.sidebar.button("🚀 Optimize Portfolio", use_container_width=True)

# Main Application Logic
if generate:
    with st.spinner("Optimizing Portfolio..."):
        # Fetch Recommendations
        allocation, cash, summary, weights, snapshot_date  = generate_recommendation(budget, risk, horizon)
        
        # Data Cleaning, Sorting & Precision Preprocessing
        allocation = allocation.sort_values("Weight (%)", ascending=False).reset_index(drop=True)
        allocation["Investment (₹)"] = allocation["Investment (₹)"].round(0)
        allocation["Weight (%)"] = allocation["Weight (%)"].round(2)
        
        # Asset Allocation Descriptive Labels
        allocation["Recommendation"] = allocation["Weight (%)"].apply(
            lambda x: "🟢 High Allocation" if x >= 15 else ("🟡 Medium Allocation" if x >= 7 else "⚪ Low Allocation")
        )
        
        # 1. Adaptable Universe Scaling: Calculate purely from returned dataset dimensions
        total_universe = len(allocation)
        allocated_count = (allocation["Weight (%)"] > 0.5).sum()
        
        # KPI Metric Cards
        col1, col2, col3, col4 = st.columns(4)
        col1.metric("💰 Investment Budget", f"₹{budget:,.0f}")
        col2.metric("📈 Recommended Investment", f"₹{summary['Total Investment']:,.0f}")
        col3.metric("💵 Cash Reserve", f"₹{cash:,.0f}")
        col4.metric("📊 Stocks Allocated", f"{allocated_count} / {total_universe}")
        
        # Split Layout: Pie Chart and Recommendation Summary Panels
        left, right = st.columns([3, 2])
        
        with left:
            st.subheader("Recommended Portfolio Allocation")
            fig = px.pie(
                allocation,
                names="Stock",
                values="Investment (₹)",
                hole=0.45
            )
            fig.update_layout(
                height=550,
                legend_title="Stocks"
            )
            st.plotly_chart(fig, use_container_width=True, config={"displayModeBar": False})
            
        with right:
            st.subheader("Recommendation Summary")
            st.success(f"💵 Cash Reserve\n\n₹{cash:,.0f}")
            
            diversification = (allocation["Weight (%)"] > 5).sum() / total_universe
            st.write("Diversification Index")
            st.progress(diversification)
            
            # Additional User Context Metrics
            meta_col1, meta_col2 = st.columns(2)
            with meta_col1:
                st.markdown(f"**Recommendation Based On:** {snapshot_date}")
            with meta_col2:
                st.markdown(f"**Investment Horizon:** {horizon} Days")

        # Horizontal Bar Chart (Sorted)
        st.divider()
        st.subheader("📊 Stock Allocation Balance")
        
        bar_fig = px.bar(
            allocation,
            x="Weight (%)",
            y="Stock",
            orientation="h",
            text="Weight (%)"
        )
        bar_fig.update_traces(
            texttemplate="%{text:.1f} %",
            textposition="outside"
        )
        bar_fig.update_layout(
            yaxis=dict(categoryorder="total ascending"),
            height=500,
            xaxis_title="Portfolio Weight (%)",
            yaxis_title=""
        )
        st.plotly_chart(bar_fig, use_container_width=True, config={"displayModeBar": False})

        # PPO vs Equal Weight Comparison Strategy Chart
        st.divider()
        st.subheader("⚖ PPO vs Equal Weight Comparison")
        
        comparison = allocation.copy()
        comparison["Equal Weight"] = 100 / total_universe
        comparison = comparison.rename(columns={"Weight (%)": "PPO"})
        
        comparison_long = comparison.melt(
            id_vars="Stock",
            value_vars=["PPO", "Equal Weight"],
            var_name="Strategy",
            value_name="Weight"
        )
        
        comp_fig = px.bar(
            comparison_long,
            x="Stock",
            y="Weight",
            color="Strategy",
            barmode="group",
            text="Weight"
        )
        comp_fig.update_traces(
            texttemplate="%{text:.1f}%",
            textposition="outside"
        )
        st.plotly_chart(comp_fig, use_container_width=True, config={"displayModeBar": False})
        st.caption("_Equal Weight assigns an identical fixed allocation to every stock, while the PPO Reinforcement Learning agent dynamically adjusts allocations based on learned market risk patterns and current portfolio states._")
        
        # Structural Asset Allocation Statistics
        st.divider()
        c1, c2, c3 = st.columns(3)
        c1.metric("Largest Allocation", f"{allocation.iloc[0]['Weight (%)']:.1f}%")
        c2.metric("Average Allocation", f"{allocation['Weight (%)'].mean():.1f}%")
        c3.metric("Diversified Stocks (>5% Weight)", (allocation["Weight (%)"] > 5).sum())

        # Cleaned Standard Header Title "Recommended Portfolio"
        st.divider()
        st.subheader("📋 Recommended Portfolio")
        
        table_display = allocation[["Stock", "Weight (%)", "Investment (₹)", "Recommendation"]].copy()
        cash_weight = (cash / budget) * 100
        cash_row = pd.DataFrame([{
            "Stock": "💵 Cash Reserve",
            "Weight (%)": round(cash_weight, 2),
            "Investment (₹)": round(cash, 0),
            "Recommendation": "Liquidity Reserve"
        }])
        table_display = pd.concat([table_display, cash_row], ignore_index=True)
        
        st.dataframe(
            table_display,
            use_container_width=True,
            hide_index=True
        )

        # Highlight Top Asset Pick
        st.divider()
        top = allocation.iloc[0]
        st.success(f"""
        ## ⭐ Top Recommendation
        ### {top['Stock']}
        
        Recommended Investment
        ### ₹{top['Investment (₹)']:,.0f}
        
        Portfolio Weight
        ### {top['Weight (%)']:.2f}%
        
        Allocation Strategy
        ### {top['Recommendation']}
        """)

        # Dynamic Portfolio Insights Panel
        st.divider()
        st.subheader("💡 Dynamic Portfolio Insights")
        
        diversified_count = (allocation["Weight (%)"] > 5).sum()
        overweighted_df = allocation[allocation["Weight (%)"] > (100 / total_universe)]
        overweighted_names = overweighted_df["Stock"].head(3).tolist()
        overweighted_str = ", ".join(overweighted_names) if overweighted_names else "None"
        
        st.write(f"""
        * 📈 **Highest Allocation Asset:** **{top['Stock']}** received the highest portfolio allocation from the PPO agent at **{top['Weight (%)']:.2f}%**.
        * 📉 **Lowest Allocation Asset:** **{allocation.iloc[-1]['Stock']}** has the lowest framework allocation at **{allocation.iloc[-1]['Weight (%)']:.2f}%**.
        * 💵 **Liquidity Management:** Cash balance retention is securely held at **₹{cash:,.0f}** as a tactical Liquidity Reserve.
        * 🧩 **Diversification Scope:** **{diversified_count} out of {total_universe}** asset blocks successfully crossed the 5% concentration limit index.
        * 🧠 **RL Comparison Profile:** Compared with an equal-weight strategy, the PPO agent allocated larger weights to: **{overweighted_str}**.
        """)

        # Capital Utilization Gauge Graphic
        st.divider()
        st.subheader("💰 Capital Utilization")
        
        invested = summary["Total Investment"]
        utilization_percent = (invested / budget) * 100
        
        gauge_fig = go.Figure(
            go.Indicator(
                mode="gauge+number",
                value=utilization_percent,
                number={"suffix": "%"},
                title={"text": f"₹{invested:,.0f} / ₹{budget:,.0f}"},
                gauge={
                    "axis": {"range": [0, 100]},
                    "bar": {"color": "#10B981"}
                }
            )
        )
        gauge_fig.update_layout(height=350)
        st.plotly_chart(gauge_fig, use_container_width=True, config={"displayModeBar": False})

        # Natural Download Branding Action String
        st.divider()
        csv = table_display.to_csv(index=False)
        st.download_button(
            label="📥 Download Recommendation",
            data=csv,
            file_name="Dynamic_Allocation_Recommendation.csv",
            mime="text/csv",
            use_container_width=True
        )

else:
    # Humanized Welcome Screen Layout
    st.info("""
    👋 **Welcome to the Dynamic Allocation System.** Configure your investment preferences using the sidebar and click **"Optimize Portfolio"** to receive a personalized portfolio allocation generated by the PPO reinforcement learning agent.
    """)

# Polished Software About Layout View
st.markdown("<br>", unsafe_allow_html=True)
with st.expander("ℹ️ About This Project"):
    st.write(f"""
    This application demonstrates a machine learning-based dynamic portfolio allocation system.
    
    The pipeline consists of:
    * **Feature Engineering Framework:** Extracts trend, momentum, and risk indicators.
    * **Gradient Boosting Prediction Engine:** Estimates expected asset trajectory behaviors.
    * **PPO Reinforcement Learning Strategy:** Evaluates continuous state transitions to optimize portfolio allocations.
    
    ---
    * **Deployment Status:** Historical Static Dataset for Reproducibility
    * **System Architecture:** Pipeline Integration v1.0
    
    _Future versions will support live market data streaming and programmatic broker APIs. Recommendations are generated from the historical dataset.._
    """)