import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.lines import Line2D
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.cm import ScalarMappable

def visualize_results(y_hat, q_hat, ylabel):
    y_hat = pd.Series(y_hat.flatten())
    taqr_results_corrected_plot = pd.DataFrame(np.array(q_hat).T, index=y_hat.index)
    m = taqr_results_corrected_plot.shape[1]  # ensemble size

    # Define the color gradient from dark blue to light cyan
    colors = [
        (173/255, 217/255, 229/255), 
        (19/255, 25/255, 148/255),  
        (173/255, 217/255, 229/255), 
    ]
    cmap = LinearSegmentedColormap.from_list("blue_to_cyan", colors, N=100)
    norm = plt.Normalize(vmin=0, vmax=m-2)  # Normalize for the ensemble size
    sm = ScalarMappable(cmap=cmap, norm=norm)

    plt.figure(figsize=(6, 4))
    for i in range(m-1):
        color = sm.to_rgba(i)
        plt.fill_between(taqr_results_corrected_plot.index, taqr_results_corrected_plot.iloc[:, i], taqr_results_corrected_plot.iloc[:, i+1], color=color, alpha=1)
    plt.plot(y_hat.index, y_hat, color="white", linewidth=3)  # White outline
    plt.plot(y_hat, color="black", label="Actuals", linewidth=1)  # Actual line
    plt.xlim(y_hat.index[0], y_hat.index[-1])
    line = Line2D([0], [0], color="black", lw=2, label="Actuals")
    contour = Line2D([0], [0], color=sm.to_rgba(m//2), lw=5, alpha=0.9, label="Prediction Interval")
    plt.legend(handles=[line, contour])
    plt.xlabel("Time")
    plt.ylabel(ylabel)
    locator = mdates.AutoDateLocator(minticks=6, maxticks=8)
    formatter = mdates.ConciseDateFormatter(locator)
    ax = plt.gca()
    ax.xaxis.set_major_locator(locator)
    ax.xaxis.set_major_formatter(formatter)
    plt.tight_layout()
    plt.savefig("TEST_NABQR_taqr_pi_plot.pdf")
    plt.show()
