import streamlit as st
import pandas as pd
import numpy as np
import yfinance as yf
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import ta
from datetime import datetime, timedelta
import pytz

def calculate_technical_indicators(data):
    """Calculate comprehensive technical indicators"""
    # Trend Indicators
    data['SMA_20'] = ta.trend.sma_indicator(data['Close'], window=20)
    data['SMA_50'] = ta.trend.sma_indicator(data['Close'], window=50)
    data['SMA_200'] = ta.trend.sma_indicator(data['Close'], window=200)
    data['EMA_20'] = ta.trend.ema_indicator(data['Close'], window=20)
    
    # Momentum Indicators
    data['RSI'] = ta.momentum.rsi(data['Close'], window=14)
    data['MACD'] = ta.trend.macd_diff(data['Close'])
    data['ROC'] = ta.momentum.roc(data['Close'], window=12)
    
    # Volume Indicators
    data['OBV'] = ta.volume.on_balance_volume(data['Close'], data['Volume'])
    data['MFI'] = ta.volume.money_flow_index(data['High'], data['Low'], data['Close'], data['Volume'])
    
    # Volatility Indicators
    bb = ta.volatility.BollingerBands(close=data['Close'])
    data['BB_upper'] = bb.bollinger_hband()
    data['BB_middle'] = bb.bollinger_mavg()
    data['BB_lower'] = bb.bollinger_lband()
    data['ATR'] = ta.volatility.average_true_range(data['High'], data['Low'], data['Close'])
    
    # Additional Indian Market Specific Calculations
    data['Daily_Return'] = data['Close'].pct_change() * 100
    data['52W_High'] = data['Close'].rolling(window=252).max()
    data['52W_Low'] = data['Close'].rolling(window=252).min()
    data['Distance_From_52W_High'] = ((data['Close'] - data['52W_High'])/data['52W_High']) * 100
    
    return data

def plot_charts(data, symbol, interval):
    """Create interactive charts using plotly"""
    # Determine the number of rows based on the interval
    minute_intervals = ["1m", "2m", "5m", "15m", "30m", "60m", "90m", "1h"]
    if interval in minute_intervals:
        rows = 3
        subplot_titles = ('Price & Volume', 'RSI', 'MACD')
        row_heights = [0.5, 0.25, 0.25]
    else:
        rows = 4
        subplot_titles = ('Price & Volume', 'RSI', 'MACD', 'Technical Indicators')
        row_heights = [0.4, 0.2, 0.2, 0.2]
    
    fig = make_subplots(rows=rows, cols=1, 
                       shared_xaxes=True,
                       vertical_spacing=0.05,
                       subplot_titles=subplot_titles,
                       row_heights=row_heights)
    
    # Main candlestick chart
    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)
    
    # Add Moving Averages
    if 'SMA_20' in data.columns:
        fig.add_trace(go.Scatter(x=data.index, y=data['SMA_20'], name='SMA 20', line=dict(color='orange')), row=1, col=1)
    if 'SMA_50' in data.columns:
        fig.add_trace(go.Scatter(x=data.index, y=data['SMA_50'], name='SMA 50', line=dict(color='blue')), row=1, col=1)
    if 'SMA_200' in data.columns:
        fig.add_trace(go.Scatter(x=data.index, y=data['SMA_200'], name='SMA 200', line=dict(color='red')), row=1, col=1)
    
    if interval not in minute_intervals:
        # Volume bars for daily and higher intervals
        colors = ['red' if row['Open'] - row['Close'] >= 0 else 'green' for index, row in data.iterrows()]
        fig.add_trace(go.Bar(x=data.index, y=data['Volume'], name='Volume', marker_color=colors), row=1, col=1)
    
    # RSI
    if 'RSI' in data.columns:
        fig.add_trace(go.Scatter(x=data.index, y=data['RSI'], name='RSI', line=dict(color='purple')), row=2, col=1)
        fig.add_hline(y=70, line_dash="dash", line_color="red", row=2, col=1)
        fig.add_hline(y=30, line_dash="dash", line_color="green", row=2, col=1)
    
    # MACD
    if 'MACD' in data.columns:
        fig.add_trace(go.Scatter(x=data.index, y=data['MACD'], name='MACD', line=dict(color='blue')), row=3, col=1)
    
    if rows == 4:
        # Bollinger Bands
        if 'BB_upper' in data.columns and 'BB_middle' in data.columns and 'BB_lower' in data.columns:
            fig.add_trace(go.Scatter(x=data.index, y=data['BB_upper'], name='BB Upper',
                                    line=dict(color='gray', dash='dash')), row=4, col=1)
            fig.add_trace(go.Scatter(x=data.index, y=data['BB_middle'], name='BB Middle',
                                    line=dict(color='gray')), row=4, col=1)
            fig.add_trace(go.Scatter(x=data.index, y=data['BB_lower'], name='BB Lower',
                                    line=dict(color='gray', dash='dash')), row=4, col=1)
    
    # Update layout
    fig.update_layout(
        title=f'{symbol} Technical Analysis Dashboard',
        yaxis_title='Price',
        height=800 if rows == 4 else 600,
        showlegend=True,
        xaxis_rangeslider_visible=False
    )
    
    return fig

def fetch_stock_data(symbol, start_date, end_date, interval):
    """Fetch stock data with error handling and data validation"""
    try:
        # Convert dates to IST timezone for accuracy
        ist = pytz.timezone('Asia/Kolkata')
        end_date_ist = datetime.now(ist).date()
        
        ticker = yf.Ticker(symbol)
        data = ticker.history(start=start_date.strftime('%Y-%m-%d'), end=end_date.strftime('%Y-%m-%d'), interval=interval)
        
        if data.empty:
            st.error(f"No data available for {symbol} with interval '{interval}'. Please verify the symbol, interval, or try a different date range.")
            return None
            
        # Calculate all technical indicators
        data = calculate_technical_indicators(data)
        
        # Save data to CSV
        filename = f'{symbol}_{interval}.csv'
        data.to_csv(filename)
        
        return data
        
    except Exception as e:
        st.error(f"Error fetching data for {symbol}: {e}")
        return None

def main():
    st.set_page_config(layout="wide")
    st.title("Indian Stock Market Technical Analysis Dashboard")
    
    # Define minute-level and daily/higher intervals
    minute_intervals = ["1m", "2m", "5m", "15m", "30m", "60m", "90m", "1h"]
    daily_or_higher_intervals = ["1d", "5d", "1wk", "1mo", "3mo"]
    all_intervals = minute_intervals + daily_or_higher_intervals
    
    # Sidebar for inputs
    st.sidebar.header("Configuration")
    symbol = st.sidebar.text_input("Enter Stock Symbol (e.g., RELIANCE.NS, TCS.NS):", "RELIANCE.NS")
    
    interval = st.sidebar.selectbox(
        "Select Data Interval",
        all_intervals,
        index=8  # Default to '1d'
    )
    
    # Timezone and current date in IST
    ist = pytz.timezone('Asia/Kolkata')
    end_date_ist = datetime.now(ist).date()
    
    # Handle start date based on interval
    if interval in minute_intervals:
        # Automatically set start_date to last 7 days
        start_date = end_date_ist - timedelta(days=7)
        st.sidebar.write(f"**Note:** For interval '{interval}', only the last 7 days of data are available.")
        st.sidebar.write(f"**Start Date:** {start_date.strftime('%Y-%m-%d')}")
    else:
        # Allow user to select start_date for daily and higher intervals
        # Define a reasonable maximum history, e.g., 5 years
        max_history = 5 * 365  # 5 years
        default_start = end_date_ist - timedelta(days=365)  # Default to 1 year ago
        start_date = st.sidebar.date_input(
            "Start Date",
            value=default_start,
            min_value=end_date_ist - timedelta(days=max_history),
            max_value=end_date_ist - timedelta(days=1)
        )
    
    # Define end_date as today
    end_date = end_date_ist
    
    # Add fetch button
    if st.sidebar.button("Fetch Data and Update Charts"):
        with st.spinner('Fetching data...'):
            data = fetch_stock_data(symbol, start_date, end_date, interval)
            
            if data is not None:
                # Display basic info
                st.subheader(f"Analysis for {symbol}")
                current_price = data['Close'][-1]
                if len(data) >= 2:
                    daily_change = ((data['Close'][-1] - data['Close'][-2])/data['Close'][-2]) * 100
                else:
                    daily_change = 0.0  # Not enough data to calculate change
                
                col1, col2, col3, col4 = st.columns(4)
                col1.metric("Current Price", f"₹{current_price:.2f}", f"{daily_change:.2f}%")
                
                if '52W_High' in data.columns and '52W_Low' in data.columns:
                    col2.metric("52 Week High", f"₹{data['52W_High'][-1]:.2f}")
                    col3.metric("52 Week Low", f"₹{data['52W_Low'][-1]:.2f}")
                
                if 'RSI' in data.columns:
                    col4.metric("RSI", f"{data['RSI'][-1]:.2f}")
                
                # Plot interactive charts
                fig = plot_charts(data, symbol, interval)
                st.plotly_chart(fig, use_container_width=True)
                
                # Technical Analysis Summary
                st.subheader("Technical Analysis Summary")
                
                # Simple trading signals
                signals = []
                if 'RSI' in data.columns:
                    rsi = data['RSI'][-1]
                    if rsi < 30:
                        signals.append("RSI indicates oversold conditions")
                    elif rsi > 70:
                        signals.append("RSI indicates overbought conditions")
                
                if 'MACD' in data.columns:
                    macd = data['MACD'][-1]
                    if macd > 0:
                        signals.append("MACD is positive - Bullish momentum")
                    else:
                        signals.append("MACD is negative - Bearish momentum")
                
                if 'SMA_20' in data.columns and 'SMA_50' in data.columns:
                    current_price = data['Close'][-1]
                    sma_20 = data['SMA_20'][-1]
                    sma_50 = data['SMA_50'][-1]
                    
                    if current_price > sma_20 and current_price > sma_50:
                        signals.append("Price is above major moving averages - Bullish")
                    elif current_price < sma_20 and current_price < sma_50:
                        signals.append("Price is below major moving averages - Bearish")
                
                for signal in signals:
                    st.write(f"• {signal}")
                
                if not signals:
                    st.write("No clear technical signals detected.")
                
                # Display raw data in expandable section
                with st.expander("View Raw Data"):
                    st.dataframe(data)
                    
                # Provide download link for CSV
                csv = data.to_csv(index=True)
                st.download_button(
                    label="Download data as CSV",
                    data=csv,
                    file_name=f'{symbol}_{interval}.csv',
                    mime='text/csv',
                )

if __name__ == "__main__":
    main()