##############################
# NE PAS MODIFIER CE FICHIER #
##############################
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import convolve2d, get_window


tol = 1e-3  # Threshold to assume 0

def plot_xn_1(n: np.ndarray, xn: np.ndarray) -> None:

    _, ax = plt.subplots(nrows=1, ncols=1, figsize=(15.0, 5.0))

    ax.plot(n, xn, label='signal')
    ax.set_title(r'signal $x(n)$')
    ax.set_xlabel('#Sample')
    ax.set_ylabel(r'$x(n)$')
    ax.legend(loc=1)
    plt.show()

def plot_xwn_dft_1(xw, xwf, xwf_shift, wn, fs, df, n_pts, win_type: str) -> None:
    """ Plot windowed signal, its DFT and centered DFT"""
    
    # Visualization
    _, axes = plt.subplots(nrows=3, ncols=1, figsize=(15.0, 8.0))

    axes[0].plot(xw, label='windowed signal')
    axes[0].plot(wn, label='window {}'.format(win_type))
    axes[0].set_title(r'Windowed signal $x_w(n) = w(n) \cdot x(n)$')
    axes[0].set_xlabel('#Sample')
    axes[0].set_ylabel(r'$x_w(n)$')
    axes[0].legend()

    # DFT
    axes[1].stem(np.abs(xwf))
    axes[1].set_title(r'DFT of $x_w(n)$, $\Delta f$ ${:.3f}$ Hz'.format(fs / n_pts))
    axes[1].set_xlabel('Frequency index')
    axes[1].set_ylabel(r'Magnitude $|X_w(f)|$')

    # DFT
    axes[2].stem(df, np.abs(xwf_shift))
    axes[2].set_title(r'Centered DFT of $x_w(n)$'.format(fs / n_pts))
    axes[2].set_xlabel('Frequency [Hz]')
    axes[2].set_ylabel(r'Magnitude $|X_w(f)|$')
    
    plt.tight_layout()
    
    n_nz = np.sum(np.abs(xwf) > tol)
    print("Non-zero terms({}, N={}): {}".format(win_type, n_pts, n_nz))
    
def plot_gkernel(h: np.ndarray, var: float) -> None:
    """
    Plot a 2D kernel (image) with imshow and optional colorbar.

    Parameters
    ----------
    h : np.ndarray (N1, N2)
        Kernel to plot.
    var: float
        Variance of the Gaussian kernel.
    """
    fig, ax = plt.subplots(1, 1, figsize=(5, 5))

    im = ax.imshow(h, cmap='jet', interpolation='none')
    ax.set_title(r'Gaussian kernel with $\sigma^2 = {}$'.format(var))
    ax.set_xlabel('m')
    ax.set_ylabel('n')
    plt.colorbar(im, ax=ax, shrink=0.8)
    plt.tight_layout()
    
def plot_2_1(vt: np.ndarray) -> None:
    fig, ax = plt.subplots(1, 1, figsize=(8, 8))
    ax.imshow(vt, cmap = 'gray', interpolation=None)
    ax.set_title(r'Output - $v_t$ (temporal)')
    ax.axis('off')
    plt.show()
    
def plot_2_2(vt, vf, Y, H, V):
    """
    Trace tous les résultats du filtrage en domaine fréquentiel :
      - spectre de Y, H et V en log-échelle
      - comparaison temporelle vs fréquentielle (vt vs vf)

    Parameters
    ----------
    vt : np.ndarray
        Résultat convolution dans le domaine temporel
    vf : np.ndarray
        Résultat convolution (ifft du domaine fréquentiel)
    Y, H, V : np.ndarray
        TF 2D centrées de y, h et leur produit
    """
    
    # ---- Spectres ----
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))
    im0 = axs[0].imshow(np.log(np.abs(Y)+1e-12),
                        cmap='jet', origin='lower')
    axs[0].set_title(r'$|Y|$ (log)')
    axs[0].axis('off')
    plt.colorbar(im0, ax=axs[0], shrink=0.7)

    im1 = axs[1].imshow(np.log(np.abs(H)+1e-12),
                        cmap='jet', origin='lower')
    axs[1].set_title(r'$|H|$ (log)')
    axs[1].axis('off')
    plt.colorbar(im1, ax=axs[1], shrink=0.7)

    im2 = axs[2].imshow(np.log(np.abs(V)+1e-12),
                        cmap='jet', origin='lower')
    axs[2].set_title(r'$|V|=|Y\cdot H|$ (log)')
    axs[2].axis('off')
    plt.colorbar(im2, ax=axs[2], shrink=0.7)

    plt.tight_layout()
    plt.show()

    # ---- Comparaison temporel vs fréquentiel ----
    fig, axs = plt.subplots(1, 2, figsize=(12, 6))
    axs[0].imshow(vt, cmap='gray')
    axs[0].set_title(r'Time domain result $v_t$')
    axs[0].axis('off')

    axs[1].imshow(np.abs(vf), cmap='gray')
    axs[1].set_title(r'Frequency domain result $v_f$')
    axs[1].axis('off')

    plt.tight_layout()
    plt.show()
    
def plot_2_2_2(Y, H, V, vt, vf):
    """
    Trace les spectres de Fourier (Y, H, V) et compare
    les résultats temporel (vt) et fréquentiel (vf).

    Parameters
    ----------
    Y : np.ndarray
        TF 2D de l'entrée y
    H : np.ndarray
        TF 2D du filtre h
    V : np.ndarray
        Produit fréquentiel Y*H
    vt : np.ndarray
        Résultat convolution temporelle
    vf : np.ndarray
        Résultat convolution fréquentielle (ifft de V)
    """

    # --- Spectres ---
    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
    im_Y = ax[0].imshow(np.log(np.abs(Y)), cmap='jet')
    ax[0].axis('off')
    ax[0].set_title(r'$Y$')
    plt.colorbar(mappable=im_Y, ax=ax[0], shrink=0.7)

    im_H = ax[1].imshow(np.log(np.abs(H)), cmap='jet')
    ax[1].axis('off')
    ax[1].set_title(r'$H$')
    plt.colorbar(mappable=im_H, ax=ax[1], shrink=0.7)

    im_V = ax[2].imshow(np.log(np.abs(V)), cmap='jet')
    ax[2].axis('off')
    ax[2].set_title(r'$V$')
    plt.colorbar(mappable=im_V, ax=ax[2], shrink=0.7)

    plt.show()

    # --- Résultats temporel vs fréquentiel ---
    fig, ax = plt.subplots(1, 2, figsize=(12, 8))
    ax[0].imshow(vt, cmap='gray')
    ax[0].set_title(r'Time domain result - $v_t$')
    ax[0].axis('off')

    ax[1].imshow(np.abs(vf), cmap='gray')
    ax[1].set_title(r'Frequency domain result - $v_f$')
    ax[1].axis('off')

    plt.show()