import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import tensorflow as tf
import numpy as np
import cv2
from tensorflow.keras.models import load_model
import numpy as np
import cv2
from tensorflow.keras.models import load_model
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.utils import get_custom_objects
from torch.autograd import Variable
import torch
import torchvision.transforms as transforms

from models import UNet


def mse(y_true, y_pred):
    return MeanSquaredError()(y_true, y_pred)

get_custom_objects().update({'mse': mse})

# Chargement des modèles
face_aging_model = load_model("face_aging_model.h5", custom_objects={"mse": mse})

device = "cuda" if torch.cuda.is_available() else "cpu"
face_aging_autoencoder = UNet().to(device)
face_aging_autoencoder.load_state_dict(torch.load("best_unet_model_dl.pth", map_location=device))


def load_image():
    file_path = filedialog.askopenfilename(title="Sélectionner une image", filetypes=[("Image Files", "*.jpg;*.png;*.jpeg")])
    if file_path:
        img = Image.open(file_path)
        img = img.resize((256, 256))  # Redimensionner l'image
        img_tk = ImageTk.PhotoImage(img)
        original_panel.configure(image=img_tk)
        original_panel.image = img_tk
        global image_data, image_path
        image_path = file_path
        image_data = np.array(img) / 255.0  # Normalisation de l'image
        image_data = np.expand_dims(image_data, axis=0) 

def predict_age():
    if 'image_path' in globals():
        predicted_age = predict_age_from_model(face_aging_model, image_path)
        messagebox.showinfo("Prédiction d'âge", f"L'âge prédit du visage est : {predicted_age:.2f} ans")
    else:
        messagebox.showerror("Erreur", "Veuillez d'abord charger une image.")

def predict_age_from_model(model, image_path):
    img = Image.open(image_path).resize((128, 128))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    prediction = model.predict(img_array)
    return prediction[0][0]

def sliding_window_tensor(input_tensor, your_model):
    window_size,stride=512,256
    input_tensor = input_tensor.to(next(your_model.parameters()).device)

    n, c, h, w = input_tensor.size()
    output_tensor = torch.zeros((n, 3, h, w), dtype=input_tensor.dtype, device=input_tensor.device)
    count_tensor = torch.zeros((n, 3, h, w), dtype=torch.float32, device=input_tensor.device)


    for y in range(0, h - window_size + 1, stride):
        for x in range(0, w - window_size + 1, stride):
            window = input_tensor[:, :, y:y + window_size, x:x + window_size]

            input_variable = Variable(window, requires_grad=False) 
            with torch.no_grad():
                output = your_model(input_variable)

            output_tensor[:, :, y:y + window_size, x:x + window_size] += output
            count_tensor[:, :, y:y + window_size, x:x + window_size] += 1

    count_tensor = torch.clamp(count_tensor, min=1.0)

    output_tensor /= count_tensor

    return output_tensor.cpu()

def process_image(your_model, image, source_age, target_age=80, ):

    image = np.array(image)
    image_original = image.copy()

    image = transforms.ToTensor()(image)

    source_age_channel = torch.full_like(image[:1, :, :], source_age / 100)
    target_age_channel = torch.full_like(image[:1, :, :], target_age / 100)
    input_tensor = torch.cat([image, source_age_channel, target_age_channel], dim=0).unsqueeze(0)

    image_original = transforms.ToTensor()(image_original)

    # performing actions on image
    aged_image = sliding_window_tensor(input_tensor, your_model)

    image_original += aged_image.squeeze(0)
    image_original = torch.clamp(image_original, 0, 1)

    return transforms.functional.to_pil_image(image_original)

def apply_aging_effect(model, image_path, age_source, age_cible):
    image = Image.open(image_path)
    return process_image(model, image, source_age=age_source, target_age=age_cible)

def show_aged_image():
    if 'image_path' in globals():
        predict_age = predict_age_from_model(face_aging_model, image_path)
        target_age = target_age_var.get()  # Récupérer l'âge cible depuis le slider
        aged_image = apply_aging_effect(face_aging_autoencoder, image_path, predict_age, target_age)
        aged_image_tk = ImageTk.PhotoImage(aged_image.resize((256, 256)))
        aged_panel.configure(image=aged_image_tk)
        aged_panel.image = aged_image_tk
    else:
        messagebox.showerror("Erreur", "Veuillez d'abord charger une image.")

# Création de la fenêtre principale Tkinter
root = tk.Tk()
root.title("Face Age Prediction and Aging")

# Affichage de l'image originale
original_panel = tk.Label(root)
original_panel.pack(pady=10)

# Boutons et widgets
load_button = tk.Button(root, text="Charger une image", command=load_image)
load_button.pack(pady=5)

predict_button = tk.Button(root, text="Prédire l'âge", command=predict_age)
predict_button.pack(pady=5)

# Ajout d'une variable pour stocker l'âge cible
target_age_var = tk.IntVar(value=80)

# Ajout d'une barre de défilement pour choisir l'âge cible
age_slider = tk.Scale(root, from_=5, to=100, orient="horizontal", label="Âge cible", variable=target_age_var, resolution=5)
age_slider.pack(pady=5)

age_button = tk.Button(root, text="Afficher l'image transformée", command=show_aged_image)
age_button.pack(pady=5)

# Affichage de l'image vieillie
aged_panel = tk.Label(root)
aged_panel.pack(pady=10)

# Lancer l'interface
root.mainloop()