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

# 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', '@'
]


def preprocess_image(image, new_width=150):
    """
    This function performs the following operations on the input image:
    -enhance contrast
    -perform edge detection (Sobel)
    -denoise edges
    -enhance contrast
    -resize image
    """
    # Open the image and convert to grayscale
    image = Image.fromarray(image)
    image = image.convert("L")

    # Contrast limited adaptive histogram equalization
    opencv_image = np.array(image)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    equalized_image = clahe.apply(opencv_image)

    # Apply Sobel edge detection
    sobelx = cv2.Sobel(equalized_image, cv2.CV_64F, 1, 0, ksize=5)
    sobely = cv2.Sobel(equalized_image, cv2.CV_64F, 0, 1, ksize=5)
    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))

    # Denoise the edges using Non-Local Means Denoising
    edges = cv2.fastNlMeansDenoising(edges, None, h=30, templateWindowSize=7, searchWindowSize=21)

    # Improve contrast of edges
    edges_image = Image.fromarray(edges)
    enhancer = ImageEnhance.Contrast(edges_image)
    enhanced_edges = enhancer.enhance(1.5)
    enhanced_edges = np.array(enhanced_edges)

    # Resize the image for ASCII aspect ratio
    width, height = image.size
    aspect_ratio = height / width * 0.45  # Adjusting for ASCII aspect ratio
    new_height = int(aspect_ratio * new_width)
    resized_image = Image.fromarray(enhanced_edges).resize((new_width, new_height))

    return resized_image


def image_to_ascii(image):
    """
    This function generates an ASCII representation of images.
    """
    pixels = np.array(image)
    ascii_str = ""
    for row in pixels:
        for pixel in row:
            # Mapping pixel intensity to ASCII
            ascii_str += ASCII_CHARS[pixel // 10]
        ascii_str += "\n"
    return ascii_str


def ascii_to_image(ascii_art, font_size=12, output_file='ascii_art.png'):
    """"
    This function converts ASCII text art to a png image.
    """

    # Specify font
    font_path = "DejaVuSansMono.ttf"
    font = ImageFont.truetype(font_path, font_size)

    # Get the required image dimensions
    lines = ascii_art.splitlines()
    max_width = max(font.getbbox(line)[2] for line in lines)  # Get max width
    img_height = len(lines) * font_size

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

    # Draw each character from the ASCII text on the image
    for y, line in enumerate(lines):
        draw.text((0, y * font_size), line, fill="black", font=font)

    image.save(output_file)
    return image


def compare_ssim(image1, image2):
    """
    This function generates the structural similarity index
    measure for the two input images.
    """
    # Convert original color image to grayscale
    image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    
    # Convert ASCII image to grayscale and cv2 format
    image2 = np.array(image2.convert("L"))

    # Resize so both images have the same dimensions.
    image2 = cv2.resize(image2, (image1.shape[1], image1.shape[0]))

    score, _ = ssim(image1, image2, full=True)
    return score


def process_image(image):
    """
    This function takes an image and generates:
        -ASCII art as a string
        -ASCII text file
        -Compute SSIM for the ASCII art image and the original image
    """
    # Get image edges and convert to ASCII
    edges = preprocess_image(image)
    ascii_art = image_to_ascii(edges)

    # Convert ASCII art back to an image for SSIM calculation
    ascii_image = ascii_to_image(ascii_art)

    # Calculate SSIM
    ssim_score = compare_ssim(image, ascii_image)

    # Write ASCII art to txt file
    ascii_txt_path = "ascii_art_output.txt"
    with open(ascii_txt_path, "w") as f:
        f.write(ascii_art)

    return ascii_art, ssim_score, ascii_txt_path


def gradio_interface(image):
    """
    This function generates the interface outputs for a given input image.
    The outputs include the formatted SSIM, ASCII art, and ASCII art txt file path.
    """
    ascii_art, ssim_score, ascii_txt_path = process_image(image)
    ssim_formatted = f'<div style="font-size: 20px;">The structural similarity index measure is {ssim_score:.3f}</div>'
    # Wrap the ASCII art in HTML with custom font size
    styled_ascii_art = f'<pre style="font-family: monospace; font-size: 7px;">{ascii_art}</pre>'
    
    return styled_ascii_art, ssim_formatted, ascii_txt_path

# Setup Gradio interface
with gr.Blocks() as interface:
    # Add title and description
    gr.Markdown("<h1 style='font-size: 40px;'>Image to ASCII Art Converter</h1>")
    gr.Markdown("<p style='font-size: 18px;'>Upload an image, and this tool will generate an ASCII art version of the image. It also calculates the Structural Similarity Index (SSIM) to evaluate the quality.</p>")

    # Image input and ASCII art + SSIM score output
    image_input = gr.Image(label="Upload an image")
    ascii_art_output = gr.HTML()
    ssim_score_output = gr.HTML()

    # File output for downloading, initially hidden
    ascii_file_output = gr.File(label="Download ASCII Art")

    # Update interface when image is uploaded
    image_input.change(fn=gradio_interface,
                       inputs=image_input,
                       outputs=[ascii_art_output, ssim_score_output, ascii_file_output])

# Launch interface
interface.launch()