import streamlit as st
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
from utils.data_loader import load_nifty50_symbols, fetch_stock_data
from utils.technical_indicators import (
    calculate_rsi,
    calculate_macd,
    calculate_bollinger_bands,
    calculate_support_resistance
)

st.title("Technical Analysis")

# Stock selection
symbol = st.selectbox("Select Stock", load_nifty50_symbols())
timeframe = st.selectbox(
    "Select Timeframe",
    options=["1w", "1mo", "3mo", "6mo", "1y", "3y", "5y"],
    format_func=lambda x: {
        "1w": "1 Week",
        "1mo": "1 Month",
        "3mo": "1 Quarter",
        "6mo": "6 Months",
        "1y": "1 Year",
        "3y": "3 Years",
        "5y": "5 Years"
    }[x]
)

# Fetch data
data = fetch_stock_data(symbol, period=timeframe)

if data is not None:
    # Create tabs for different analysis
    tab1, tab2, tab3 = st.tabs(["Price & Volume", "Technical Indicators", "Support & Resistance"])

    with tab1:
        # Candlestick chart with volume
        fig = make_subplots(rows=2, cols=1, shared_xaxes=True, 
                           vertical_spacing=0.03, 
                           row_heights=[0.7, 0.3])

        # Candlestick
        fig.add_trace(go.Candlestick(
            x=data.index,
            open=data['Open'],
            high=data['High'],
            low=data['Low'],
            close=data['Close'],
            name='OHLC'
        ), row=1, col=1)

        # Volume
        fig.add_trace(go.Bar(
            x=data.index,
            y=data['Volume'],
            name='Volume',
            marker_color='rgba(0,184,148,0.3)'
        ), row=2, col=1)

        fig.update_layout(
            title=f"{symbol} Price and Volume",
            yaxis_title="Price",
            yaxis2_title="Volume",
            xaxis_rangeslider_visible=False,
            height=800
        )

        st.plotly_chart(fig, use_container_width=True)

    with tab2:
        # Technical indicators
        col1, col2 = st.columns(2)

        with col1:
            # RSI
            rsi = calculate_rsi(data)
            fig_rsi = go.Figure()
            fig_rsi.add_trace(go.Scatter(
                x=data.index,
                y=rsi,
                name='RSI',
                line=dict(color='#00B894')
            ))
            fig_rsi.add_hline(y=70, line_dash="dash", line_color="red")
            fig_rsi.add_hline(y=30, line_dash="dash", line_color="green")
            fig_rsi.update_layout(title="RSI (14)", height=400)
            st.plotly_chart(fig_rsi, use_container_width=True)

        with col2:
            # MACD
            macd, signal = calculate_macd(data)
            fig_macd = go.Figure()
            fig_macd.add_trace(go.Scatter(
                x=data.index,
                y=macd,
                name='MACD',
                line=dict(color='#00B894')
            ))
            fig_macd.add_trace(go.Scatter(
                x=data.index,
                y=signal,
                name='Signal',
                line=dict(color='#FFA500')
            ))
            fig_macd.update_layout(title="MACD", height=400)
            st.plotly_chart(fig_macd, use_container_width=True)

        # Bollinger Bands
        upper_band, middle_band, lower_band = calculate_bollinger_bands(data)
        fig_bb = go.Figure()
        fig_bb.add_trace(go.Scatter(
            x=data.index,
            y=upper_band,
            name='Upper Band',
            line=dict(color='gray', dash='dash')
        ))
        fig_bb.add_trace(go.Scatter(
            x=data.index,
            y=middle_band,
            name='Middle Band',
            line=dict(color='blue')
        ))
        fig_bb.add_trace(go.Scatter(
            x=data.index,
            y=lower_band,
            name='Lower Band',
            line=dict(color='gray', dash='dash')
        ))
        fig_bb.add_trace(go.Scatter(
            x=data.index,
            y=data['Close'],
            name='Close Price',
            line=dict(color='#00B894')
        ))
        fig_bb.update_layout(title="Bollinger Bands", height=500)
        st.plotly_chart(fig_bb, use_container_width=True)

    with tab3:
        # Support and Resistance
        support, resistance = calculate_support_resistance(data)
        fig_sr = go.Figure()

        fig_sr.add_trace(go.Scatter(
            x=data.index,
            y=data['Close'],
            name='Price',
            line=dict(color='#00B894')
        ))
        fig_sr.add_trace(go.Scatter(
            x=data.index,
            y=support,
            name='Support',
            line=dict(color='green', dash='dash')
        ))
        fig_sr.add_trace(go.Scatter(
            x=data.index,
            y=resistance,
            name='Resistance',
            line=dict(color='red', dash='dash')
        ))

        fig_sr.update_layout(
            title="Support and Resistance Levels",
            height=600
        )
        st.plotly_chart(fig_sr, use_container_width=True)

else:
    st.error("Unable to fetch data. Please try again later.")