# -*- coding: utf-8 -*-
## Étape 6 : Chargement et Test du Modèle depuis Hugging Face
# -----------------------------------------------------------------------------

import os
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from huggingface_hub import HfApi, hf_hub_download
import matplotlib.pyplot as plt
import json

print("\n--- Chargement et Test du Modèle ---")

# Configuration pour le test (doit correspondre à celle de l'entraînement)
HUGGINGFACE_REPO_ID = "Clemylia/Imurka"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TEMPERATURE = 8.7 # Nouvelle variable pour contrôler la température de génération, ajustée pour un effet plus visible

# --- 1. Télécharger la configuration et le modèle depuis Hugging Face ---
print(f"Téléchargement de la configuration depuis {HUGGINGFACE_REPO_ID}...")
config_path = hf_hub_download(repo_id=HUGGINGFACE_REPO_ID, filename="config.json", repo_type="model")
with open(config_path, 'r') as f:
    config = json.load(f)

LATENT_DIM = config['latent_dim']
IMAGE_SIZE = config['image_size']
CHANNELS = config['channels']

print(f"Configuration chargée : LATENT_DIM={LATENT_DIM}, IMAGE_SIZE={IMAGE_SIZE}, CHANNELS={CHANNELS}")

# --- 2. Redéfinir l'architecture du modèle (doit être identique à l'entraînement) ---
class Encoder(nn.Module):
    def __init__(self, latent_dim=LATENT_DIM):
        super().__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(CHANNELS, 32, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
            nn.Flatten()
        )
        # Calculer la taille de l'entrée linéaire après les convolutions
        # Pour une image de IMAGE_SIZE x IMAGE_SIZE, après 5 couches (stride=2, padding=1):
        # 128 -> 64 -> 32 -> 16 -> 8 -> 4
        self.fc_mu = nn.Linear(512 * 4 * 4, latent_dim)
        self.fc_log_var = nn.Linear(512 * 4 * 4, latent_dim)

    def forward(self, x):
        x = self.conv_layers(x)
        mu = self.fc_mu(x)
        log_var = self.fc_log_var(x)
        return mu, log_var

class Decoder(nn.Module):
    def __init__(self, latent_dim=LATENT_DIM):
        super().__init__()
        self.fc = nn.Linear(latent_dim, 512 * 4 * 4)

        self.deconv_layers = nn.Sequential(
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, CHANNELS, kernel_size=4, stride=2, padding=1),
            nn.Sigmoid()
        )

    def forward(self, z):
        x = self.fc(z)
        x = x.view(-1, 512, 4, 4) # Redimensionner pour les couches de déconvolution
        reconstruction = self.deconv_layers(x)
        return reconstruction

class VAE(nn.Module):
    def __init__(self, latent_dim=LATENT_DIM):
        super().__init__()
        self.encoder = Encoder(latent_dim)
        self.decoder = Decoder(latent_dim)

    def reparameterize(self, mu, log_var, temperature=1.0): # Ajout du paramètre de température
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std * temperature # Application de la température

    def forward(self, x, temperature=1.0): # Ajout du paramètre de température
        mu, log_var = self.encoder(x)
        z = self.reparameterize(mu, log_var, temperature) # Passage de la température à reparameterize
        reconstruction = self.decoder(z)
        return reconstruction, mu, log_var

# --- 3. Instancier le modèle et charger les poids ---
model = VAE(latent_dim=LATENT_DIM).to(DEVICE)

print(f"Téléchargement du modèle binaire depuis {HUGGINGFACE_REPO_ID}...")
model_bin_path = hf_hub_download(repo_id=HUGGINGFACE_REPO_ID, filename="pytorch_model.bin", repo_type="model")

model.load_state_dict(torch.load(model_bin_path, map_location=DEVICE))
model.eval() # Mettre le modèle en mode évaluation
print("Modèle VAE chargé avec succès.")

# --- 4. Préparer une image de test ---
# Utiliser l'image téléchargée par l'utilisateur
user_image_path = "/content/lae.jpg"

# Transformation pour les images de test (doit être la même que pour l'entraînement)
test_transform = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.CenterCrop(IMAGE_SIZE),
    transforms.ToTensor(),
])

print(f"Chargement de l'image de test depuis : {user_image_path}...")
try:
    # Charger l'image utilisateur
    display_original_pil = Image.open(user_image_path).convert("RGB")
    test_image_tensor = test_transform(display_original_pil).unsqueeze(0).to(DEVICE)
except FileNotFoundError:
    print(f"Erreur: Le fichier {user_image_path} n'a pas été trouvé. Veuillez vérifier le chemin.")
    # Créer une image noire si le fichier n'est pas trouvé pour éviter de planter
    display_original_pil = Image.new('RGB', (IMAGE_SIZE, IMAGE_SIZE), color = 'black')
    test_image_tensor = test_transform(display_original_pil).unsqueeze(0).to(DEVICE)
except Exception as e:
    print(f"Erreur lors du chargement de l'image de test: {e}")
    display_original_pil = Image.new('RGB', (IMAGE_SIZE, IMAGE_SIZE), color = 'black')
    test_image_tensor = test_transform(display_original_pil).unsqueeze(0).to(DEVICE)


# --- 5. Effectuer l'inférence ---
print("Exécution de l'inférence...")
with torch.no_grad(): # Pas besoin de calculer les gradients pour l'inférence
    # Passage du paramètre de température au modèle
    reconstructed_image_tensor, _, _ = model(test_image_tensor, temperature=TEMPERATURE)

# --- 6. Post-traitement et Visualisation ---
print("Visualisation des résultats...")

# Convertir le tenseur de sortie en image PIL
reconstructed_image_tensor = reconstructed_image_tensor.squeeze(0).cpu()
reconstructed_image_pil = transforms.ToPILImage()(reconstructed_image_tensor)

plt.figure(figsize=(10, 5))

plt.subplot(1, 2, 1)
plt.imshow(display_original_pil)
plt.title("Image d'Entrée")
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(reconstructed_image_pil)
plt.title("Image Générée par le VAE (Modifiée)")
plt.axis('off')

plt.show()

print("\n✅ Test du modèle terminé. Vous devriez voir l'image originale et sa reconstruction.")