# tk_Image_Enhance_GUI.py

# tk_Image_Enhance_GUI.py ZZZ? Most likely needs more work

from PIL import Image
from tkinter import Tk, Button, Canvas, PhotoImage, Frame
from tkinter.filedialog import askopenfilename
import os
import pickle
import time

def select_and_process_image():
    filename = askopenfilename(filetypes=[("JPG files", "*.jpg")])
    if not filename:
        print("No file selected")
        return

    if not filename.lower().endswith('.jpg'):
        print(f"Error: {filename} is not a JPG file")
        return

    try:
        original_img = Image.open(filename)
        original_img = original_img.convert('RGB')
    except Exception as e:
        print(f"Error loading image: {e}")
        return

    # Build or load patch dictionary
    dictionaries = build_patch_dictionary(original_img, filename)

    # Enhance image (default scale factor t=2, can be changed)
    enhanced_img = enhance_image(original_img, dictionaries, t=2)

    # Save enhanced image
    enhanced_filename = f"enhanced_{os.path.basename(filename)}"
    enhanced_img.save(enhanced_filename)
    print(f"Enhanced image saved as {enhanced_filename}")

    # Display on canvas (resize to fit 640x640)
    display_size = (640, 640)
    enhanced_img_display = enhanced_img.resize(display_size, Image.BICUBIC)
    photo = PhotoImage(enhanced_img_display)

    canvas.delete("all")
    canvas.create_image(320, 320, image=photo)
    canvas.image = photo  # Keep a reference to avoid garbage collection
    print("Image displayed on canvas")

def progress_meter(i):
    canvas.delete('meter')
    canvas.create_rectangle(8, 8, 632, 42, outline="black", fill="yellow", tags='meter')
    canvas.create_rectangle(10, 10, i % 630, 40, outline="", fill="lime", tags='meter')

    # Calculate text position to be centered in the progress bar
    text_x = 320
    text_y = 25

    canvas.create_text(text_x, text_y, text="Progress...", fill="black", tags='meter', font=("Arial", 14))
    canvas.update()

# Helper function to convert a 2D list patch to a PIL Image
def list_to_image(patch):
    height = len(patch)
    width = len(patch[0])
    img = Image.new('RGB', (width, height))
    for i in range(height):
        for j in range(width):
            img.putpixel((j, i), tuple(patch[i][j]))
    return img

# Helper function to convert a PIL Image to a 2D list
def image_to_list(img):
    width, height = img.size
    return [[list(img.getpixel((j, i))) for j in range(width)] for i in range(height)]

# Helper function to upsample a patch by factor t
def upsample_patch(patch, t):
    img = list_to_image(patch)
    new_size = (len(patch[0]) * t, len(patch) * t)
    img_up = img.resize(new_size, Image.BICUBIC)
    return image_to_list(img_up)

# Helper function to downscale a patch to 4x4
def downscale_patch(patch, t):
    img = list_to_image(patch)
    img_down = img.resize((4, 4), Image.BICUBIC)
    return image_to_list(img_down)

# Flatten a patch for dot product computation (4x4 RGB -> 192 elements)
def flatten_patch(patch):
    flat = []
    for i in range(4):
        for j in range(4):
            flat.extend(patch[i][j])
    return flat

# Compute dot product between two flattened patches
def dot_product(a, b):
    return sum(a[i] * b[i] for i in range(len(a)))

# Subtract two patches element-wise (4x4)
def subtract_patches(p1, p2):
    return [[[p1[i][j][k] - p2[i][j][k] for k in range(3)] for j in range(4)] for i in range(4)]

# Add two patches element-wise with clamping to [0, 255]
def add_patches(p1, p2):
    height = len(p1)
    width = len(p1[0])
    return [[[min(max(p1[i][j][k] + p2[i][j][k], 0), 255) for k in range(3)] for j in range(width)] for i in range(height)]

# Extract non-overlapping 4x4 patches from an image list
def extract_patches(img_list, width, height):
    patches = []
    positions = []
    for i in range(0, height - 3, 4):
        for j in range(0, width - 3, 4):
            patch = [[img_list[i + di][j + dj] for dj in range(4)] for di in range(4)]
            patches.append(patch)
            positions.append((i, j))
    return patches, positions

# Log message to file
def log(message):
    with open('processing_log.txt', 'a') as f:
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - {message}\n")

# build_patch_dictionary.py

def build_patch_dictionary(original_img, filename):
    cache_filename = f"cache_{os.path.splitext(os.path.basename(filename))[0]}.pkl"

    # Check cache
    if os.path.exists(cache_filename):
        print(f"Loading cached dictionary from {cache_filename}")
        with open(cache_filename, 'rb') as f:
            dictionaries = pickle.load(f)
        return dictionaries

    print(f"Building patch dictionary from {filename}")
    width, height = original_img.size
    original_list = image_to_list(original_img)

    # Dictionary to store patches for each scale
    dictionaries = {}
    scale_factors = [2, 4, 6, 8]

    for s in scale_factors:
        # Downscale image
        downscaled_size = (width // s, height // s)
        I_s = original_img.resize(downscaled_size, Image.BICUBIC)
        # Upsample back to original size
        I_s_up = I_s.resize((width, height), Image.BICUBIC)
        I_s_up_list = image_to_list(I_s_up)

        # Extract non-overlapping 4x4 patches
        high_res_patches, positions = extract_patches(original_list, width, height)
        low_res_patches, _ = extract_patches(I_s_up_list, width, height)

        # Flatten low-res patches for similarity computation
        low_res_flat = [flatten_patch(patch) for patch in low_res_patches]

        # Compute top 250 similar patches for each low-res patch
        neighbors = []

        k = 0
        for i in range(len(low_res_flat)):
            similarities = []
            for j in range(len(low_res_flat)):
                if i != j:
                    sim = dot_product(low_res_flat[i], low_res_flat[j])
                    similarities.append((sim, j))
                    k += 1
                    progress_meter(k)
            similarities.sort(reverse=True)
            top_neighbors = [idx for _, idx in similarities[:250]]
            neighbors.append(top_neighbors)
        canvas.delete('meter')

        # Store in dictionary
        dictionaries[s] = {
            'low_res_patches': low_res_patches,
            'high_res_patches': high_res_patches,
            'positions': positions,
            'neighbors': neighbors
        }

    # Save to cache
    with open(cache_filename, 'wb') as f:
        pickle.dump(dictionaries, f)
    print(f"Saved patch dictionary to {cache_filename}")

    return dictionaries

# enhance_image.py

def enhance_image(original_img, dictionaries, t):
    print(f"Enhancing image with scale factor {t}")
    width, height = original_img.size

    # Upsample original image by t
    upsampled_size = (width * t, height * t)
    I_up = original_img.resize(upsampled_size, Image.BICUBIC)
    I_up_list = image_to_list(I_up)

    # Get dictionary for scale t
    if t not in dictionaries:
        print(f"No dictionary for scale factor {t}, using nearest available")
        t_available = min(dictionaries.keys(), key=lambda x: abs(x - t))
    else:
        t_available = t
    dict_data = dictionaries[t_available]
    low_res_patches = dict_data['low_res_patches']
    high_res_patches = dict_data['high_res_patches']
    low_res_flat = [flatten_patch(patch) for patch in low_res_patches]

    # Initialize output image list
    output_list = [[I_up_list[i][j] for j in range(upsampled_size[0])] for i in range(upsampled_size[1])]

    # Extract non-overlapping patches of size (4*t)x(4*t)
    patch_size = 4 * t
    for i in range(0, upsampled_size[1] - patch_size + 1, patch_size):
        for j in range(0, upsampled_size[0] - patch_size + 1, patch_size):
            # Extract patch Q
            Q = [[I_up_list[i + di][j + dj] for dj in range(patch_size)] for di in range(patch_size)]

            # Downscale Q to 4x4
            Q_down = downscale_patch(Q, t)
            Q_down_flat = flatten_patch(Q_down)

            # Find top K similar low-res patches (K=5 for robustness)
            K = 5
            similarities = []
            for idx in range(len(low_res_flat)):
                sim = dot_product(Q_down_flat, low_res_flat[idx])
                similarities.append((sim, idx))
            similarities.sort(reverse=True)
            top_k = similarities[:K]

            # Compute weights
            total_sim = sum(sim for sim, _ in top_k) + 1e-10  # Avoid division by zero
            weights = [sim / total_sim for sim, _ in top_k]

            # Compute enhanced patch using weighted average of differences
            D_acc = [[ [0 for k in range(3)] for j in range(4)] for i in range(4)]
            for (sim, idx), w in zip(top_k, weights):
                P_low = low_res_patches[idx]
                P_high = high_res_patches[idx]
                D = subtract_patches(P_high, P_low)
                # Scale difference by weight
                D_scaled = [[[D[i][j][k] * w for k in range(3)] for j in range(4)] for i in range(4)]
                D_acc = add_patches(D_acc, D_scaled)

            # Upsample D_acc to patch_size x patch_size
            D_up = upsample_patch(D_acc, t)

            # Add to Q with clamping
            Q_enhanced = add_patches(Q, D_up)

            # Place enhanced patch in output image
            for di in range(patch_size):
                for dj in range(patch_size):
                    if i + di < upsampled_size[1] and j + dj < upsampled_size[0]:
                        output_list[i + di][j + dj] = Q_enhanced[di][dj]

    enhanced_img = list_to_image(output_list)
    print("Image enhancement completed")
    return enhanced_img

# Set up GUI
root = Tk()
root.title("Image Enhance GUI")
root.geometry("+0+0")

canvas = Canvas(root, width=640, height=640, bg="white")
canvas.pack(side="left")

right_frame = Frame(root)
right_frame.pack(side="right", padx=10, pady=10, fill="both", expand=True)

select_button = Button(right_frame, text="Select JPG Image", command=select_and_process_image)
select_button.pack(pady=10)

print("Started image enhancement application")

root.mainloop()