import streamlit as st
import plotly.graph_objects as go
import pandas as pd
import numpy as np
from utils.data_loader import load_nifty50_symbols, fetch_stock_data, get_stock_info
from utils.market_analysis import MarketAnalyzer
import asyncio
import json

st.title("Market Sentiment Heatmap")

# Initialize market analyzer
market_analyzer = MarketAnalyzer()

# Create tabs for different views
tab1, tab2 = st.tabs(["Sector-wise Sentiment", "Stock-wise Sentiment"])

with tab1:
    st.subheader("Sector Sentiment Analysis")

    # Define sectors and their stocks
    sectors = {
        "IT": ["TCS.NS", "INFY.NS", "HCLTECH.NS", "TECHM.NS", "WIPRO.NS"],
        "Banking & Finance": ["HDFCBANK.NS", "ICICIBANK.NS", "SBIN.NS", "AXISBANK.NS", "KOTAKBANK.NS"],
        "Energy & Oil": ["RELIANCE.NS", "ONGC.NS", "POWERGRID.NS", "NTPC.NS", "BPCL.NS"],
        "Automobile": ["TATAMOTORS.NS", "M&M.NS", "MARUTI.NS", "HEROMOTOCO.NS"],
        "Consumer Goods": ["HINDUNILVR.NS", "ITC.NS", "NESTLEIND.NS", "BRITANNIA.NS"],
        "Metals & Mining": ["TATASTEEL.NS", "HINDALCO.NS", "JSWSTEEL.NS", "COALINDIA.NS"],
        "Pharmaceuticals": ["SUNPHARMA.NS", "DRREDDY.NS", "CIPLA.NS", "DIVISLAB.NS"]
    }

    # Calculate sector sentiment
    sector_data = {}

    with st.spinner("Analyzing sector sentiment..."):
        for sector, symbols in sectors.items():
            valid_changes = []
            for symbol in symbols:
                try:
                    data = fetch_stock_data(symbol, period='1d')
                    if data is not None and not data.empty:
                        change = ((data['Close'].iloc[-1] - data['Open'].iloc[0]) / data['Open'].iloc[0]) * 100
                        valid_changes.append(change)
                except Exception as e:
                    continue

            if valid_changes:
                sector_data[sector] = np.mean(valid_changes)

    if sector_data:
        sectors = list(sector_data.keys())
        values = list(sector_data.values())

        fig = go.Figure(data=[go.Treemap(
            labels=sectors,
            parents=[''] * len(sectors),
            values=[abs(v) for v in values],
            textinfo="label+value+percent",
            marker=dict(
                colors=values,
                colorscale='RdYlGn',
                showscale=True,
                colorbar=dict(title="Change %")
            ),
            hovertemplate="""
            Sector: %{label}<br>
            Average Change: %{color:.2f}%<br>
            <extra></extra>
            """
        )])

        fig.update_layout(
            title="Sector-wise Market Sentiment",
            width=800,
            height=600
        )

        st.plotly_chart(fig, use_container_width=True)

        # Display sector statistics
        st.subheader("Sector Performance")
        sector_stats = pd.DataFrame({
            'Sector': sectors,
            'Average Change %': values
        }).sort_values('Average Change %', ascending=False)

        st.dataframe(sector_stats)
    else:
        st.warning("No sector data available. Please try again later.")

with tab2:
    st.subheader("Stock-wise Sentiment")

    # Select sector for stock analysis
    selected_sector = st.selectbox("Select Sector", list(sectors.keys()))

    if selected_sector:
        symbols = sectors[selected_sector]
        stock_data = []

        with st.spinner(f"Analyzing sentiment for {selected_sector} sector stocks..."):
            for symbol in symbols:
                try:
                    # Get basic stock info
                    info = get_stock_info(symbol)
                    if info is None:
                        st.warning(f"Could not fetch info for {symbol}")
                        continue

                    # Fetch stock data
                    data = fetch_stock_data(symbol, period='5d')
                    if data is None or data.empty:
                        st.warning(f"Could not fetch data for {symbol}")
                        continue

                    # Get AI sentiment
                    sentiment_response = await market_analyzer.get_ai_sentiment(symbol, data)

                    if sentiment_response.startswith("Error"):
                        st.error(f"Sentiment analysis error for {symbol}: {sentiment_response}")
                        continue

                    try:
                        # Parse sentiment response
                        sentiment_dict = json.loads(sentiment_response.replace("'", '"'))

                        sentiment_score = {
                            'bullish': 1,
                            'neutral': 0,
                            'bearish': -1
                        }.get(str(sentiment_dict.get('sentiment', '')).lower(), 0)

                        confidence = float(sentiment_dict.get('confidence', 0.5))
                        final_score = sentiment_score * confidence

                        stock_data.append({
                            'name': info['name'],
                            'symbol': symbol,
                            'sentiment_score': final_score,
                            'market_cap': info['market_cap'],
                            'sentiment': sentiment_dict.get('sentiment', 'N/A'),
                            'confidence': confidence,
                            'recommendation': sentiment_dict.get('recommendation', 'N/A'),
                            'risk_level': sentiment_dict.get('risk_level', 'N/A'),
                            'key_factors': sentiment_dict.get('key_factors', [])
                        })

                    except json.JSONDecodeError as e:
                        st.error(f"Error parsing sentiment data for {symbol}: {str(e)}")
                        continue
                    except Exception as e:
                        st.error(f"Unexpected error processing {symbol}: {str(e)}")
                        continue

            if stock_data:
                # Create heatmap
                fig = go.Figure(data=[go.Treemap(
                    labels=[f"{d['name']}<br>({d['symbol']})" for d in stock_data],
                    parents=[''] * len(stock_data),
                    values=[abs(d['market_cap']) for d in stock_data],
                    textinfo="label",
                    marker=dict(
                        colors=[d['sentiment_score'] for d in stock_data],
                        colorscale='RdYlGn',
                        showscale=True,
                        colorbar=dict(title="Sentiment Score")
                    ),
                    hovertemplate="""
                    Company: %{label}<br>
                    Sentiment Score: %{color:.2f}<br>
                    Market Cap: ₹%{value:,.0f}<br>
                    <extra></extra>
                    """
                )])

                fig.update_layout(
                    title=f"{selected_sector} Sector Sentiment Heatmap",
                    width=800,
                    height=600
                )

                st.plotly_chart(fig, use_container_width=True)

                # Create summary table
                st.subheader("Sentiment Summary")
                summary_df = pd.DataFrame([{
                    'Company': d['name'],
                    'Symbol': d['symbol'],
                    'Sentiment': d['sentiment'],
                    'Confidence': f"{d['confidence']:.2f}",
                    'Recommendation': d['recommendation'],
                    'Risk Level': d['risk_level']
                } for d in stock_data])

                st.dataframe(summary_df)

                # Show detailed insights
                st.subheader("Detailed Analysis")
                for stock in stock_data:
                    with st.expander(f"{stock['name']} ({stock['symbol']})"):
                        st.markdown(f"""
                        **Sentiment:** {stock['sentiment']}  
                        **Confidence:** {stock['confidence']:.2f}  
                        **Recommendation:** {stock['recommendation']}  
                        **Risk Level:** {stock['risk_level']}

                        **Key Factors:**
                        """)
                        for factor in stock['key_factors']:
                            st.markdown(f"- {factor}")
            else:
                st.warning("No sentiment data available. Please try again later.")
    else:
        st.info("Please select a sector to view stock-wise sentiment analysis.")