"""
Created Date: 09-26-2024
Updated Date: -
Author: Chaitanya Chadha
"""

import streamlit as st
from transformers import AutoImageProcessor, ResNetForImageClassification
import torch
from PIL import Image, ImageDraw, ImageFont
import math
import os
import io
from datetime import datetime

# Set page configuration at the very beginning
st.set_page_config(page_title="🎨 Colored ASCII Art Generator", layout="wide")

# Initialize model and processor once to improve performance
@st.cache_resource
def load_model_and_processor():
    processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
    model = ResNetForImageClassification.from_pretrained("microsoft/resnet-50")
    return processor, model

processor, model = load_model_and_processor()

def Classify_Image(image):
    """
    Classifies the image using a pre-trained ResNet-50 model.
    Returns the predicted label.
    """
    inputs = processor(image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    # model predicts one of the 1000 ImageNet classes
    predicted_label = logits.argmax(-1).item()
    return model.config.id2label[predicted_label]

def resize_image(image, new_width=100):
    width, height = image.size
    aspect_ratio = height / width
    # Adjusting height based on the aspect ratio and a scaling factor
    new_height = int(aspect_ratio * new_width * 0.55)
    resized_image = image.resize((new_width, new_height))
    return resized_image

def grayify(image):
    return image.convert("L")

def calculate_image_entropy(image):
    # Calculates the entropy of the grayscale image to determine its complexity.
    # Higher entropy indicates a more complex image with more details.
    histogram = image.histogram()
    histogram_length = sum(histogram)

    samples_probability = [float(h) / histogram_length for h in histogram if h != 0]
    entropy = -sum([p * math.log(p, 2) for p in samples_probability])

    return entropy

def select_character_set(entropy, art_gen_choice, classification):
    ASCII_CHARS_SETS = {
        "standard": [
            '@', '%', '#', '*', '+', '=', '-', ':', '.', ' '
        ],
        "detailed": [
            '$', '@', 'B', '%', '8', '&', 'W', 'M', '#', '*', 'o', 'a', 'h', 'k', 'b',
            'd', 'p', 'q', 'w', 'm', 'Z', 'O', '0', 'Q', 'L', 'C', 'J', 'U', 'Y',
            'X', 'z', 'c', 'v', 'u', 'n', 'x', 'r', 'j', 'f', 't', '/', '\\', '|',
            '(', ')', '1', '{', '}', '[', ']', '?', '-', '_', '+', '~', '<', '>',
            'i', '!', 'l', 'I', ';', ':', ',', '"', '^', '', "'", '.', ' '
        ],
        "simple": [
            '#', '*', '+', '=', '-', ':', '.', ' '
        ],
        "custom": [
            "!", "~", "@", "#", "$", "%", "¨", "&", "*", "(", ")", "_", "+", "-",
            "=", "{", "}", "[", "]", "|", "\\", "/", ":", ";", "'", "\"", ",", "<",
            ".", ">", "?", " " ] + list(set(list(str(classification))))
    }

    if art_gen_choice.lower() == "custom":
        selected_set = "custom"
        return ASCII_CHARS_SETS[selected_set]
    else:
        # Define entropy thresholds (these values can be adjusted based on experimentation)
        if entropy < 4.0:
            selected_set = "simple"
        elif 4.0 <= entropy < 5.5:
            selected_set = "standard"
        else:
            selected_set = "detailed"

        return ASCII_CHARS_SETS[selected_set]

def determine_optimal_width(image, max_width=120):
    # Determines the optimal width for the ASCII art based on the image's original dimensions.
    original_width, original_height = image.size
    if original_width > max_width:
        return max_width
    else:
        return original_width

def render_ascii_to_image(ascii_chars, img_width, img_height, font_path="fonts/DejaVuSansMono.ttf", font_size=10):
    # Renders the ASCII characters onto an image with their corresponding colors.
    if not os.path.isfile(font_path):
        st.error(f"Font file not found at {font_path}. Please provide a valid font path.")
        return None

    # Create a new image with white background
    try:
        font = ImageFont.truetype(font_path, font_size)
    except Exception as e:
        st.error(f"Error loading font: {e}")
        return None

    # Get character dimensions
    # Using getbbox which is compatible with Pillow >= 8.0
    left, top, right, bottom = font.getbbox('A')
    char_width = right - left
    char_height = bottom - top
    image_width = char_width * img_width
    image_height = char_height * img_height
    new_image = Image.new("RGB", (image_width, image_height), "white")
    draw = ImageDraw.Draw(new_image)

    for i, (char, color) in enumerate(ascii_chars):
        x = (i % img_width) * char_width
        y = (i // img_width) * char_height
        draw.text((x, y), char, fill=color, font=font)

    return new_image

def pixels_to_colored_ascii(grayscale_image, color_image, chars):
    # Maps each pixel to an ASCII character from the selected character set.
    # Returns a list of tuples: (character, (R, G, B))
    grayscale_pixels = grayscale_image.getdata()
    color_pixels = color_image.getdata()
    ascii_chars = []

    for grayscale_pixel, color_pixel in zip(grayscale_pixels, color_pixels):
        # Scale grayscale pixel value to the range of the character set
        index = grayscale_pixel * (len(chars) - 1) // 255
        ascii_char = chars[index]

        # Extract RGB values; ignore alpha if present
        if len(color_pixel) == 4:
            r, g, b, _ = color_pixel
        else:
            r, g, b = color_pixel

        ascii_chars.append((ascii_char, (r, g, b)))

    return ascii_chars

def generate_ascii_art(image, font_path="fonts/DejaVuSansMono.ttf", font_size=12):
    # image: PIL Image object

    # Classify the image
    image_class = Classify_Image(image)
    st.write(f"**Image Classification:** {image_class}")

    # Convert to grayscale and calculate entropy
    grayscale_image = grayify(image)
    entropy = calculate_image_entropy(grayscale_image)
    st.write(f"**Image Entropy:** {entropy:.2f}")

    # Select character set
    art_gen_choice = st.session_state.get('art_gen_choice', 'custom')
    selected_chars = select_character_set(entropy, art_gen_choice, image_class)
    st.write(f"**Selected Character Set:** '{art_gen_choice}' with {len(selected_chars)} characters.")

    # Determine optimal width
    optimal_width = determine_optimal_width(image)
    st.write(f"**Selected Width:** {optimal_width}")

    # Resize images
    resized_image = resize_image(image, optimal_width)
    grayscale_resized_image = grayify(resized_image)
    color_resized_image = resized_image.convert("RGB")  # Ensure image is in RGB mode

    # Convert pixels to ASCII
    ascii_chars = pixels_to_colored_ascii(grayscale_resized_image, color_resized_image, selected_chars)

    # Create ASCII string
    ascii_str = ''.join([char for char, color in ascii_chars])
    ascii_lines = [ascii_str[index: index + optimal_width] for index in range(0, len(ascii_str), optimal_width)]
    ascii_art = "\n".join(ascii_lines)

    return ascii_art, ascii_chars, resized_image.size

def main():
    # Title and Description are already set after set_page_config
    st.title("🎨 Colored ASCII Art Generator")

    st.markdown("""
    Upload an image, and this app will convert it into colored ASCII art. 
    You can view the classification, entropy, and download or copy the ASCII art.
    """)

    # Sidebar for options
    st.sidebar.header("Options")
    art_gen_choice = st.sidebar.selectbox(
        "Character Set Choice",
        options=["custom", "standard", "detailed", "simple"],
        help="Select the character set to use for ASCII art generation."
    )
    st.session_state['art_gen_choice'] = art_gen_choice

    # File uploader
    uploaded_file = st.file_uploader("Upload an Image", type=["png", "jpg", "jpeg", "bmp", "gif"])

    if uploaded_file is not None:
        try:
            # Read the image
            image = Image.open(uploaded_file).convert("RGB")
            st.image(image, caption='Uploaded Image', use_column_width=True)

            # Generate ASCII Art
            with st.spinner("Generating ASCII Art..."):
                ascii_art, ascii_chars, resized_size = generate_ascii_art(image)

            # Display ASCII Art as Image
            ascii_image = render_ascii_to_image(
                ascii_chars,
                img_width=resized_size[0],
                img_height=resized_size[1],
                font_path="fonts/DejaVuSansMono.ttf",
                font_size=12
            )

            if ascii_image:
                st.image(ascii_image, caption='Colored ASCII Art', use_column_width=True)

                # Provide Download Button for ASCII Image
                img_byte_arr = io.BytesIO()
                ascii_image.save(img_byte_arr, format='PNG')
                img_byte_arr = img_byte_arr.getvalue()

                st.download_button(
                    label="📥 Download ASCII Art Image",
                    data=img_byte_arr,
                    file_name=f"ascii_art_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png",
                    mime="image/png"
                )

            # Display ASCII Art as Text with Copy Option
            st.text_area("ASCII Art", ascii_art, height=300)

            # Provide a download button for the ASCII text
            st.download_button(
                label="📄 Download ASCII Art Text",
                data=ascii_art,
                file_name=f"ascii_art_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
                mime="text/plain"
            )

        except Exception as e:
            st.error(f"An error occurred: {e}")
    else:
        st.info("Please upload an image to get started.")

if __name__ == "__main__":
    main()