from PIL import Image, ImageDraw, ImageOps, ImageFilter, ImageFont
import numpy as np
import math
import gradio as gr

##### CONSTANTS #####
# ASCII characters used to represent image pixels, reversed for better contrast in mapping
CHARS = ' .:-=+*#%@'[::-1]
# Convert the characters to a list for easier access
CHAR_ARRAY = list(CHARS)
# Number of available ASCII characters
CHAR_LEN = len(CHAR_ARRAY)

# Grayscale level for each ASCII character, determining how many shades of gray each character represents
GRAYSCALE_LEVEL = CHAR_LEN / 256

# Scaling factor to resize the image
SCALE = 0.15

# Character dimensions (width and height in pixels) used to match image aspect ratio to character aspect ratio
CHAR_W = 6
CHAR_H = 14

##### FUNCTIONS #####
def getChar(inputInt, gamma=1.8):
    """Map a grayscale pixel value to an ASCII character with gamma correction applied."""
    # Adjust the input pixel intensity using gamma correction for perceptual brightness adjustment
    inputInt = (inputInt / 255) ** gamma * 255
    # Map the corrected pixel value to an appropriate ASCII character
    return CHAR_ARRAY[math.floor(inputInt * GRAYSCALE_LEVEL)]


def load_and_preprocess_image(image):
    """Resize and preprocess the input image, adjusting contrast and blurring for better ASCII conversion."""
    width, height = image.size
    # Resize image, adjusting aspect ratio to fit ASCII character dimensions
    im = image.resize((int(SCALE * width), int(SCALE * height * (CHAR_W / CHAR_H))))

    # Enhance contrast to bring out more detail in the ASCII representation
    im = ImageOps.equalize(im, mask=None)

    # Apply a slight blur to reduce noise and simplify pixel values
    im = im.filter(ImageFilter.GaussianBlur(radius=0.5))

    return im

def create_ascii_art(im):
    """Convert a preprocessed image into ASCII art by mapping grayscale values to ASCII characters."""
    # Convert the image to grayscale
    grayscale_image = im.convert("L")
    # Get the pixel values as a NumPy array for easier processing
    pix = np.array(grayscale_image)
    width, height = grayscale_image.size

    # Create a string that holds the ASCII art, where each character represents a pixel
    ascii_art = ""
    for i in range(height):
        for j in range(width):
            # Append the corresponding ASCII character for each pixel
            ascii_art += getChar(pix[i, j])
        # Newline after each row to maintain image structure
        ascii_art += '\n'
    
    return ascii_art


def draw_ascii_image(ascii_art_string, char_width, char_height, font_size):
    """Draw the ASCII art string onto an image."""
    # Split the ASCII art string into lines
    lines = ascii_art_string.split('\n')
    
    # Determine the dimensions of the image based on the number of characters
    width = max(len(line) for line in lines)  # Maximum line width
    height = len(lines)  # Number of lines (height)
    
    # Create a blank white image based on the ASCII art size and font size
    img_width = width * char_width
    img_height = height * char_height
    ascii_image = Image.new("RGB", (img_width, img_height), "white")
    draw = ImageDraw.Draw(ascii_image)

    # Use default or custom font (Pillow's default font in this case)
    font = ImageFont.load_default()
    # Draw the ASCII art on the image
    for i, line in enumerate(lines):
        # Draw each line of the ASCII art onto the image
        draw.text((0, i * char_height), line, font=font, fill="black")

    return ascii_image


def process_image(image):
    """Process the input image to generate both an ASCII art image and a downloadable text file."""
    # Resize and preprocess the image
    resized_image = load_and_preprocess_image(image)
    # Generate the ASCII art as text
    ascii_art = create_ascii_art(resized_image)

    # Create an image from the ASCII art characters
    output_image = draw_ascii_image(ascii_art, char_width=CHAR_W, char_height=CHAR_H, font_size=10)

    # Save the ASCII art as a text file
    ascii_txt_path = "ascii_art.txt"
    with open(ascii_txt_path, "w") as text_file:
        text_file.write(ascii_art)

    return output_image, ascii_txt_path


##### GRADIO INTERFACE #####
def gradio_interface(image):
    """Gradio interface function to handle user input and return the ASCII art image and text file."""
    ascii_image, txt_file = process_image(image)
    return ascii_image, txt_file

# Set up the Gradio interface
demo = gr.Interface(
    fn=gradio_interface, 
    inputs=gr.Image(type="pil", label="Upload an Image", height=300), 
    outputs=[
        gr.Image(type="pil", label="ASCII Art Image", height=300), 
        gr.File(label="Download ASCII Art Text File", height=50),
    ],
    title="ASCII Art Generator",
    description="Upload an image, and this tool will generate ASCII art and provide a downloadable text file of the result.",
    allow_flagging="never",
    examples=[
        ['images/building.jpg'],
        ['images/cat.webp'],
        ['images/mountain.webp'],
        ['images/people.jpg'],
        ['images/Northeastern_seal.png'],
        ['images/einstein.jpg'],
    ],

)

demo.launch(share=True)