import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
import io
import base64
from statsmodels.tsa.holtwinters import ExponentialSmoothing
import pmdarima as pm
from prophet import Prophet


def create_acf_pacf_plots(data):
    # Create ACF and PACF plots using matplotlib
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

    plot_acf(data, ax=ax1, lags=40)
    ax1.set_title('Autocorrelation Function')

    plot_pacf(data, ax=ax2, lags=40)
    ax2.set_title('Partial Autocorrelation Function')

    # Convert matplotlib plot to Plotly
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close(fig)
    buf.seek(0)
    img_str = base64.b64encode(buf.getvalue()).decode('utf-8')

    # Create Plotly figure with image
    fig_plotly = go.Figure()
    fig_plotly.add_layout_image(
        dict(
            source=f'data:image/png;base64,{img_str}',
            x=0,
            y=1,
            xref="paper",
            yref="paper",
            sizex=1,
            sizey=1,
            sizing="stretch",
            opacity=1,
            layer="below"
        )
    )
    fig_plotly.update_layout(
        height=600,
        showlegend=False,
        xaxis=dict(visible=False),
        yaxis=dict(visible=False)
    )
    return pio.to_html(fig_plotly, full_html=False)


def create_comparison_plot(filepath, forecast_history, selected_indices):
    # Read data
    if filepath.endswith('.csv'):
        df = pd.read_csv(filepath)
    else:
        df = pd.read_excel(filepath)

    date_col = df.columns[0]
    value_col = df.columns[1]
    df[date_col] = pd.to_datetime(df[date_col])
    df.set_index(date_col, inplace=True)

    # Create Plotly figure
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df.index, y=df[value_col], name='Historical', line=dict(color='black')))

    # Use Plotly qualitative colors
    colors = px.colors.qualitative.Plotly

    # Generate forecasts for selected indices
    for idx, run_idx in enumerate(selected_indices):
        entry = forecast_history[run_idx]
        train_percent = entry['train_percent'] / 100
        forecast_periods = entry['forecast_periods']
        model_type = entry['model_type']

        # Split data
        train_size = int(len(df) * train_percent)
        test_size = len(df) - train_size
        train_data = df[value_col].iloc[:train_size]
        test_data = df[value_col].iloc[train_size:] if test_size > 0 else pd.Series()
        forecast_dates = pd.date_range(start=df.index[-1], periods=forecast_periods + 1, freq=df.index.inferred_freq)[
                         1:]

        # Run model based on model_type
        forecast = None
        if model_type == 'ARIMA':
            model = pm.auto_arima(train_data,
                                  seasonal=True,
                                  m=12,
                                  start_p=0, start_q=0,
                                  max_p=3, max_q=3,
                                  start_P=0, start_Q=0,
                                  max_P=2, max_Q=2,
                                  d=1, D=1,
                                  trace=False,
                                  error_action='ignore',
                                  suppress_warnings=True,
                                  stepwise=True)
            model_fit = model.fit(train_data)
            forecast = model_fit.predict(n_periods=forecast_periods)

        elif model_type == 'Exponential Smoothing':
            model = ExponentialSmoothing(train_data,
                                         trend='add',
                                         seasonal='add',
                                         seasonal_periods=12)
            model_fit = model.fit()
            forecast = model_fit.forecast(forecast_periods)

        elif model_type == 'Prophet':
            prophet_df = train_data.reset_index().rename(columns={date_col: 'ds', value_col: 'y'})
            model = Prophet(yearly_seasonality=True, weekly_seasonality=False, daily_seasonality=False)
            model.add_seasonality(name='monthly', period=30.5, fourier_order=5)
            model_fit = model.fit(prophet_df)
            future = model.make_future_dataframe(periods=forecast_periods, freq=df.index.inferred_freq)
            forecast_full = model_fit.predict(future)
            forecast = forecast_full['yhat'].iloc[-forecast_periods:].values

        # Add test data if available (only once to avoid clutter)
        if test_size > 0 and idx == 0:
            fig.add_trace(go.Scatter(x=df.index[train_size:], y=test_data, name='Test Data', line=dict(color='green')))

        # Add forecast
        label = f"Forecast Run {run_idx + 1}: {model_type}, {entry['train_percent']:.0f}/{entry['test_percent']:.0f}, {forecast_periods} periods"
        fig.add_trace(go.Scatter(x=forecast_dates, y=forecast, name=label,
                                 line=dict(dash='dash', color=colors[idx % len(colors)])))

    fig.update_layout(title='Forecast Comparison', height=400, showlegend=True)
    return pio.to_html(fig, full_html=False)