import gradio as gr
import yfinance as yf
import plotly.graph_objects as go
import plotly.express as px
from datetime import datetime
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# Function to construct the pie chart
def get_pie_graph(data):
  data['Change'] = data['Open'] - data['Close'].shift(1)
  data['Movement'] = data['Change'].apply(lambda x : 'Up' if x > 0 else 'Down')
  fig, ax = plt.subplots()
  ax.pie(data['Movement'].value_counts(), labels=[f"{label}: {count}" for label, count in zip(data['Movement'].unique(), data['Movement'].value_counts().tolist())])
  return fig

# Function to construct the chart with markers
def get_trading_markers(data, t, tll):
    val = data['Close']
    th = []
    marker = []

    for i in range(len(data)):
        if(i == 0):
            th.append(0)
        else:
            x = val[i] - val[i-1]
            x =  (x / val[i-1] )*10
            th.append(abs(x))

    data['Value'] = th
    data['Marker'] = np.where(data['Value'] > t, 'yes', 'no')

    df = data[data['Marker'] == 'yes']
    df = df.Close

    x = pd.DataFrame({tll: data["Close"]})
    plt.plot(x)
    plt.scatter(df.index, df, marker = 'x', c = 'b' , s = 60, zorder=2)
    plt.title(tll)

    return plt

def get_stock_data(start_date, end_date, stock_ticker, graph_type, t):
    # Validate date format
    try:
        start_date = datetime.strptime(start_date, "%Y-%m-%d").date()
        end_date = datetime.strptime(end_date, "%Y-%m-%d").date()
    except ValueError:
        return "Invalid date format. Please use the YYYY-MM-DD format."

    # Fetch stock data using Yahoo Finance API
    data = yf.download(stock_ticker, start=start_date, end=end_date)
    data.reset_index(inplace=True)  # Reset index to get separate "Date" column

    # Return a different graph type depending on which option the user selected
    match graph_type:
      case 'Line chart of open prices':
        # Create the line plot using Plotly Express
        line_fig = px.line(data, x='Date', y='Open', title='Open Price')
        return line_fig

      case 'Candle chart of stocks':
        candle_fig = go.Figure(data=[go.Candlestick(x=data['Date'],
                open=data['Open'],
                high=data['High'],
                low=data['Low'],
                close=data['Close'])])
        return candle_fig

      case 'Chart of volume traded':
        bar_fig = px.line(data, x = 'Date', y = 'Volume', title = 'Volume Traded')
        return bar_fig

      case 'Pie chart of stock movement':
        pie_fig = get_pie_graph(data)
        return pie_fig

      case 'Chart of trading markers':
        trade_markers_fig = get_trading_markers(data,t, stock_ticker)
        return trade_markers_fig

outputs = [gr.Plot(label="Plot")]

iface = gr.Interface(
    fn=get_stock_data,
    inputs=[
        gr.inputs.Textbox(placeholder="YYYY-MM-DD", label="Start Date"),
        gr.inputs.Textbox(placeholder="YYYY-MM-DD", label="End Date"),
        gr.inputs.Textbox(placeholder="AAPL", label="Stock Ticker"),
        # Dropdown of different choices of graphs to display
        gr.inputs.Dropdown(choices=['Line chart of open prices',
                                    'Candle chart of stocks',
                                    'Chart of volume traded',
                                    'Pie chart of stock movement',
                                    'Chart of trading markers'],
                            label='Graph type'),
        gr.inputs.Number(label="Marker value")

    ],
    outputs=outputs,
    title="Stock Data Viewer",
    description="Enter the start date, end date, and stock ticker to view the stock data.",
)

iface.launch()