import gradio as gr
import pandas as pd
import yfinance as yf
import plotly.graph_objects as go
import numpy as np

# Functions for calculating indicators
def calculate_sma(df, window):
    return df['Close'].rolling(window=window).mean()

def calculate_ema(df, window):
    return df['Close'].ewm(span=window, adjust=False).mean()

def calculate_macd(df):
    short_ema = df['Close'].ewm(span=12, adjust=False).mean()
    long_ema = df['Close'].ewm(span=26, adjust=False).mean()
    macd = short_ema - long_ema
    signal = macd.ewm(span=9, adjust=False).mean()
    return macd, signal

def calculate_rsi(df):
    delta = df['Close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

def calculate_bollinger_bands(df):
    middle_bb = df['Close'].rolling(window=20).mean()
    upper_bb = middle_bb + 2 * df['Close'].rolling(window=20).std()
    lower_bb = middle_bb - 2 * df['Close'].rolling(window=20).std()
    return middle_bb, upper_bb, lower_bb

def calculate_stochastic_oscillator(df):
    lowest_low = df['Low'].rolling(window=14).min()
    highest_high = df['High'].rolling(window=14).max()
    slowk = ((df['Close'] - lowest_low) / (highest_high - lowest_low)) * 100
    slowd = slowk.rolling(window=3).mean()
    return slowk, slowd

def calculate_cmf(df, window=20):
    mfv = ((df['Close'] - df['Low']) - (df['High'] - df['Close'])) / (df['High'] - df['Low']) * df['Volume']
    cmf = mfv.rolling(window=window).sum() / df['Volume'].rolling(window=window).sum()
    return cmf

def calculate_cci(df, window=20):
    typical_price = (df['High'] + df['Low'] + df['Close']) / 3
    sma = typical_price.rolling(window=window).mean()
    mean_deviation = (typical_price - sma).abs().rolling(window=window).mean()
    cci = (typical_price - sma) / (0.015 * mean_deviation)
    return cci

# Function to adjust thresholds based on sensitivity
def adjust_thresholds_by_sensitivity(sensitivity):
    """
    Convert a single sensitivity value (1-10) to appropriate thresholds
    1 = Most sensitive (more signals)
    10 = Least sensitive (fewer, stronger signals)
    """
    # Map sensitivity to thresholds
    if sensitivity == 1:  # Most sensitive
        return {
            'SMA': 5,
            'RSI_lower': 30,
            'RSI_upper': 70,
            'BB': 0.5,
            'Stochastic_lower': 20,
            'Stochastic_upper': 80,
            'CMF': 0.1,
            'CCI': 100
        }
    elif sensitivity == 10:  # Least sensitive
        return {
            'SMA': 50,
            'RSI_lower': 5,
            'RSI_upper': 95,
            'BB': 5,
            'Stochastic_lower': 5,
            'Stochastic_upper': 95,
            'CMF': 0.6,
            'CCI': 300
        }
    else:
        # Linear interpolation between extremes
        factor = (sensitivity - 1) / 9  # 0 to 1
        return {
            'SMA': int(5 + (50 - 5) * factor),
            'RSI_lower': int(30 - (30 - 5) * factor),
            'RSI_upper': int(70 + (95 - 70) * factor),
            'BB': 0.5 + (5 - 0.5) * factor,
            'Stochastic_lower': int(20 - (20 - 5) * factor),
            'Stochastic_upper': int(80 + (95 - 80) * factor),
            'CMF': 0.1 + (0.6 - 0.1) * factor,
            'CCI': int(100 + (300 - 100) * factor)
        }

def generate_trading_signals(df, thresholds, enabled_signals):
    # Calculate various indicators
    df['SMA_30'] = calculate_sma(df, 30)
    df['SMA_100'] = calculate_sma(df, 100)
    df['EMA_12'] = calculate_ema(df, 12)
    df['EMA_26'] = calculate_ema(df, 26)
    df['RSI'] = calculate_rsi(df)
    df['MiddleBB'], df['UpperBB'], df['LowerBB'] = calculate_bollinger_bands(df)
    df['SlowK'], df['SlowD'] = calculate_stochastic_oscillator(df)
    df['CMF'] = calculate_cmf(df)
    df['CCI'] = calculate_cci(df)

    # Initialize all signals as 0 (no signal)
    signal_columns = ['SMA_Signal', 'MACD_Signal', 'RSI_Signal', 'BB_Signal', 
                      'Stochastic_Signal', 'CMF_Signal', 'CCI_Signal']
    for col in signal_columns:
        df[col] = 0
    
    # Only generate signals for enabled indicators
    
    # SMA Signal
    if 'SMA' in enabled_signals:
        sma_threshold = thresholds['SMA']
        df['SMA_Diff_Pct'] = (df['SMA_30'] - df['SMA_100']) / df['SMA_100'] * 100
        df['SMA_Signal'] = np.where(df['SMA_Diff_Pct'] > sma_threshold, 1, 0)
        df['SMA_Signal'] = np.where(df['SMA_Diff_Pct'] < -sma_threshold, -1, df['SMA_Signal'])
    
    # MACD Signal
    if 'MACD' in enabled_signals:
        macd, signal = calculate_macd(df)
        df['MACD'] = macd
        df['MACD_Signal_Line'] = signal
        df['MACD_Signal'] = np.select([(macd > signal) & (macd.shift(1) <= signal.shift(1)),
                                      (macd < signal) & (macd.shift(1) >= signal.shift(1))], [1, -1], default=0)
    
    # RSI Signals
    if 'RSI' in enabled_signals:
        rsi_lower = thresholds['RSI_lower']
        rsi_upper = thresholds['RSI_upper']
        df['RSI_Signal'] = np.where(df['RSI'] < rsi_lower, 1, 0)
        df['RSI_Signal'] = np.where(df['RSI'] > rsi_upper, -1, df['RSI_Signal'])
    
    # Bollinger Bands
    if 'BB' in enabled_signals:
        bb_buffer = thresholds['BB'] / 100  # Convert percentage to decimal
        df['BB_Signal'] = np.where(
            (df['Close'] < df['LowerBB'] * (1 - bb_buffer)) & 
            (df['Close'].shift(1) < df['LowerBB'].shift(1) * (1 - bb_buffer)) & 
            (df['Close'].shift(2) < df['LowerBB'].shift(2) * (1 - bb_buffer)), 1, 0
        )
        df['BB_Signal'] = np.where(
            (df['Close'] > df['UpperBB'] * (1 + bb_buffer)) & 
            (df['Close'].shift(1) > df['UpperBB'].shift(1) * (1 + bb_buffer)) & 
            (df['Close'].shift(2) > df['UpperBB'].shift(2) * (1 + bb_buffer)), -1, df['BB_Signal']
        )
    
    # Stochastic signals
    if 'Stochastic' in enabled_signals:
        stoch_lower = thresholds['Stochastic_lower']
        stoch_upper = thresholds['Stochastic_upper']
        df['Stochastic_Signal'] = np.where((df['SlowK'] < stoch_lower) & (df['SlowD'] < stoch_lower), 1, 0)
        df['Stochastic_Signal'] = np.where((df['SlowK'] > stoch_upper) & (df['SlowD'] > stoch_upper), -1, df['Stochastic_Signal'])
    
    # CMF Signals
    if 'CMF' in enabled_signals:
        cmf_threshold = thresholds['CMF']
        df['CMF_Signal'] = np.where(df['CMF'] > cmf_threshold, -1, np.where(df['CMF'] < -cmf_threshold, 1, 0))
    
    # CCI Signals
    if 'CCI' in enabled_signals:
        cci_threshold = thresholds['CCI']
        df['CCI_Signal'] = np.where(df['CCI'] < -cci_threshold, 1, 0)
        df['CCI_Signal'] = np.where(df['CCI'] > cci_threshold, -1, df['CCI_Signal'])

    return df

def plot_simplified_signals(df, ticker, enabled_signals):
    # Create a figure with improved styling
    fig = go.Figure()
    
    # Use a line chart instead of candlestick for simplicity
    fig.add_trace(go.Scatter(
        x=df.index, 
        y=df['Close'],
        mode='lines', 
        name='Price', 
        line=dict(color='#26a69a', width=2),
        opacity=0.9
    ))
    
    # Add SMA lines
    fig.add_trace(go.Scatter(
        x=df.index, y=df['SMA_30'], 
        mode='lines', 
        name='SMA 30', 
        line=dict(color='#42a5f5', width=1.5, dash='dot')
    ))
    
    fig.add_trace(go.Scatter(
        x=df.index, y=df['SMA_100'], 
        mode='lines', 
        name='SMA 100', 
        line=dict(color='#5e35b1', width=1.5, dash='dot')
    ))

    # Add bollinger bands with lighter appearance
    if 'BB' in enabled_signals:
        fig.add_trace(go.Scatter(
            x=df.index, y=df['UpperBB'],
            mode='lines',
            name='Upper BB',
            line=dict(color='rgba(250, 250, 250, 0.3)', width=1),
            showlegend=True
        ))
        
        fig.add_trace(go.Scatter(
            x=df.index, y=df['LowerBB'],
            mode='lines',
            name='Lower BB',
            line=dict(color='rgba(250, 250, 250, 0.3)', width=1),
            fill='tonexty',
            fillcolor='rgba(173, 216, 230, 0.1)',
            showlegend=True
        ))

    # Group signals by type to reduce legend clutter
    buy_signals_df = pd.DataFrame(index=df.index)
    sell_signals_df = pd.DataFrame(index=df.index)
    
    signal_names = [f"{signal}_Signal" for signal in enabled_signals]
    
    # Collect all buy and sell signals
    for signal in signal_names:
        buy_signals_df[signal] = np.where(df[signal] == 1, df['Close'], np.nan)
        sell_signals_df[signal] = np.where(df[signal] == -1, df['Close'], np.nan)
    
    # Add hover data
    buy_hovers = []
    for idx in buy_signals_df.index:
        signals_on_day = [col.split('_')[0] for col in buy_signals_df.columns 
                         if not pd.isna(buy_signals_df.loc[idx, col])]
        if signals_on_day:
            hover_text = f"Buy Signals: {', '.join(signals_on_day)}<br>Date: {idx.strftime('%Y-%m-%d')}<br>Price: ${df.loc[idx, 'Close']:.2f}"
            buy_hovers.append((idx, df.loc[idx, 'Close'], hover_text))
    
    sell_hovers = []
    for idx in sell_signals_df.index:
        signals_on_day = [col.split('_')[0] for col in sell_signals_df.columns 
                         if not pd.isna(sell_signals_df.loc[idx, col])]
        if signals_on_day:
            hover_text = f"Sell Signals: {', '.join(signals_on_day)}<br>Date: {idx.strftime('%Y-%m-%d')}<br>Price: ${df.loc[idx, 'Close']:.2f}"
            sell_hovers.append((idx, df.loc[idx, 'Close'], hover_text))
    
    # Add buy signals (single trace for all buy signals)
    if buy_hovers:
        buy_x, buy_y, buy_texts = zip(*buy_hovers)
        fig.add_trace(go.Scatter(
            x=buy_x, 
            y=[y * 0.995 for y in buy_y],  # Position slightly below price for visibility
            mode='markers', 
            marker=dict(symbol='triangle-up', size=10, color='#00e676', line=dict(color='white', width=1)), 
            name='Buy Signals',
            hoverinfo='text',
            hovertext=buy_texts
        ))

    # Add sell signals (single trace for all sell signals)
    if sell_hovers:
        sell_x, sell_y, sell_texts = zip(*sell_hovers)
        fig.add_trace(go.Scatter(
            x=sell_x, 
            y=[y * 1.005 for y in sell_y],  # Position slightly above price for visibility
            mode='markers', 
            marker=dict(symbol='triangle-down', size=10, color='#ff5252', line=dict(color='white', width=1)), 
            name='Sell Signals',
            hoverinfo='text',
            hovertext=sell_texts
        ))

    # Improve the layout with larger dimensions
    fig.update_layout(
        title=dict(
            text=f'{ticker}: Technical Analysis & Trading Signals',
            font=dict(size=24, color='white'),
            x=0.5
        ),
        xaxis=dict(
            title='Date',
            gridcolor='rgba(255, 255, 255, 0.1)',
            linecolor='rgba(255, 255, 255, 0.2)'
        ),
        yaxis=dict(
            title='Price',
            side='right',
            gridcolor='rgba(255, 255, 255, 0.1)',
            linecolor='rgba(255, 255, 255, 0.2)',
            tickprefix='$'
        ),
        plot_bgcolor='#1e1e1e',
        paper_bgcolor='#1e1e1e',
        font=dict(color='white'),
        hovermode='closest',
        legend=dict(
            bgcolor='rgba(30, 30, 30, 0.8)',
            bordercolor='rgba(255, 255, 255, 0.2)',
            borderwidth=1,
            font=dict(color='white', size=10),
            orientation='h',
            yanchor='bottom',
            y=1.02,
            xanchor='center',
            x=0.5
        ),
        margin=dict(l=50, r=50, b=100, t=100, pad=4),
        height=800,  # Increased height
        width=1200   # Increased width
    )
    
    # Add range selector for better time navigation
    fig.update_xaxes(
        rangeslider_visible=True,
        rangeselector=dict(
            buttons=list([
                dict(count=1, label="1m", step="month", stepmode="backward"),
                dict(count=3, label="3m", step="month", stepmode="backward"),
                dict(count=6, label="6m", step="month", stepmode="backward"),
                dict(count=1, label="YTD", step="year", stepmode="todate"),
                dict(count=1, label="1y", step="year", stepmode="backward"),
                dict(step="all")
            ]),
            bgcolor='rgba(30, 30, 30, 0.8)',
            activecolor='#536dfe',
            font=dict(color='white')
        )
    )
    
    return fig

def stock_analysis(ticker, start_date, end_date, 
                  sensitivity,  # New simplified parameter
                  use_sma, use_macd, use_rsi, use_bb, 
                  use_stoch, use_cmf, use_cci):
    try:
        # Download stock data from Yahoo Finance
        df = yf.download(ticker, start=start_date, end=end_date)
        
        # Check if data was retrieved
        if df.empty:
            fig = go.Figure()
            fig.add_annotation(
                text="No data found for this ticker and date range",
                xref="paper", yref="paper", 
                x=0.5, y=0.5, 
                showarrow=False,
                font=dict(color="white", size=16)
            )
            fig.update_layout(
                plot_bgcolor='#1e1e1e',
                paper_bgcolor='#1e1e1e',
                height=800,
                width=1200
            )
            return fig

        # If the DataFrame has a MultiIndex for columns, handle it
        if isinstance(df.columns, pd.MultiIndex):
            df.columns = df.columns.droplevel(1) if len(df.columns.levels) > 1 else df.columns
        
        # Create list of enabled signals
        enabled_signals = []
        if use_sma: enabled_signals.append('SMA')
        if use_macd: enabled_signals.append('MACD')
        if use_rsi: enabled_signals.append('RSI')
        if use_bb: enabled_signals.append('BB')
        if use_stoch: enabled_signals.append('Stochastic')
        if use_cmf: enabled_signals.append('CMF')
        if use_cci: enabled_signals.append('CCI')
        
        # If no signals are enabled, enable all by default
        if not enabled_signals:
            enabled_signals = ['SMA', 'MACD', 'RSI', 'BB', 'Stochastic', 'CMF', 'CCI']
        
        # Get thresholds from sensitivity
        thresholds = adjust_thresholds_by_sensitivity(sensitivity)
        
        # Generate signals
        df = generate_trading_signals(df, thresholds, enabled_signals)
        
        # Last 360 days for plotting (or all data if less than 360 days)
        df_last_360 = df.tail(min(360, len(df)))

        # Plot simplified signals
        fig = plot_simplified_signals(df_last_360, ticker, enabled_signals)
        
        return fig
        
    except Exception as e:
        # Create error figure
        fig = go.Figure()
        fig.add_annotation(
            text=f"Error: {str(e)}",
            xref="paper", yref="paper", 
            x=0.5, y=0.5, 
            showarrow=False,
            font=dict(color="#ff5252", size=16)
        )
        fig.update_layout(
            plot_bgcolor='#1e1e1e',
            paper_bgcolor='#1e1e1e',
            font=dict(color='white'),
            height=800,
            width=1200
        )
        return fig

# Define Gradio interface with improved styling
custom_theme = gr.themes.Monochrome(
    primary_hue="blue",
    secondary_hue="purple",
    neutral_hue="gray",
    radius_size=gr.themes.sizes.radius_sm,
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
)

with gr.Blocks(theme=custom_theme) as demo:
    gr.Markdown("# Technical Analysis")
    gr.Markdown("This app helps you analyze stocks with technical indicators and generates trading signals.")

    with gr.Row():
        with gr.Column(scale=1):
            ticker_input = gr.Textbox(
                label="Stock Ticker Symbol",
                placeholder="e.g., AAPL, NVDA, MSFT",
                value="NVDA"
            )
            start_date_input = gr.Textbox(
                label="Start Date",
                placeholder="YYYY-MM-DD",
                value="2022-01-01"
            )
            end_date_input = gr.Textbox(
                label="End Date",
                placeholder="YYYY-MM-DD",
                value="2026-01-01"  # Updated to current date
            )
            
            gr.Markdown("### Choose Indicators")
            with gr.Row():
                use_sma = gr.Checkbox(label="SMA", value=True)
                use_macd = gr.Checkbox(label="MACD", value=True)
                use_rsi = gr.Checkbox(label="RSI", value=True)
                use_bb = gr.Checkbox(label="Bollinger", value=True)
                use_stoch = gr.Checkbox(label="Stochastic", value=True)
                use_cmf = gr.Checkbox(label="CMF", value=True)
                use_cci = gr.Checkbox(label="CCI", value=True)
            
            gr.Markdown("### Signal Sensitivity")
            with gr.Row():
                sensitivity = gr.Slider(
                    label="Signal Sensitivity", 
                    minimum=1, 
                    maximum=10, 
                    step=1, 
                    value=5, 
                    info="1 = (sensitive), 10 = (strict)"
                )
            
            # Create a submit button with styling
            button = gr.Button("Analyze Stock", variant="primary")
    
    # Output: Signals plot with increased height
    signals_output = gr.Plot(label="Technical Analysis & Trading Signals")

    # Link button to function with updated parameters
    button.click(
        stock_analysis, 
        inputs=[
            ticker_input, start_date_input, end_date_input,
            sensitivity,  # Single threshold parameter
            use_sma, use_macd, use_rsi, use_bb, 
            use_stoch, use_cmf, use_cci
        ], 
        outputs=[signals_output]
    )

    gr.Markdown("""
    ## 📈 Trading Signals Legend
    - **Green Triangle Up (▲)** indicates Buy signals
    - **Red Triangle Down (▼)** indicates Sell signals
    - Hover over signals to see which indicators triggered them
    
    ## 🔍 Signal Sensitivity Explained
    - **Lower values (1-3)**: More frequent signals, good for short-term trading
    - **Medium values (4-6)**: Balanced approach, moderate number of signals
    - **Higher values (7-10)**: Fewer but potentially stronger signals, good for long-term investors
    
    ## 🛠️ Trading Strategy Tips
    - **Day Trading**: Use lower sensitivity with multiple indicators
    - **Swing Trading**: Use medium sensitivity with 3-4 indicators
    - **Long-term Investing**: Use higher sensitivity focusing on trend indicators
    - **Combine**: Using multiple indicators helps confirm signals and reduce false positives
    """)

# Launch the interface
demo.launch()