# ela_hybrid.py
import numpy as np
import cv2 as cv
from PIL import Image
import matplotlib.pyplot as plt


def compress_jpg(image, quality=75):
    """Compress image using JPEG compression (shared from ela.py)."""
    encode_param = [int(cv.IMWRITE_JPEG_QUALITY), quality]
    _, buffer = cv.imencode('.jpg', image, encode_param)
    return cv.imdecode(buffer, cv.IMREAD_COLOR)


def generate_ela_hybrid(image_path: str, quality: int = 95, scale_factor: int = 150):
    """
    Generate a 6-channel hybrid image combining RGB and ELA (3 channels each).
    
    Args:
        image_path (str): Path to the input image.
        quality (int): JPEG compression quality (1-100).
        scale_factor (int): Scale factor for ELA contrast.
        
    Returns:
        np.ndarray: 6-channel (RGB + ELA) image with shape (H, W, 6), dtype float32, normalized [0,1]
    """
    # Load original image
    original = cv.imread(image_path, cv.IMREAD_COLOR)
    original = original.astype(np.float32) / 255  # Normalize to [0, 1]

    # Compress and reload image
    compressed = compress_jpg(original, quality)
    compressed = compressed.astype(np.float32) / 255  # Normalize to [0, 1]

    # Generate ELA as absolute difference between original and compressed
    ela = cv.absdiff(original, compressed)

    # Apply scale factor to enhance ELA differences
    ela = cv.convertScaleAbs(ela, alpha=scale_factor / 100, beta=0)
    ela = ela.astype(np.float32) / 255  # Normalize back to [0, 1]

    # Stack RGB and ELA (3 channels each) into 6-channel input
    hybrid_image = np.concatenate([original, ela], axis=-1)  # Shape: H×W×6

    return hybrid_image


def save_hybrid_image(hybrid_array, save_path: str):
    """Save the 6-channel hybrid image as a .npy file for model input."""
    np.save(save_path, hybrid_array)
    print(f"Saved hybrid ELA image to {save_path}")

def visualize_hybrid(hybrid_array: np.ndarray):
    """Return a list of 2 images: [RGB Image, ELA Map (grayscale)], as PIL Images."""
    from PIL import Image

    # Extract RGB (3 channels at front)
    rgb_image = Image.fromarray((hybrid_array[:, :, :3] * 255).astype(np.uint8))
    
    # Extract ELA channels (last 3)
    ela_image = Image.fromarray((hybrid_array[:, :, 3:] * 255).astype(np.uint8))

    return [rgb_image, ela_image]
    # ela_hybrid.py
    
def generate_hybrid_ela_func(img_input):
    """
    Returns a list of 2 PIL Images: [Original RGB Image, Hybrid ELA Output].
    """
    from PIL import Image
    import numpy as np

    # Convert input to numpy array if needed
    if not isinstance(img_input, np.ndarray):
        img_input = np.array(img_input)  # Ensure numpy format

    # Generate ELA
    hybrid_array = generate_ela_hybrid(img_input, quality=75, scale_factor=100)
    visualizations = visualize_hybrid(hybrid_array)
    return list(visualizations)  # Returns [RGB PIL Image, ELA PIL Image]


def generate_concatenated_hybrid(img_array):
    """
    Concatenate original RGB and ELA map into a single image.
    Returns a single PIL Image (side-by-side).
    """
    # Extract RGB and ELA channels
    original_rgb = Image.fromarray((img_array[:, :, :3] * 255).astype(np.uint8))  # 3 channels
    ela_map = Image.fromarray((img_array[:, :, 3:] * 255).astype(np.uint8))  # 3 channels

    # Resize to match height
    ela_map = ela_map.resize(original_rgb.size)

    # Concatenate and return as a single PIL Image
    combined = Image.new("RGB", (ela_map.width + original_rgb.width, original_rgb.height))  # Single row
    combined.paste(original_rgb, (0, 0))
    combined.paste(ela_map, (original_rgb.width, 0))
    return combined