import matplotlib.pyplot as plt
import numpy as np
from matplotlib.pyplot import color_sequences
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.stattools import acf, pacf

from ..styling import PLOT_COLORS, PLOT_TITLE_FONTSIZE, TITLE_FONTSIZE, CMAP

def explore_ts(y, trend, seasonal, resid, lags=100):
    """
    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
    """

    plot_decomposition(seasonal, trend, y)

    plot_residual(resid)
    plot_residual_histogram(resid)

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

def plot_decomposition(seasonal, trend, y):
    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"{col}")
    else:
        plt.plot(seasonal, label="Seasonal")

    plt.title("Decomposition", fontsize=PLOT_TITLE_FONTSIZE, fontweight="bold")
    plt.legend()
    plt.show()

def plot_residual(resid):
    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("Residual", fontsize=PLOT_TITLE_FONTSIZE, fontweight="bold")
    plt.legend()
    plt.show()

def plot_residual_histogram(resid):
    plt.figure(figsize=(24, 4))
    plt.hist(resid, bins=50, color=PLOT_COLORS[0], edgecolor="black")
    plt.title("Residual Histogram", fontsize=PLOT_TITLE_FONTSIZE, fontweight="bold")
    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])

    ax.set_title("Residual Autocorrelation", fontsize=PLOT_TITLE_FONTSIZE, fontweight="bold")
    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])

    ax.set_title(f"Partial Residual Autocorrelation ({method})", fontsize=PLOT_TITLE_FONTSIZE, fontweight="bold")
    plt.show()





