import requests
import pandas as pd
import plotly.graph_objects as go
from ultralytics import YOLO
import cv2
import time
import gradio as gr

API_KEY = "ITWJ6NDTF45CBTDO"

def get_stock_candlestick_data(symbol, interval="5min", output_size="compact"):
    """
    Fetch stock candlestick data from Alpha Vantage.
    """
    url = f"https://www.alphavantage.co/query?function=TIME_SERIES_INTRADAY&symbol={symbol}&interval={interval}&apikey={API_KEY}&outputsize={output_size}"
    print(f"Fetching data from: {url}")  # Debugging
    response = requests.get(url)
    if response.status_code == 200:
        data = response.json()
        print("API Response:", data)  # Debugging
        if f"Time Series ({interval})" in data:
            return data[f"Time Series ({interval})"]
        else:
            print("Error: No candlestick data found in response.")
            print(data)
            return None
    else:
        print(f"Error fetching data: {response.status_code}")
        print(response.text)
        return None

def process_stock_candlestick_data(data):
    """
    Process Alpha Vantage stock candlestick data into a DataFrame.
    """
    rows = []
    for timestamp, values in data.items():
        rows.append({
            "timestamp": timestamp,
            "open": float(values["1. open"]),
            "high": float(values["2. high"]),
            "low": float(values["3. low"]),
            "close": float(values["4. close"]),
            "volume": float(values["5. volume"])
        })
    return pd.DataFrame(rows)

def generate_candlestick_chart(df, n=50):
    """
    Generate a candlestick chart using Plotly with the last n data points.
    """
    df = df.tail(n)  # Use only the last n rows
    fig = go.Figure(data=[go.Candlestick(
        x=df["timestamp"],
        open=df["open"],
        high=df["high"],
        low=df["low"],
        close=df["close"]
    )])
    fig.update_layout(
        title="Candlestick Chart",
        xaxis_title="Time",
        yaxis_title="Price",
        xaxis_rangeslider_visible=False
    )
    # Removed fig.show() since Gradio will display the image
    fig.write_image("candlestick.png")

def yolo_model(img_path, model):
    """
    Run YOLO model on the image and count GAP UP and GAP DOWN patterns.
    """
    results = model(img_path)
    gap_up_count = 0
    gap_down_count = 0
    for result in results:
        classes = result.boxes.cls
        for cls in classes:
            if cls == 0:
                gap_down_count += 1
            elif cls == 1:
                gap_up_count += 1
        annotated_image = result.plot()
    return annotated_image, gap_up_count, gap_down_count

def detect_gap_patterns(symbol):
    """
    Main function to fetch data, generate charts, and detect GAP patterns in near-real-time.
    """
    model = YOLO("/content/best.pt")  # Load model once outside the loop
    while True:
        data = get_stock_candlestick_data(symbol)
        if not data:
            print("Failed to fetch data. Retrying in 15 seconds.")
            time.sleep(15)
            continue  # Retry instead of exiting

        df = process_stock_candlestick_data(data)
        generate_candlestick_chart(df, n=50)  # Generate chart with last 50 candles
        annotated_image, gap_up_count, gap_down_count = yolo_model("candlestick.png", model)
        cv2.imwrite("annotated_output.png", annotated_image)
        yield "annotated_output.png", gap_up_count, gap_down_count
        time.sleep(15)  # Wait 15 seconds to respect API rate limits

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# GAP Pattern Detection in Real-Time Stock Charts")
    gr.Markdown("Enter a stock symbol (e.g., AAPL) to detect GAP UP and GAP DOWN patterns in near-real-time candlestick charts.")
    
    with gr.Row():
        symbol_input = gr.Textbox(label="Stock Symbol", placeholder="Enter a stock symbol (e.g., AAPL)")
        submit_button = gr.Button("Start Real-Time Detection")
    
    with gr.Row():
        output_image = gr.Image(label="Annotated Candlestick Chart")
        gap_up_output = gr.Textbox(label="GAP UP Count")
        gap_down_output = gr.Textbox(label="GAP DOWN Count")
    
    # Start real-time detection when the button is clicked
    submit_button.click(
        fn=detect_gap_patterns,
        inputs=symbol_input,
        outputs=[output_image, gap_up_output, gap_down_output]
    )

# Launch the Gradio app
demo.launch(share=True, debug=True)