import gradio as gr
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageEnhance
from skimage.metrics import structural_similarity as ssim
import tempfile

# ASCII characters to map pixel intensity
ASCII_CHARS = [
    ' ', '.', ',', ':', ';', 'i', 'l', '!', 'I', '1', 't', 'o', 'r', 'x', 'z', 'v', 
    'u', 'n', 'm', 'w', 'Q', 'B', 'N', 'M', '@'
]

# Step 1: Resize image and perform general Sobel edge detection

def preprocess_image(image_path, new_width=150):
    # Open the image and convert to grayscale
    image = Image.open(image_path)
    print(type(image))
    # Step 1: Resize the image for ASCII aspect ratio
    width, height = image.size
    aspect_ratio = height / width * 0.55  # Adjusting for ASCII aspect ratio
    new_height = int(aspect_ratio * new_width)
    
    # # Step 3: Show the equalized image
    # equalized_image.show(title="Equalized Image")
    opencv_image = np.array(image)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    equalized_image = clahe.apply(opencv_image)

    # Convert to OpenCV format for Sobel edge detection
    opencv_image = np.array(equalized_image)

    # Step 4: Sobel edge detection
    sobelx = cv2.Sobel(opencv_image, cv2.CV_64F, 1, 0, ksize=5)  # Horizontal gradients
    sobely = cv2.Sobel(opencv_image, cv2.CV_64F, 0, 1, ksize=5)  # Vertical gradients
    edges = np.hypot(sobelx, sobely)  # Combine gradients

    # Normalize the result to fit into the 0-255 grayscale range
    edges = np.uint8(255 * edges / np.max(edges))

    # Step 5: Sharpen the edges
    edges_image = Image.fromarray(edges)
    edges_image.show(title="Sharpened Edges")
    enhancer = ImageEnhance.Contrast(edges_image)
    enhanced_edges = enhancer.enhance(3.0)  # Increase contrast by a factor (2.0 is an example)

    sharpened_edges = enhanced_edges.filter(ImageFilter.SHARPEN)

    # sharpened_edges = edges_image.filter(ImageFilter.SHARPEN)

    # Step 6: Show the sharpened edges image
    sharpened_edges.show(title="Sharpened Edges")

    sharpened_edges = sharpened_edges.resize((new_width, new_height))
    sharpened_edges.show(title="Sharpened Edges")
    return sharpened_edges, sharpened_edges

# Step 2: Convert image to ASCII art
def image_to_ascii(image):
    pixels = np.array(image)
    ascii_str = ""
    for row in pixels:
        for pixel in row:
            ascii_str += ASCII_CHARS[pixel // 10]  # Mapping pixel intensity to ASCII (256 levels / 25 characters)
        ascii_str += "\n"
    return ascii_str

# Step 3: Convert ASCII text back to image
def ascii_to_image(ascii_art, font_size=12, output_file='ascii_art.png'):
    font_path = "DejaVuSansMono.ttf"  # Update with correct path if needed
    font = ImageFont.truetype(font_path, font_size)
    
    lines = ascii_art.splitlines()
    max_width = max(font.getbbox(line)[2] for line in lines)  # Get max width
    img_height = len(lines) * font_size

    image = Image.new("RGB", (max_width, img_height), "white")
    draw = ImageDraw.Draw(image)

    for y, line in enumerate(lines):
        draw.text((0, y * font_size), line, fill="black", font=font)

    image.save(output_file)
    return image

# Step 4: Calculate Structural Similarity Index (SSIM)
def compare_ssim(image_path_1, image_path_2):
    image1 = cv2.imread(image_path_1, cv2.IMREAD_GRAYSCALE)
    image2 = cv2.imread(image_path_2, cv2.IMREAD_GRAYSCALE)

    # Resize images to the same size if needed
    image2 = cv2.resize(image2, (image1.shape[1], image1.shape[0]))

    # Compute SSIM between the two images
    score, _ = ssim(image1, image2, full=True)
    return score

# Usage Example
if __name__ == "__main__":
    # Path to input image
    input_image_path = "image8a.jpg"

    # Step 1: Preprocess and edge detection with Sobel filter
    resized_image, edges = preprocess_image(input_image_path)

    # Step 2: Convert the edge-detected image to ASCII
    ascii_art = image_to_ascii(edges)
    print("ASCII Art:\n", ascii_art)

    # Step 3: Convert ASCII back to image
    ascii_image_path = "ascii_image.jpg"
    ascii_img = ascii_to_image(ascii_art, 12, ascii_image_path)

    # Step 4: Compare SSIM
    ssim_score = compare_ssim(input_image_path, ascii_image_path)
    print(f"SSIM between original and ASCII image: {ssim_score}")

# def process_image(image):
#     # Step 1: Preprocess and convert to ASCII
#     edges, _ = preprocess_image(image)
#     ascii_art = image_to_ascii(edges)
    
#     # Step 2: Convert ASCII art back to an image for SSIM calculation
#     ascii_image = ascii_to_image(ascii_art)
    
#     # Step 3: Calculate SSIM
#     ssim_score = compare_ssim(image, ascii_image)
    
#     # Step 4: Return ASCII art and SSIM score
#     with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as f:
#         f.write(ascii_art.encode())
#         ascii_txt_path = f.name
    
#     return ascii_art, ssim_score, ascii_txt_path

# # Gradio Interface
# def gradio_interface(image):
#     print(type(image))
#     ascii_art, ssim_score, ascii_txt_path = process_image(image)
#     return ascii_art, ssim_score, ascii_txt_path

# # Setup Gradio interface
# interface = gr.Interface(
#     fn=gradio_interface,
#     inputs="image",
#     outputs=["text", "number", "file"],
#     title="Image to ASCII Art with SSIM",
#     description="Upload an image, see the ASCII art, the SSIM score, and download the ASCII art as a text file."
# )

interface.launch()