import matplotlib.pyplot as plt
import numpy as np
from statsmodels.tsa.stattools import acf, pacf

from ..string_helper_functions import snake_to_title
from ..styling import PLOT_COLORS, plt_title_styled


def explore_ts(y, trend, seasonal, resid, lags=100, plot_decomp_separately = False, x_label = "Frequency", y_label = "Time"):
    """
    Explore the time series decomposition.

    @param y: Original time series
    @param trend: Trend component
    @param seasonal: Seasonal component
    @param resid: Residuals
    @param plot_lags: Number of lags to plot ACF and PACF
    """

    if plot_decomp_separately:
        plot_decomposition_separately(seasonal, trend, y, x_label, y_label)
    else:
        plot_decomposition(seasonal, trend, y, x_label, y_label)

    plot_residual(resid, x_label, y_label)
    plot_residual_histogram(resid, y_label)

    plot_autocorrelation(resid, lags)
    plot_partial_autocorrelation(resid, lags)

def plot_decomposition(seasonal, trend, y, x_label, y_label):
    plt.figure(figsize=(24, 8))
    plt.plot(y, label="Original", color=PLOT_COLORS[2], alpha=0.6)
    plt.plot(trend, label="Trend", color=PLOT_COLORS[1], linewidth=2)

    if hasattr(seasonal, "columns"):
        for col in seasonal.columns:
            plt.plot(seasonal[col], label=f"{snake_to_title(col)}")
    else:
        plt.plot(seasonal, label="Seasonal")

    plt_title_styled("Decomposition", y=1.05)
    plt.legend()
    add_labels(x_label, y_label)
    plt.show()

def plot_decomposition_separately(seasonal, trend, y, x_label, y_label):
    plt.figure(figsize=(24, 8))
    plt.plot(y, label="Original", color=PLOT_COLORS[2], alpha=0.6)
    plt.plot(trend, label="Trend", color=PLOT_COLORS[1], linewidth=2)
    plt_title_styled("Trend")
    plt.legend()
    add_labels(x_label, y_label)
    plt.show()


    if hasattr(seasonal, "columns"):
        for col in seasonal.columns:
            plot_seasonal_component(seasonal[col], snake_to_title(col), x_label, y_label, y)
    else:
        plot_seasonal_component(seasonal, "Seasonal", x_label, y_label, y)

def plot_seasonal_component(component, season_label, x_label, y_label, original = None):
    plt.figure(figsize=(24, 8))
    if original is not None:
        plt.plot(original, label="Original", color=PLOT_COLORS[2], alpha=0.6)
    plt.plot(component, label=f"{season_label}", color=PLOT_COLORS[1], linewidth=2)
    plt_title_styled(f"{season_label}")
    plt.legend()
    add_labels(x_label, y_label)
    plt.show()

def plot_residual(resid, x_label, y_label):
    plt.figure(figsize=(24, 4))
    plt.plot(resid, label="Residual", color=PLOT_COLORS[0])
    plt.hlines(0, resid.index[0], resid.index[-1], colors=PLOT_COLORS[2])
    plt_title_styled("Residual")
    plt.legend()
    add_labels(x_label, y_label)
    plt.show()

def plot_residual_histogram(resid, y_label):
    plt.figure(figsize=(24, 4))
    plt.hist(resid, bins=50, color=PLOT_COLORS[0], edgecolor="black")
    plt_title_styled("Residual Histogram")
    add_labels("Frequency", y_label)
    plt.show()
    plt.close()

def plot_autocorrelation(resid, lags: int):
    acf_vals, confint = acf(resid, nlags=lags, alpha=0.05)
    fig, ax = plt.subplots(figsize=(24,4))

    ax.fill_between(np.arange(len(acf_vals)), confint[:,0]-acf_vals, confint[:,1]-acf_vals,
                    color=PLOT_COLORS[0], alpha=0.2)

    ax.vlines(np.arange(len(acf_vals)), 0, acf_vals, color=PLOT_COLORS[0], lw=2)
    ax.axhline(0, color=PLOT_COLORS[2], lw=2)
    ax.plot(acf_vals, 'o', color=PLOT_COLORS[0])

    plt_title_styled(f"Residual Autocorrelation", ax)
    add_labels("Steps (Lag)", "Autocorrelation", ax)
    plt.show()


def plot_partial_autocorrelation(resid, lags: int, method='yw'):
    pacf_vals, confint = pacf(resid, nlags=lags, alpha=0.05, method=method)

    fig, ax = plt.subplots(figsize=(24, 4))

    ax.fill_between(np.arange(len(pacf_vals)), confint[:, 0] - pacf_vals, confint[:, 1] - pacf_vals,
                    color=PLOT_COLORS[0], alpha=0.2)

    ax.vlines(np.arange(len(pacf_vals)), 0, pacf_vals, color=PLOT_COLORS[0], lw=2)
    ax.axhline(0, color=PLOT_COLORS[2], lw=2)
    ax.plot(pacf_vals, 'o', color=PLOT_COLORS[0])

    plt_title_styled(f"Partial Residual Autocorrelation ({method})", ax)
    add_labels("Steps (Lag)", "Partial Autocorrelation", ax)
    plt.show()

def add_labels(x_label, y_label, ax=None):
    if ax is not None:
        ax.set_xlabel(x_label, fontsize=14, labelpad=10)
        ax.set_ylabel(y_label, fontsize=14, labelpad=4)
    else:
        plt.xlabel("Frequency", fontsize=14, labelpad=10)
        plt.ylabel(y_label, fontsize=14, labelpad=4)