import cv2
import os 
import numpy as np 

patches_folder = "Patches"
pred_patches = "Patch_pred"
os.makedirs(patches_folder, exist_ok=True)
os.makedirs(pred_patches,exist_ok=True)


def split(image, destination = patches_folder, patch_size = 256):
    img = cv2.imread(image)
    h,w,_ = img.shape
    for y in range(0, h, patch_size):
        for x in range(0, w, patch_size):
            patch = img[y:y+patch_size, x:x+patch_size]

            
            patch_filename = f"patch_{y}_{x}.png" 
            patch_path = os.path.join(destination, patch_filename)
            cv2.imwrite(patch_path, patch)

def merge(patch_folder , dest_image = 'out.png', image_shape = None):
    merged = np.zeros(image_shape[:-1] + (3,), dtype=np.uint8)
    for filename in os.listdir(patch_folder):
        if filename.endswith(".png"):
            patch_path = os.path.join(patch_folder, filename)
            patch = cv2.imread(patch_path)
            patch_height, patch_width, _ = patch.shape

            # Extract patch coordinates from filename
            parts = filename.split("_")
            x, y = None, None
            for part in parts:
                if part.endswith(".png"):
                    x = int(part.split(".")[0])
                elif part.isdigit():
                    y = int(part)
            if x is None or y is None:
                raise ValueError(f"Invalid filename: {filename}")

            # Check if patch fits within image boundaries
            if x + patch_width > image_shape[1] or y + patch_height > image_shape[0]:
                # Adjust patch position to fit within image boundaries
                if x + patch_width > image_shape[1]:
                    x = image_shape[1] - patch_width
                if y + patch_height > image_shape[0]:
                    y = image_shape[0] - patch_height

            # Merge patch into the main image
            merged[y:y+patch_height, x:x+patch_width, :] = patch

    cv2.imwrite(dest_image, merged)