import gradio as gr
from PIL import Image
import numpy as np

def image_to_ascii(image, width=100, chars="@%#*+=-:. "):
    """
    Convert image to ASCII art
    """
    # Convert image to grayscale
    if image.mode != 'L':
        image = image.convert('L')
    
    # Resize image while maintaining aspect ratio
    original_width, original_height = image.size
    aspect_ratio = original_height / original_width
    height = int(width * aspect_ratio * 0.55)  # 0.55 to account for character aspect ratio
    
    # Resize image
    img_resized = image.resize((width, height))
    
    # Convert to numpy array
    pixels = np.array(img_resized)
    
    # Normalize pixels to 0-1 range
    pixels = pixels / 255.0
    
    # Map pixels to ASCII characters
    ascii_chars = list(chars)
    num_chars = len(ascii_chars)
    
    # Create ASCII art
    ascii_art = []
    for row in pixels:
        line = ''.join(ascii_chars[min(int(pixel * num_chars), num_chars - 1)] for pixel in row)
        ascii_art.append(line)
    
    return '\n'.join(ascii_art)

def process_image(image, width, character_set):
    """
    Process the uploaded image and convert to ASCII art
    """
    if image is None:
        return "Please upload an image first!"
    
    try:
        # Define character sets
        char_sets = {
            "Detailed (dense)": "@%#*+=-:. ",
            "Simple (sparse)": "@#*+-. ",
            "Blocks": "█▓▒░ ",
            "Inverted": " .:-=+*#%@",
            "Custom": character_set
        }
        
        chars = char_sets.get(character_set, "@%#*+=-:. ")
        
        ascii_result = image_to_ascii(image, width, chars)
        return f"```\n{ascii_result}\n```"
    except Exception as e:
        return f"Error processing image: {str(e)}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Image to ASCII Art Converter")
    gr.Markdown("Convert your images to ASCII art! Upload an image, adjust the width and character set, and see the text-based version of your image.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image")
            width_slider = gr.Slider(minimum=50, maximum=200, value=100, step=10, label="ASCII Art Width")
            char_set_dropdown = gr.Dropdown(
                choices=["Detailed (dense)", "Simple (sparse)", "Blocks", "Inverted", "Custom"],
                value="Detailed (dense)",
                label="Character Set"
            )
            custom_chars = gr.Textbox(
                value="@%#*+=-:. ",
                label="Custom Character Set (used when 'Custom' is selected)",
                placeholder="Enter characters from darkest to lightest"
            )
            submit_btn = gr.Button("Convert to ASCII")
        
        with gr.Column():
            output_text = gr.Textbox(
                label="ASCII Art Result",
                lines=20,
                max_lines=50,
                show_copy_button=True
            )
    
    # Update character set based on dropdown selection
    def update_char_set(character_set, custom_chars):
        char_sets = {
            "Detailed (dense)": "@%#*+=-:. ",
            "Simple (sparse)": "@#*+-. ",
            "Blocks": "█▓▒░ ",
            "Inverted": " .:-=+*#%@",
            "Custom": custom_chars
        }
        return char_sets.get(character_set, "@%#*+=-:. ")
    
    # Process function that uses the selected character set
    def process_with_char_set(image, width, character_set, custom_chars):
        char_sets = {
            "Detailed (dense)": "@%#*+=-:. ",
            "Simple (sparse)": "@#*+-. ",
            "Blocks": "█▓▒░ ",
            "Inverted": " .:-=+*#%@",
            "Custom": custom_chars
        }
        chars = char_sets.get(character_set, "@%#*+=-:. ")
        
        if image is None:
            return "Please upload an image first!"
        
        try:
            ascii_result = image_to_ascii(image, width, chars)
            return f"\n{ascii_result}\n"
        except Exception as e:
            return f"Error processing image: {str(e)}"
    
    # Connect the components
    char_set_dropdown.change(
        update_char_set,
        inputs=[char_set_dropdown, custom_chars],
        outputs=custom_chars
    )
    
    submit_btn.click(
        process_with_char_set,
        inputs=[image_input, width_slider, char_set_dropdown, custom_chars],
        outputs=[output_text]
    )

if __name__ == "__main__":
    demo.launch()