import torch
import numpy as np
from pathlib import Path
import os
from torchvision.transforms import ToTensor, ToPILImage
from PIL import Image, ImageDraw, ImageFont

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#Get num closest to 8
def cl8(num):
    rem = num % 8
    if rem <= 4:
        return round(num - rem)
    else:
        return round(num + (8 - rem))

def closest_lcm(n, div1, div2):
    # Find the LCM
    lcm = np.lcm(int(div1), int(div2))

    # Find the multiple of LCM closest to n
    lower_multiple = (n // lcm) * lcm  # Largest multiple of lcm less than or equal to n
    upper_multiple = lower_multiple + lcm  # Smallest multiple of lcm greater than n

    # Find which multiple is closer to n
    if n - lower_multiple > upper_multiple - n:
        return upper_multiple
    else:
        return lower_multiple
    
def normalize_size(images):
    refimage = images[0]
    refimage = refimage.resize((cl8(refimage.width), cl8(refimage.height)), Image.Resampling.LANCZOS)
    return_images = []
    for i in range(len(images)):
        if images[i].size != refimage.size:
            images[i] = images[i].resize(refimage.size, Image.Resampling.LANCZOS)
        return_images.append(images[i])
    return return_images

def constrain_image(image, max_width, max_height):
    width, height = image.size
    aspect_ratio = width / float(height)

    if width > max_width or height > max_height:
        if width / float(max_width) > height / float(max_height):
            new_width = max_width
            new_height = int(new_width / aspect_ratio)
        else:
            new_height = max_height
            new_width = int(new_height * aspect_ratio)
        image = image.resize((cl8(new_width), cl8(new_height)), Image.Resampling.LANCZOS)

    return image

def padlist(lst, targetsize):
    if targetsize <= len(lst):
            return lst[:targetsize]

    last_elem = lst[-1]
    num_repeats = targetsize - len(lst)
    
    return lst + [last_elem] * num_repeats

def MakeGrid(images, rows, cols):
    widths, heights = zip(*(i.size for i in images))

    grid_width = max(widths) * cols
    grid_height = max(heights) * rows
    cell_width = grid_width // cols
    cell_height = grid_height // rows
    final_image = Image.new('RGB', (grid_width, grid_height))
    x_offset = 0
    y_offset = 0
    for i in range(len(images)):
        final_image.paste(images[i], (x_offset, y_offset))
        x_offset += cell_width
        if x_offset == grid_width:
            x_offset = 0
            y_offset += cell_height

    # Save the final image
    return final_image

def BreakGrid(grid, rows, cols):
    width = grid.width // cols
    height = grid.height // rows
    outimages = []
    for row in range(rows):
            for col in range(cols):
                left = col * width
                top = row * height
                right = left + width
                bottom = top + height
                current_img = grid.crop((left, top, right, bottom))
                outimages.append(current_img)
    return outimages

def ImgLabeler(img, text, size=72, color=(255,255,255)):
    font = None
    # A list of common font names for cross-platform compatibility
    font_names = [
        "arial.ttf",              # Windows, macOS
        "DejaVuSans.ttf",         # Common on Linux
        "LiberationSans-Regular.ttf" # Another common Linux font
    ]

    for font_name in font_names:
        try:
            font = ImageFont.truetype(font_name, size)
            break  # Stop searching once a font is found
        except IOError:
            continue # Try the next font in the list

    if font is None:
        # If no common fonts are found, fall back to the default bitmap font.
        print(f"Could not find any of {font_names}. Falling back to default font. Text may be very small.")
        font = ImageFont.load_default()

    draw = ImageDraw.Draw(img)

    # Use textbbox for better centering if available (Pillow >= 8.0.0)
    if hasattr(draw, 'textbbox'):
        bbox = draw.textbbox((0, 0), text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        x = (img.width - text_width) / 2
        y = (img.height - text_height) / 2
        # Adjust position based on the bounding box's top-left corner
        text_position = (x - bbox[0], y - bbox[1])
    else: # Fallback for older Pillow versions
        # textsize is deprecated but works as a fallback.
        text_width, text_height = draw.textsize(text, font=font)
        x = (img.width - text_width) / 2
        y = (img.height - text_height) / 2
        text_position = (x, y)

    # Draw outline for better visibility
    outline_color = (0, 0, 0)
    for offset in [(-1, -1), (-1, 1), (1, -1), (1, 1), (-1, 0), (1, 0), (0, -1), (0, 1)]:
        draw.text((text_position[0] + offset[0], text_position[1] + offset[1]), text, font=font, fill=outline_color)

    # Draw the main text onto the image
    draw.text(text_position, text, font=font, fill=color)
    
    return img

def load_and_preprocess(image_path):
    with Image.open(image_path) as img:
        return torch.tensor(np.array(img.convert('L')), device=device, dtype=torch.float32)

def compute_histogram(tensor):
    # Min and max values for grayscale images
    min_val, max_val = 0, 255
    # Compute the histogram by counting the number of occurrences within each bin
    hist = torch.histc(tensor, bins=255, min=min_val, max=max_val)
    return hist / tensor.numel()  # Normalize by the number of elements

def get_iterated_path(directory, base_filename, extension='.png'):
    # Construct the full file path
    count = 0
    while True:
        # Append a count to the filename if it's not the first file
        if count == 0:
            unique_filename = f"{base_filename}{extension}"
        else:
            unique_filename = f"{base_filename}_{count}{extension}"
        full_file_path = os.path.join(directory, unique_filename)

        # Check if a file with this name already exists
        if not os.path.exists(full_file_path):
            break  # Exit the loop once the image is saved
        count += 1
    
    return full_file_path

def CheckMakeDir(dir):
    try:
        if not Path.exists(Path(dir)):
            Path.mkdir(Path(dir), parents=True, exist_ok=True)
        return True, "Success"
    except:
        st_out = f"Error: could not locate nor create directory {dir}."
        print(st_out)
        return False, st_out

def apply_conditional_ema_pytorch(images: list[Image.Image], alpha: float, threshold: float) -> list[Image.Image]:

    processed_images = []
    ema_tensor = None
    to_tensor = ToTensor()
    to_pil = ToPILImage()

    for image in images:
        frame_tensor = to_tensor(image).to(device)

        if ema_tensor is None:
            # For the first frame, EMA is just the frame itself
            ema_tensor = frame_tensor.clone()
        else:
            # Calculate EMA for each subsequent frame
            deviation = torch.abs(frame_tensor - ema_tensor)
            update_mask = deviation > threshold  # Mask to identify where to update

            # EMA update
            ema_tensor = alpha * frame_tensor + (1 - alpha) * ema_tensor
            ema_tensor[update_mask] = frame_tensor[update_mask]  # Apply conditional update

        processed_images.append(to_pil(ema_tensor))

    return processed_images

def cat_to_pils(tensor):
    to_pil = ToPILImage()
    tensor = tensor.permute(0, 3, 1, 2)
    pils = [to_pil(tensor[i]) for i in range(tensor.shape[0])]
    return pils

def pil_to_tens(pimg):
    to_tensor = ToTensor()
    tensor = to_tensor(pimg).unsqueeze(0).permute(0, 2, 3, 1)
    return tensor

def get_grid_aspect(num_images: int, image_width: int, image_height: int) -> (int, int):

    if num_images == 0:
        return 0, 0

    min_diff = float('inf')
    best_layout = (1, num_images)

    if image_width > image_height:
        for cols in range(1, num_images + 1):
            rows = -(-num_images // cols)
            grid_width = cols * image_width
            grid_height = rows * image_height
            diff = abs(grid_width - grid_height)

            if diff < min_diff:
                min_diff = diff
                best_layout = (rows, cols)

            if cols > num_images / cols:
                break
    else:
        for rows in range(1, num_images + 1):
            cols = -(-num_images // rows)
            grid_width = cols * image_width
            grid_height = rows * image_height
            diff = abs(grid_height - grid_width)

            if diff < min_diff:
                min_diff = diff
                best_layout = (rows, cols)

            if rows > num_images / rows:
                break

    return best_layout

# --- Image Difference Calculation Functions ---

def calculate_mse_diff(frames1_t, frames2_t):
    """Calculates the Mean Squared Error between two batches of frames."""
    diffs = frames1_t.float() - frames2_t.float()
    return torch.mean(diffs ** 2, dim=(1, 2, 3))

def calculate_ssim_diff(frames1_p, frames2_p):
    """Calculates the Structural Similarity Index (SSIM) difference between two lists of PIL images."""
    from skimage.metrics import structural_similarity as ssim
    diffs = []
    for i in range(len(frames1_p)):
        # Convert images to grayscale numpy arrays for SSIM
        im1 = np.array(frames1_p[i].convert('L'))
        im2 = np.array(frames2_p[i].convert('L'))
        # SSIM returns a value between -1 and 1, where 1 is perfect similarity.
        # We want the difference, so we use 1 - ssim.
        similarity = ssim(im1, im2, data_range=im2.max() - im2.min())
        diffs.append(1.0 - similarity)
    return torch.tensor(diffs, device=device)

def calculate_phash_diff(frames1_p, frames2_p):
    """Calculates the perceptual hash (pHash) difference between two lists of PIL images."""
    import imagehash
    diffs = []
    for i in range(len(frames1_p)):
        hash1 = imagehash.phash(frames1_p[i])
        hash2 = imagehash.phash(frames2_p[i])
        # The difference is the number of bits that are different (Hamming distance).
        # We normalize it by the total number of bits in the hash (usually 64).
        diffs.append((hash1 - hash2) / len(hash1))
    return torch.tensor(diffs, device=device)

def calculate_contrast(frame_p):
    """Calculates the RMS contrast of a single PIL image."""
    # Convert to grayscale numpy array
    im = np.array(frame_p.convert('L'), dtype=float)
    return im.std()

def calculate_edge_diff(frame1_p, frame2_p, focus_ratio=0.3):
    """Calculates the difference between the edges of the central focused area of two PIL images."""
    import cv2

    def process_frame(frame):
        # Crop to central focus area
        w, h = frame.size
        crop_w, crop_h = int(w * focus_ratio), int(h * focus_ratio)
        left, top = (w - crop_w) // 2, (h - crop_h) // 2
        right, bottom = left + crop_w, top + crop_h
        cropped = frame.crop((left, top, right, bottom))
        
        # Convert to grayscale numpy array
        gray = np.array(cropped.convert('L'))
        
        # Apply Laplacian operator to detect edges
        edges = cv2.Laplacian(gray, cv2.CV_64F)
        return edges

    edges1 = process_frame(frame1_p)
    edges2 = process_frame(frame2_p)
    
    # Calculate Mean Squared Error between the edge maps
    return np.mean((edges1 - edges2) ** 2)