import pandas as pd
import numpy as np
import plotly.graph_objects as go
from datetime import timedelta
from statsmodels.tsa.arima.model import ARIMA
from config import FORECAST_PERIOD, ticker_dict, CONFIDENCE_INTERVAL
from data_fetcher import get_stock_data, get_company_info

def is_business_day(a_date):
    return a_date.weekday() < 5

def forecast_series(series, model="ARIMA", forecast_period=FORECAST_PERIOD):
    predictions = []
    confidence_intervals = []
    
    if series.shape[1] > 1:
        series = series['Close'].values.tolist()
    
    if model == "ARIMA":
        model = ARIMA(series, order=(5, 1, 0))
        model_fit = model.fit()
        forecast = model_fit.forecast(steps=forecast_period, alpha=(1 - CONFIDENCE_INTERVAL))
        predictions = forecast.predicted_mean
        confidence_intervals = forecast.conf_int()
    elif model == "Prophet":
        # Implement Prophet forecasting method
        pass
    elif model == "LSTM":
        # Implement LSTM forecasting method
        pass

    return predictions, confidence_intervals

def get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method, start_date, end_date):
    stock_name, ticker_name = stock.split(":")
    
    if ticker_dict[idx] == 'FTSE 100':
        ticker_name += '.L' if ticker_name[-1] != '.' else 'L'
    elif ticker_dict[idx] == 'CAC 40':
        ticker_name += '.PA'

    series = get_stock_data(ticker_name, interval, start_date, end_date)
    predictions, confidence_intervals = forecast_series(series, model=forecast_method)

    last_date = pd.to_datetime(series['Date'].values[-1])
    forecast_dates = pd.date_range(start=last_date + timedelta(days=1), periods=FORECAST_PERIOD)
    forecast_dates = [date for date in forecast_dates if is_business_day(date)]

    forecast = pd.DataFrame({
        "Date": forecast_dates,
        "Forecast": predictions,
        "Lower_CI": confidence_intervals.iloc[:, 0],
        "Upper_CI": confidence_intervals.iloc[:, 1]
    })

    if graph_type == 'Line Graph':
        fig = go.Figure()
        fig.add_trace(go.Scatter(x=series['Date'], y=series['Close'], mode='lines', name='Historical'))
        fig.add_trace(go.Scatter(x=forecast['Date'], y=forecast['Forecast'], mode='lines', name='Forecast'))
        fig.add_trace(go.Scatter(
            x=forecast['Date'].tolist() + forecast['Date'].tolist()[::-1],
            y=forecast['Upper_CI'].tolist() + forecast['Lower_CI'].tolist()[::-1],
            fill='toself',
            fillcolor='rgba(0,100,80,0.2)',
            line=dict(color='rgba(255,255,255,0)'),
            hoverinfo="skip",
            showlegend=False
        ))
    else:  # Candlestick Graph
        fig = go.Figure(data=[go.Candlestick(x=series['Date'],
                                             open=series['Open'],
                                             high=series['High'],
                                             low=series['Low'],
                                             close=series['Close'],
                                             name='Historical')])
        fig.add_trace(go.Scatter(x=forecast['Date'], y=forecast['Forecast'], mode='lines', name='Forecast'))
        fig.add_trace(go.Scatter(
            x=forecast['Date'].tolist() + forecast['Date'].tolist()[::-1],
            y=forecast['Upper_CI'].tolist() + forecast['Lower_CI'].tolist()[::-1],
            fill='toself',
            fillcolor='rgba(0,100,80,0.2)',
            line=dict(color='rgba(255,255,255,0)'),
            hoverinfo="skip",
            showlegend=False
        ))

    fig.update_layout(title=f"Stock Price of {stock_name}",
                      xaxis_title="Date",
                      yaxis_title="Price")

    fundamentals = get_company_info(ticker_name)

    return fig, fundamentals