main.py

import os
import numpy as np
import matplotlib.pyplot as plt
import compress_image as ci

def load_dataset(dataset_name):
    """
    Carica il dataset specificato.
    :param dataset_name: nome del dataset
    :return: images, k_values (immagini e valori di k da testare)
    """
    match dataset_name:

        case 'cifar10':
            from keras.datasets import cifar10
            (x_train, _), _ = cifar10.load_data()
            images = x_train[:100].astype(float)
            k_values = [2, 5, 10, 15, 20]

        case 'cifar100':
            from keras.datasets import cifar100
            (x_train, _), _ = cifar100.load_data()
            images = x_train[:100].astype(float)
            k_values = [2, 5, 10, 15, 20]

        case 'skimage':
            from skimage import data
            images = [data.astronaut().astype(float),
                      data.chelsea().astype(float),
                      data.coffee().astype(float)]
            k_values = [10, 25, 50, 75, 100]

        case _:
            raise ValueError(f"Dataset '{dataset_name}' non supportato.")
        
    return images, k_values

def calculate_elbow(k_values, error_values):
    """
    Identifica il valore k ottimale usando il metodo geometrico (distanza massima dalla retta secante).
    Normalizza i dati per evitare che la scala dell'MSE o di k domini il calcolo.
    
    :param k_values: lista dei valori k
    :param error_values: lista dei valori MSE corrispondenti
    :return: (best_k, best_mse, index)
    """

    x = np.array(k_values)
    y = np.array(error_values)

    # 1. Normalizzazione min-max
    x_norm = (x - x.min()) / (x.max() - x.min())
    y_norm = (y - y.min()) / (y.max() - y.min())

    # 2. Definiamo il vettore della retta che congiunge inizio e fine
    p1 = np.array([x_norm[0], y_norm[0]])
    p2 = np.array([x_norm[-1], y_norm[-1]])
    line_vec = p2 - p1
    
    # 3. Calcolo distanza di ogni punto dalla retta
    distances = []
    for i in range(len(x)):
        p_curr = np.array([x_norm[i], y_norm[i]])
        vec_p1_curr = p_curr - p1
        cross_prod = np.cross(line_vec, vec_p1_curr)
        dist = np.abs(cross_prod) / np.linalg.norm(line_vec)
        distances.append(dist)

    # 4. Troviamo l'indice con la distanza massima
    best_idx = np.argmax(distances)
    
    return k_values[best_idx]

def plot_results(dataset, results, best_k_pca, best_k_svd):
    """
    Genera e salva i grafici di confronto tra PCA e SVD per
    MSE, PSNR, SSIM e Compression Ratio (CR).
    """
    
    # --- Estrazione Dati ---
    # Dati PCA
    pca_data = results["pca"]
    k_values = [r[0] for r in pca_data]
    cr_pca = [r[1] for r in pca_data]
    mse_pca = [r[2] for r in pca_data]
    psnr_pca = [r[3] for r in pca_data]
    ssim_pca = [r[4] for r in pca_data]

    # Dati SVD
    svd_data = results["svd"]
    cr_svd = [r[1] for r in svd_data]
    mse_svd = [r[2] for r in svd_data]
    psnr_svd = [r[3] for r in svd_data]
    ssim_svd = [r[4] for r in svd_data]

    # --- Colori Personalizzati ---
    colore_pca = '#1A9988'
    colore_svd = '#EB5600'

    # --- Creazione Grafici (Griglia 2x2) ---
    fig, axs = plt.subplots(2, 2, figsize=(16, 12))
    fig.suptitle('Confronto Compressione Immagini: PCA vs SVD', fontsize=18)

    # 1. Grafico MSE (Minore è meglio)
    axs[0, 0].plot(k_values, mse_pca, marker='o', linestyle='--', label='PCA', color=colore_pca)
    axs[0, 0].plot(k_values, mse_svd, marker='x', linestyle='-', label='SVD', color=colore_svd)
    axs[0, 0].scatter(best_k_pca, mse_pca[k_values.index(best_k_pca)], color='red', s=150, zorder=5, edgecolors='black', label=f'Best PCA (k={best_k_pca})')
    axs[0, 0].scatter(best_k_svd, mse_svd[k_values.index(best_k_svd)], color='red', s=150, marker='X', zorder=5, edgecolors='black', label=f'Best SVD (k={best_k_svd})')
    axs[0, 0].set_title('Mean Squared Error (MSE)')
    axs[0, 0].set_xlabel('k (Numero Componenti)')
    axs[0, 0].set_ylabel('MSE (Medio)')
    axs[0, 0].legend()
    axs[0, 0].grid(True, linestyle=':', alpha=0.7)

    # 2. Grafico PSNR (Maggiore è meglio)
    axs[0, 1].plot(k_values, psnr_pca, marker='o', linestyle='--', label='PCA', color=colore_pca)
    axs[0, 1].plot(k_values, psnr_svd, marker='x', linestyle='-', label='SVD', color=colore_svd)
    axs[0, 1].set_title('Peak Signal-to-Noise Ratio (PSNR)')
    axs[0, 1].set_xlabel('k (Numero Componenti)')
    axs[0, 1].set_ylabel('PSNR (dB)')
    axs[0, 1].legend()
    axs[0, 1].grid(True, linestyle=':', alpha=0.7)

    # 3. Grafico SSIM (Maggiore è meglio, max 1.0)
    axs[1, 0].plot(k_values, ssim_pca, marker='o', linestyle='--', label='PCA', color=colore_pca)
    axs[1, 0].plot(k_values, ssim_svd, marker='x', linestyle='-', label='SVD', color=colore_svd)
    axs[1, 0].set_title('Structural Similarity Index (SSIM)')
    axs[1, 0].set_xlabel('k (Numero Componenti)')
    axs[1, 0].set_ylabel('SSIM (Medio)')
    axs[1, 0].legend()
    axs[1, 0].grid(True, linestyle=':', alpha=0.7)

    # 4. Grafico Compression Ratio (CR)
    axs[1, 1].plot(k_values, cr_pca, marker='o', linestyle='--', label='PCA', color=colore_pca)
    axs[1, 1].plot(k_values, cr_svd, marker='x', linestyle='-', label='SVD', color=colore_svd)
    axs[1, 1].set_title('Compression Ratio (CR)')
    axs[1, 1].set_xlabel('k (Numero Componenti)')
    axs[1, 1].set_ylabel('CR (Medio)')
    axs[1, 1].legend()
    axs[1, 1].grid(True, linestyle=':', alpha=0.7)

    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    
    # Salva il grafico come immagine PNG
    output_filename = f'compressione_{dataset}.png'
    plt.savefig(output_filename)
    
    print(f"\nGrafici di confronto salvati come '{output_filename}'")

def test(k_values):
    """
    Esegue la compressione per una lista di k e salva i risultati su disco.
    """
    from matplotlib.image import imread

    image_path = "lamp_ORIGINAL.dng"
    output_dir = "risultati_compressione"
    results = {"pca": [], "svd": []}
    
    # Crea la cartella di output se non esiste
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Cartella '{output_dir}' creata.")

    # Lettura immagine originale
    try:
        original = imread(image_path).astype(float)
    except FileNotFoundError:
        print(f"Errore: Immagine '{image_path}' non trovata.")
        return

    print(f"--- Inizio elaborazione e salvataggio per {len(k_values)} valori di k ---\n")

    for k in k_values:
        print(f"Sto elaborando k = {k}...")

        # --- 1. Compressione PCA ---
        compressed_img_pca, cr_pca = ci.compress_image('pca', original, k)
        img_to_save_pca = np.clip(compressed_img_pca, 0, 255).astype('uint8')
        filename_pca = os.path.join(output_dir, f"comp_pca_k{k}.jpg")
        plt.imsave(filename_pca, img_to_save_pca)
        mse_pca, psnr_pca, ssim_pca = ci.evaluate_compression(original, compressed_img_pca)
        print(f"   -> PCA: MSE={mse_pca:.2f}, PSNR={psnr_pca:.2f}, SSIM={ssim_pca:.4f}")
        results["pca"].append((k, cr_pca, mse_pca, psnr_pca, ssim_pca))

        # --- 2. Compressione SVD ---
        compressed_img_svd, cr_svd = ci.compress_image('svd', original, k)
        img_to_save_svd = np.clip(compressed_img_svd, 0, 255).astype('uint8')
        filename_svd = os.path.join(output_dir, f"comp_svd_k{k}.jpg")
        plt.imsave(filename_svd, img_to_save_svd)
        mse_svd, psnr_svd, ssim_svd = ci.evaluate_compression(original, compressed_img_svd)
        print(f"   -> SVD: MSE={mse_svd:.2f}, PSNR={psnr_svd:.2f}, SSIM={ssim_svd:.4f}")
        results["svd"].append((k, cr_svd, mse_svd, psnr_svd, ssim_svd))

        # Stampa rapida info
        print(f"   -> Salvato: {filename_pca} (Ratio: {cr_pca:.2f}x)")
        print(f"   -> Salvato: {filename_svd} (Ratio: {cr_svd:.2f}x)")
        print("-" * 30)

    print(f"\Tutte le immagini sono nella cartella '{output_dir}'.")

    best_k_pca = calculate_elbow(k_values, [r[2] for r in results["pca"]])
    best_k_svd = calculate_elbow(k_values, [r[2] for r in results["svd"]])

    return results, best_k_pca, best_k_svd

def main(dataset_name):

    images, k_values = load_dataset(dataset_name)
    results = {"pca": [], "svd": []}
    
    for k in k_values:
        print(f"k={k} -------------------------")
        mse_pca_list, psnr_pca_list, ssim_pca_list = [], [], []
        mse_svd_list, psnr_svd_list, ssim_svd_list = [], [], []

        for img in images:
            # PCA
            comp_pca, cr_pca = ci.compress_image("pca", img, k)
            mse, psnr, ssim_val = ci.evaluate_compression(img, comp_pca)
            mse_pca_list.append(mse)
            psnr_pca_list.append(psnr)
            ssim_pca_list.append(ssim_val)
            # SVD
            comp_svd, cr_svd = ci.compress_image("svd", img, k)
            mse, psnr, ssim_val = ci.evaluate_compression(img, comp_svd)
            mse_svd_list.append(mse)
            psnr_svd_list.append(psnr)
            ssim_svd_list.append(ssim_val)
    
        results["pca"].append((k, cr_pca, np.mean(mse_pca_list), np.mean(psnr_pca_list), np.mean(ssim_pca_list)))
        results["svd"].append((k, cr_svd, np.mean(mse_svd_list), np.mean(psnr_svd_list), np.mean(ssim_svd_list)))

    # Identificazione del miglior k usando il metodo del gomito
    mse_values_pca = [r[2] for r in results["pca"]]
    mse_values_svd = [r[2] for r in results["svd"]]
    best_k_pca = calculate_elbow(k_values, mse_values_pca)
    best_k_svd = calculate_elbow(k_values, mse_values_svd)

    # Stampa risultati
    print("\n=== RISULTATI MEDI ===")
    print("\n--- PCA ---")
    for r in results["pca"]:
        print(f"k={r[0]} | CR={r[1]:.2f} | MSE={r[2]:.2f} | PSNR={r[3]:.2f} | SSIM={r[4]:.4f}")
    print(f"\nMiglior k PCA: {best_k_pca}")
    print("\n--- SVD ---")
    for r in results["svd"]:
        print(f"k={r[0]} | CR={r[1]:.2f} | MSE={r[2]:.2f} | PSNR={r[3]:.2f} | SSIM={r[4]:.4f}")
    print(f"\nMiglior k SVD: {best_k_svd}")

    return results, best_k_pca, best_k_svd

if __name__ == "__main__":

    os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
    dataset = 'skimage'

    risultati_finali, best_k_pca, best_k_svd = test([10, 20, 50, 100, 200])
    plot_results(dataset, risultati_finali, best_k_pca, best_k_svd)