import gradio as gr
import pandas as pd
import yfinance as yf
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from statsforecast import StatsForecast
from statsforecast.models import AutoARIMA

def calculate_indicators(df):
    # RSI Calculation
    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
    df['RSI'] = 100 - (100 / (1 + rs))
    
    # MACD Calculation
    df['EMA12'] = df['Close'].ewm(span=12, adjust=False).mean()
    df['EMA26'] = df['Close'].ewm(span=26, adjust=False).mean()
    df['MACD'] = df['EMA12'] - df['EMA26']
    df['Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
    return df

def get_pro_chart(df, forecast_df, ticker):
    # Professional Subplots: Price + Indicators
    fig = make_subplots(rows=3, cols=1, shared_xaxes=True, 
                        vertical_spacing=0.05, row_heights=[0.5, 0.25, 0.25],
                        subplot_titles=(f"{ticker} Forecast & Confidence Zone", "RSI Momentum", "MACD Trend"))

    # --- 1. Price + Forecast + CONFIDENCE ZONE ---
    # Shaded Confidence Area (AutoARIMA-hi-80 and AutoARIMA-lo-80)
    fig.add_trace(go.Scatter(
        x=forecast_df['ds'].tolist() + forecast_df['ds'].tolist()[::-1],
        y=forecast_df['AutoARIMA-hi-80'].tolist() + forecast_df['AutoARIMA-lo-80'].tolist()[::-1],
        fill='toself',
        fillcolor='rgba(0, 210, 255, 0.2)',
        line=dict(color='rgba(255,255,255,0)'),
        hoverinfo="skip",
        showlegend=True,
        name='80% Confidence'
    ), row=1, col=1)

    # Historical Price
    fig.add_trace(go.Scatter(x=df.index, y=df['Close'], name='Historical Price', line=dict(color='#00d2ff', width=2)), row=1, col=1)
    
    # AI Mean Forecast Line
    fig.add_trace(go.Scatter(x=forecast_df['ds'], y=forecast_df['AutoARIMA'], name='AI Target', line=dict(color='#F23645', dash='dot')), row=1, col=1)

    # --- 2. RSI ---
    fig.add_trace(go.Scatter(x=df.index, y=df['RSI'], name='RSI', line=dict(color='#FF9800')), row=2, col=1)
    fig.add_hline(y=70, line_dash="dot", line_color="red", row=2, col=1)
    fig.add_hline(y=30, line_dash="dot", line_color="green", row=2, col=1)

    # --- 3. MACD ---
    fig.add_trace(go.Bar(x=df.index, y=df['MACD'] - df['Signal'], name='Momentum', marker_color='#9d50bb'), row=3, col=1)
    
    # Styling to match your Aether UI
    fig.update_layout(
        template='plotly_dark', 
        height=800, 
        showlegend=True, 
        paper_bgcolor='rgba(0,0,0,0)', # Transparent to show your terminal background
        plot_bgcolor='rgba(0,0,0,0)',
        font=dict(family="Lato", size=12)
    )
    fig.update_xaxes(showgrid=False, zeroline=False)
    fig.update_yaxes(showgrid=False, zeroline=False)
    
    return fig

def analyze(ticker, horizon):
    try:
        # Step 1: Data
        stock = yf.Ticker(ticker)
        df = stock.history(period="1y")
        if df.empty: return None, "Symbol Error"
        
        info = stock.info
        target_high = info.get('targetHighPrice', 'N/A')
        current_price = df['Close'].iloc[-1]
        
        # Step 2: Indicators
        df = calculate_indicators(df)
        
        # Step 3: StatsForecast with Levels
        data = df.reset_index()[['Date', 'Close']]
        data.columns = ['ds', 'y']
        data['unique_id'] = ticker
        
        sf = StatsForecast(models=[AutoARIMA(season_length=5)], freq='B')
        sf.fit(data)
        
        # We request 80% level to get the 'lo-80' and 'hi-80' columns
        forecast = sf.predict(h=horizon, level=[80])
        forecast = forecast.reset_index()
        
        # Step 4: Signal Analysis
        rsi_val = df['RSI'].iloc[-1]
        rsi_stat = "OVERSOLD" if rsi_val < 30 else "OVERBOUGHT" if rsi_val > 70 else "NEUTRAL"
        
        signal_html = f"""
        <div style='background: rgba(255,255,255,0.05); padding: 20px; border-radius: 15px; border: 1px solid rgba(255,255,255,0.1); font-family: "Lato", sans-serif;'>
            <p style='margin:0; color:#888; font-size:12px;'>CURRENT VALUE</p>
            <h2 style='margin:0; color:#00d2ff;'>${current_price:.2f}</h2>
            <hr style='border:0; border-top:1px solid #333; margin:15px 0;'>
            <p>RSI Analysis: <b style='color:#FF9800;'>{rsi_stat} ({rsi_val:.1f})</b></p>
            <p>AI Forecast ({horizon}d): <b style='color:#F23645;'>${forecast['AutoARIMA'].iloc[-1]:.2f}</b></p>
            <p style='font-size:11px; color:#555;'>Confidence Zone: ${forecast['AutoARIMA-lo-80'].iloc[-1]:.2f} - ${forecast['AutoARIMA-hi-80'].iloc[-1]:.2f}</p>
        </div>
        """
        
        return get_pro_chart(df, forecast, ticker), signal_html
    except Exception as e:
        return None, f"<p style='color:red;'>Error: {str(e)}</p>"

# --- UI Layout ---
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="slate")) as demo:
    gr.Markdown("<h1 style='text-align: center; letter-spacing: 5px;'>AETHER BACKEND ENGINE</h1>")
    with gr.Row():
        with gr.Column(scale=1):
            t_in = gr.Textbox(label="Ticker", value="NVDA")
            h_in = gr.Slider(7, 90, value=30, label="Forecast Days")
            btn = gr.Button("UPDATE ENGINE", variant="primary")
            info_out = gr.HTML()
        with gr.Column(scale=4):
            plot_out = gr.Plot()

    btn.click(analyze, [t_in, h_in], [plot_out, info_out])

if __name__ == "__main__":
    demo.launch()