import numpy as np
import cv2
from keras.models import load_model
from keras_preprocessing.image import img_to_array, load_img
import tkinter as tk
from tkinter import filedialog
from tkinter import messagebox
import matplotlib.pyplot as plt

# Load U-Net and CNN-LSTM models
unet_model = load_model(r"C:\Users\jack8\OneDrive\LSTM and ARIMA\LSTM\lstm_model.h5")
cnn_lstm_model = load_model(r'C:\Users\jack8\OneDrive\LSTM and ARIMA\LSTM\cnn_lstm_model.h5')

# Function to split image into patches
def image_to_patches(image, patch_size, overlap=0):
    patches = []
    positions = []
    img_height, img_width = image.shape[:2]
    patch_height, patch_width = patch_size
    step_y = patch_height - overlap
    step_x = patch_width - overlap
    for y in range(0, img_height - overlap + 1, step_y):
        for x in range(0, img_width - overlap + 1, step_x):
            patch = image[y:y + patch_height, x:x + patch_width]
            if patch.shape[0] != patch_height or patch.shape[1] != patch_width:
                pad_height = patch_height - patch.shape[0]
                pad_width = patch_width - patch.shape[1]
                patch = np.pad(patch, ((0, pad_height), (0, pad_width), (0, 0)), mode='constant')
            patches.append(patch)
            positions.append((y, x))
    return np.array(patches), positions

# Function to combine patches into a single image
def patches_to_image(patches, positions, image_shape, overlap=0):
    reconstructed_image = np.zeros(image_shape, dtype=np.float32)
    weight_mask = np.zeros(image_shape, dtype=np.float32)
    patch_height, patch_width = patches.shape[1:3]

    for patch, (y, x) in zip(patches, positions):
        y_end = min(y + patch_height, image_shape[0])
        x_end = min(x + patch_width, image_shape[1])
        h = y_end - y
        w = x_end - x
        
        patch_crop = patch[:h, :w, 0]  # Shape: (h, w)
        patch_crop = patch_crop[..., np.newaxis]  # Shape: (h, w, 1)

        reconstructed_image[y:y_end, x:x_end, 0] += patch_crop[:, :, 0]
        weight_mask[y:y_end, x:x_end, 0] += 1

    weight_mask[weight_mask == 0] = 1
    reconstructed_image /= weight_mask

    return reconstructed_image

# File selection dialog functions
def select_input_file():
    file_path = filedialog.askopenfilename(filetypes=[("Image Files", "*.png;*.jpg")])
    input_entry.delete(0, tk.END)
    input_entry.insert(0, file_path)

def select_output_folder():
    folder_path = filedialog.askdirectory()
    output_entry.delete(0, tk.END)
    output_entry.insert(0, folder_path)

# Image processing and saving function
def process_and_save():
    try:
        input_img_path = input_entry.get()
        output_folder = output_entry.get()
        model_type = model_var.get()

        if not input_img_path or not output_folder:
            messagebox.showerror("錯誤", "請選擇輸入文件和輸出文件夾")
            return
        
        # Load and preprocess image
        new_img = load_img(input_img_path)
        new_img_array = img_to_array(new_img) / 255.0
        original_height, original_width = new_img_array.shape[:2]

        # Convert to grayscale for edge detection
        if new_img_array.shape[2] == 3:  # Check if the image has three channels (RGB)
            gray_img = cv2.cvtColor((new_img_array * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
            # Apply viridis colormap
            colormap_img = cv2.applyColorMap(gray_img, cv2.COLORMAP_VIRIDIS)
            colormap_img = colormap_img / 255.0  # Normalize again
        else:
            colormap_img = new_img_array  # If it's already single-channel, use it directly

        # Set patch size and overlap based on model type
        if model_type == "U-Net":
            patch_size = (640, 640)
            overlap = 0
        else:  # For CNN, LSTM and CNN-LSTM
            patch_size = (128, 128)
            overlap = 0 if model_type != "CNN-LSTM" else 32
            
        patches, positions = image_to_patches(colormap_img, patch_size, overlap)

        predicted_patches = []
        if model_type == "U-Net":
            for patch in patches:
                patch_input = np.expand_dims(patch, axis=0)
                prediction = unet_model.predict(patch_input)
                predicted_patches.append(prediction[0])
        elif model_type in ["CNN", "LSTM"]:
            for patch in patches:
                patch_input = np.expand_dims(patch, axis=0)
                prediction = cnn_lstm_model.predict(patch_input)  # Assuming this handles both CNN and LSTM inputs
                predicted_patches.append(prediction[0])
        elif model_type == "CNN-LSTM":
            timesteps = 3
            for patch in patches:
                patch_sequence = np.array([patch] * timesteps)
                patch_input = np.expand_dims(patch_sequence, axis=0)
                prediction = cnn_lstm_model.predict(patch_input)
                predicted_patches.append(prediction[0])
        
        predicted_patches = np.array(predicted_patches)
        prediction_mask = patches_to_image(predicted_patches, positions, (original_height, original_width, 1), overlap)

        # Generate closed mask
        prediction_mask_uint8 = (prediction_mask[:, :, 0] * 255).astype(np.uint8)
        kernel_size = 20 if model_type == 'CNN-LSTM' else 29
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
        closed_mask = cv2.morphologyEx(prediction_mask_uint8, cv2.MORPH_CLOSE, kernel)
        closed_mask_binary = (closed_mask > 127).astype(np.uint8)

        output_path = f"{output_folder}/closed_mask_binary_{model_type}.png"
        cv2.imwrite(output_path, closed_mask_binary * 255)
        messagebox.showinfo("完成", f"Closed Mask Binary 已保存到: {output_path}")

        # Visualize result
        plt.figure(figsize=(10, 10))
        plt.imshow(closed_mask_binary, cmap='gray')
        plt.axis('off')
        plt.show()

    except Exception as e:
        messagebox.showerror("錯誤", f"發生錯誤: {e}")

# Create Tkinter window
root = tk.Tk()
root.title("模型選擇器")

# Configure grid to center elements
def center_window():
    window_width = 600
    window_height = 400
    screen_width = root.winfo_screenwidth()
    screen_height = root.winfo_screenheight()
    x = (screen_width // 2) - (window_width // 2)
    y = (screen_height // 2) - (window_height // 2)
    root.geometry(f'{window_width}x{window_height}+{x}+{y}')

center_window()

# Model selection radio buttons
model_var = tk.StringVar(value="U-Net")  # Default to U-Net
tk.Label(root, text="選擇模型:").grid(row=0, column=0, padx=10, pady=10, columnspan=4)
tk.Radiobutton(root, text="U-Net", variable=model_var, value="CNN-LSTM").grid(row=1, column=0, padx=5, pady=5)
tk.Radiobutton(root, text="CNN-LSTM", variable=model_var, value="U-Net").grid(row=1, column=1, padx=5, pady=5)
tk.Radiobutton(root, text="CNN", variable=model_var, value="CNN").grid(row=1, column=2, padx=5, pady=5)
tk.Radiobutton(root, text="LSTM", variable=model_var, value="LSTM").grid(row=1, column=3, padx=5, pady=5)

# Input file selection
tk.Label(root, text="選擇輸入圖片:").grid(row=2, column=0, padx=10, pady=10, columnspan=4)
input_entry = tk.Entry(root, width=50)
input_entry.grid(row=3, column=0, columnspan=2, padx=10, pady=10)
input_button = tk.Button(root, text="瀏覽", command=select_input_file)
input_button.grid(row=3, column=2, padx=10)

# Output folder selection
tk.Label(root, text="選擇輸出文件夾:").grid(row=4, column=0, padx=10, pady=10, columnspan=4)
output_entry = tk.Entry(root, width=50)
output_entry.grid(row=5, column=0, columnspan=2, padx=10, pady=10)
output_button = tk.Button(root, text="瀏覽", command=select_output_folder)
output_button.grid(row=5, column=2, padx=10)

# Execute button
process_button = tk.Button(root, text="執行並保存", command=process_and_save)
process_button.grid(row=6, column=1, padx=10, pady=20, columnspan=2)

# Run Tkinter main loop
root.mainloop()


###輸出的二值圖先不要儲存，再增加一個使用canny法過濾黑白二值圖之後輸出線圖