import subprocess
import sys
import os

# --- Funzione per installare librerie mancanti ---
def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

required_packages = ["torch", "torchvision", "compressai", "Pillow", "numpy", "scikit-image", "matplotlib"]
for pkg in required_packages:
    try:
        __import__(pkg)
    except ImportError:
        print(f"Libreria {pkg} mancante, installazione in corso...")
        install(pkg)

# --- Import principali ---
import torch
from compressai.zoo import bmshj2018_factorized
from PIL import Image
from torchvision import transforms
from torchvision.transforms.functional import to_pil_image
import numpy as np
from skimage.metrics import peak_signal_noise_ratio as psnr, structural_similarity as ssim
import matplotlib.pyplot as plt

# --- Configurazione ---
image_path = "immagine.png"  # Percorso della tua immagine
quality = 4                  # Qualità compressione IA (1=alta compressione, 10=bassa compressione)
jpeg_quality = 40            # Qualità JPEG per confronto

# --- Caricamento immagine originale ---
img = Image.open(image_path).convert("RGB")
W, H = img.size

# Dimensione non compressa (RGB a 8 bit = 3 byte per pixel)
size_original = W * H * 3  

# --- Compressione JPEG ---
img.save("img_jpeg.jpg", quality=jpeg_quality)
img_jpeg = Image.open("img_jpeg.jpg")
size_jpeg = os.path.getsize("img_jpeg.jpg")

# --- Compressione IA ---
device = torch.device("cpu")
model = bmshj2018_factorized(quality=quality, pretrained=True).to(device).eval()

transform = transforms.ToTensor()
img_tensor = transform(img).unsqueeze(0).to(device)

with torch.no_grad():
    compressed = model.compress(img_tensor)
    decompressed = model.decompress(**compressed)["x_hat"]

img_ai = to_pil_image(decompressed.squeeze())

# Calcolo dimensione bitstream IA (in byte)
size_ai = sum(len(s) for s in compressed["strings"][0])  

# --- Ridimensionamento per confronto visivo ---
img_jpeg = img_jpeg.resize(img.size, Image.BICUBIC)
img_ai = img_ai.resize(img.size, Image.BICUBIC)

# --- Mostra immagini ---
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
axes[0].imshow(img); axes[0].set_title(f"Originale\n{size_original/1024:.1f} KB"); axes[0].axis('off')
axes[1].imshow(img_jpeg); axes[1].set_title(f"JPEG Q={jpeg_quality}\n{size_jpeg/1024:.1f} KB"); axes[1].axis('off')
axes[2].imshow(img_ai); axes[2].set_title(f"IA Q={quality}\n{size_ai/1024:.1f} KB stimati"); axes[2].axis('off')
plt.show()

# --- Calcolo metriche ---
img_np = np.array(img)
img_jpeg_np = np.array(img_jpeg)
img_ai_np = np.array(img_ai)

print("PSNR JPEG:", psnr(img_np, img_jpeg_np))
print("PSNR IA:", psnr(img_np, img_ai_np))
print("SSIM JPEG:", ssim(img_np, img_jpeg_np, channel_axis=-1, win_size=3))
print("SSIM IA:", ssim(img_np, img_ai_np, channel_axis=-1, win_size=3))

print(f"\nDimensione stimata immagine non compressa: {size_original/1024:.1f} KB")
print(f"Dimensione file JPEG: {size_jpeg/1024:.1f} KB")
print(f"Dimensione stream AI: {size_ai/1024:.1f} KB (stimati dai bitstream CompressAI)")
