main.py

from pywebio import start_server
from pywebio.output import put_text, use_scope
from pywebio.input import select
from pywebio.output import put_image, put_markdown, put_loading, put_row, put_button
from diversity_search import diversity_reranking, similarity_reranking
from pywebio.session import set_env, run_js
from pywebio.output import put_buttons, put_image, popup, close_popup, put_html
import io
import base64
from data_loader import load_embeddings_and_scores
import numpy as np
from PIL import Image
import time
import os
from datetime import datetime


def pil_to_base64(pil_img, size=(260, 260)):
    """Convert PIL image to base64 with consistent sizing"""
    img_byte_arr = io.BytesIO()
    
    # Convert problematic modes to RGB
    if pil_img.mode in ('RGBA', 'P'):
        pil_img = pil_img.convert('RGB')
    
    # Resize image to consistent dimensions while maintaining aspect ratio
    # Use thumbnail to maintain aspect ratio, then pad to exact size
    original_size = pil_img.size
    pil_img.thumbnail(size, Image.Resampling.LANCZOS)
    
    # Create a new image with the exact target size and white background
    new_img = Image.new('RGB', size, (255, 255, 255))
    
    # Calculate position to center the resized image
    x = (size[0] - pil_img.size[0]) // 2
    y = (size[1] - pil_img.size[1]) // 2
    
    # Paste the resized image onto the white background
    new_img.paste(pil_img, (x, y))
    
    new_img.save(img_byte_arr, format='JPEG', quality=85)
    img_byte_arr = img_byte_arr.getvalue()
    return f"data:image/jpeg;base64,{base64.b64encode(img_byte_arr).decode()}"


def show_final_image_popup(img_idx):
    """Show the final selected image in a pop-up with navigation buttons"""
    global s, S, N, M_values, embeddings, scores, indices, reranked_indices, images, prev_idx
    global exploration_start_time, selected_image_id
    
    if img_idx is not None:
        root = reranked_indices[img_idx]
    else:
        root = None
        print("No root image selected {}".format(img_idx))
    
    # Get the image to display
    final_image = images[root]
    
    def handle_popup_action(action):
        # Close the popup first
        close_popup()
        
        if action == 'end':
            # Calculate exploration time
            exploration_time = 0
            if exploration_start_time is not None:
                exploration_time = time.time() - exploration_start_time
            
            # Set selected image ID
            selected_image_id = root
            
            # Log the session data (use class_name instead of search_query)
            log_exploration_session(exploration_time, selected_image_id, current_class_name, "div")
            
            with use_scope('my_scope', clear=True):
                put_markdown("## Session ended")
                put_markdown("Displaying the selected image")
                
                # Display the main selected image prominently
                put_html(f'''
                <div style="text-align: center; margin: 20px 0;">
                    <h3 style="color: #333; margin-bottom: 15px;">Final selected image</h3>
                    <div style="display: inline-block; padding: 20px; background: white; border-radius: 12px; box-shadow: 0 4px 16px rgba(0,0,0,0.1); border: 3px solid #4CAF50;">
                        <img src="{pil_to_base64(final_image, size=(300, 300))}" width="300" height="300" style="border-radius: 8px;">
                    </div>
                </div>
                ''')
                
                # Add restart button
                def restart_session():
                    # Reload the page to return to query selection
                    put_html('<script>window.location.reload();</script>')
                
                # Center the restart button
                put_html('<div style="text-align: center; margin: 20px 0;">')
                put_button("Start over", onclick=restart_session, color='primary')
                put_html('</div>')
    
    # Create the pop-up
    popup('Final selected image', [
        put_markdown("### This is your final selected image"),
        put_image(pil_to_base64(final_image, size=(400, 400)), width="400px"),
        put_markdown("Press **Confirm** if this is your final choice. Press the X button to continue exploring."),
        put_html('<div style="text-align: center;">'),
        put_buttons([
            dict(label='Confirm', value='end', color='danger'),
        ], onclick=handle_popup_action),
        put_html('</div>')
    ])


def diversity_iter_with_scroll(idx):
    """Wrapper for diversity_iter that includes scroll functionality"""
    diversity_iter(idx)
    # Add scroll to selected query section
    run_js("""
    setTimeout(function() {
        var targetElement = document.getElementById('selected-query-display');
        if (targetElement) {
            targetElement.scrollIntoView({ behavior: 'smooth', block: 'start' });
        }
    }, 100);
    """)

def diversity_iter(idx):
    global s, S, N, M_values, embeddings, scores, indices, reranked_indices, images
    
    print(f"diversity_iter called with idx={idx}, s={s}, S={S}")

    if s < S:

        #progress = s/S # For linear diversity
        progress = progress_bar[s] # For non-linear diversity (exponential decay)
        w_r = weight_r_bar[s]
        
        if idx is not None:
            root = reranked_indices[idx]
        else:
            root = 0
            print("No root image selected {}".format(idx))
        
        # In the last step, show similar images instead of diverse images
        current_M = M_values[s] if s < len(M_values) else M_values[-1]
        if s == S - 1:
            reranked_indices = similarity_reranking(root, base, embeddings, scores, indices, N*current_M)
        else:
            reranked_indices = diversity_reranking(root, base, embeddings, scores, indices, progress, w_d, w_r, N*current_M)
        prev_idx.append(root)
        
        # Convert reranked indices back to images
        images_div = []
        for img_idx in reranked_indices:
            if img_idx < len(images):
                images_div.append(images[img_idx])
            else:
                print(f"❌ WARNING: Index {img_idx} is out of bounds for images array (length: {len(images)})")
        
        s += 1

        with put_loading():
            # Clear and refresh the display
            with use_scope('my_scope', clear=True):
                # Add progress indicator
                put_html(f"""
                <div style="text-align: center; margin: 20px 0;">
                    <div style="display: inline-block; padding: 10px 20px; background: #e3f2fd; border-radius: 20px; color: #1976d2; font-weight: bold;">
                        Step {s}/{S}
                    </div>
                </div>
                """)
                
                # Add navigation buttons if not at the beginning
                if s > 1:
                    put_row([
                        put_button("← Back", onclick=back_iter, color='secondary'),
                        put_button("To top", onclick=back_to_top, color='info')
                    ])
                
                # Re-display the grid of images
                for i in range(0, len(images_div), N):
                    # Display images
                    put_row([
                        put_html(f'''
                        <div style="text-align: center; margin: 6px; width: 280px;">
                            <div class="image-grid-button">
                                <img src="{pil_to_base64(img, size=(280, 280))}" width="280" height="280" style="object-fit: cover;">
                            </div>
                        </div>
                        ''') for j, img in enumerate(images_div[i:i+N])
                    ])
                    
                    # Create button pairs for each image with more space
                    button_pairs = []
                    for j, img in enumerate(images_div[i:i+N]):
                        img_idx = i+j
                        buttons_def = [
                            dict(label="Choose this", value='end', color='danger',
                                 style="font-size: 10px; padding: 4px 8px;")
                        ] if s == S else [
                            dict(label="See more", value='continue', color='success', 
                                 style="font-size: 10px; padding: 4px 8px; margin-right: 2px;"),
                            dict(label="Choose this", value='end', color='danger',
                                 style="font-size: 10px; padding: 4px 8px;")
                        ]
                        button_pair = put_buttons(
                            buttons_def,
                            onclick=lambda choice, img_index=img_idx: (
                                diversity_iter_with_scroll(img_index) if choice == 'continue' 
                                else show_final_image_popup(img_index)
                            )
                        )
                        button_pairs.append(button_pair)
                    put_row(button_pairs)
    else:
        # Show final image in pop-up
        show_final_image_popup(idx)
    print(prev_idx)


def back_to_top():
    """Go back to step 1 directly"""
    global s, S, N, M_values, embeddings, scores, indices, reranked_indices, images, prev_idx
    
    # Reset to step 1
    s = 1
    prev_idx = prev_idx[:1]  # Keep only the initial root
    
    # Use initial parameters for step 1
    progress = progress_bar[0]
    w_r = weight_r_bar[0]
    root_for_this_step = prev_idx[0]
    
    # Regenerate step 1
    current_M = M_values[0]
    reranked_indices = diversity_reranking(root_for_this_step, base, embeddings, scores, indices, progress, w_d, w_r, N*current_M)
    
    # Convert reranked indices back to images
    images_div = []
    for img_idx in reranked_indices:
        if img_idx < len(images):
            images_div.append(images[img_idx])
        else:
            print(f"❌ WARNING (back_to_top): Index {img_idx} is out of bounds for images array (length: {len(images)})")

    with put_loading():
        # Clear and refresh the display
        with use_scope('my_scope', clear=True):
            # Add progress indicator
            put_html(f"""
            <div style="text-align: center; margin: 20px 0;">
                <div style="display: inline-block; padding: 10px 20px; background: #e3f2fd; border-radius: 20px; color: #1976d2; font-weight: bold;">
                    Step {s}/{S}
                </div>
            </div>
            """)
            
            # Add navigation buttons
            put_row([
                put_button("← Back", onclick=back_iter, color='secondary'),
                put_button("To top", onclick=back_to_top, color='info')
            ])
            
            # Re-display the grid of images
            for i in range(0, len(images_div), N):
                # Display images
                put_row([
                    put_html(f'''
                    <div style="text-align: center; margin: 6px; width: 280px;">
                        <div class="image-grid-button">
                            <img src="{pil_to_base64(img, size=(280, 280))}" width="280" height="280" style="object-fit: cover;">
                        </div>
                    </div>
                    ''') for j, img in enumerate(images_div[i:i+N])
                ])
                
                # Create button pairs for each image with more space
                button_pairs = []
                for j, img in enumerate(images_div[i:i+N]):
                    img_idx = i+j
                    buttons_def = [
                        dict(label="Choose this", value='end', color='danger',
                             style="font-size: 10px; padding: 4px 8px;")
                    ] if s == S else [
                        dict(label="See more", value='continue', color='success', 
                             style="font-size: 10px; padding: 4px 8px; margin-right: 2px;"),
                        dict(label="Choose this", value='end', color='danger',
                             style="font-size: 10px; padding: 4px 8px;")
                    ]
                    button_pair = put_buttons(
                        buttons_def,
                        onclick=lambda choice, img_index=img_idx: (
                            diversity_iter_with_scroll(img_index) if choice == 'continue' 
                            else show_final_image_popup(img_index)
                        )
                    )
                    button_pairs.append(button_pair)
                put_row(button_pairs)


def back_iter():
    global s, S, N, M_values, embeddings, scores, indices, reranked_indices, images, prev_idx

    s -= 1
    
    # Re-calculate diversity with previous step
    prev_idx.pop()
    
    # Use the previous step's root (the one that originally generated this step)
    root_for_this_step = prev_idx[-1]

    #progress = s/S # For linear diversity
    progress = progress_bar[s-1] # Use the step that originally created this step
    w_r = weight_r_bar[s-1]

    # When going backward, always use diversity reranking (never similarity)
    current_M = M_values[s-1]
    reranked_indices = diversity_reranking(root_for_this_step, base, embeddings, scores, indices, progress, w_d, w_r, N*current_M)

    # Convert reranked indices back to images
    images_div = []
    for img_idx in reranked_indices:
        if img_idx < len(images):
            images_div.append(images[img_idx])
        else:
            print(f"❌ WARNING (back_iter): Index {img_idx} is out of bounds for images array (length: {len(images)})")

    with put_loading():
        # Clear previous elements
        with use_scope('my_scope', clear=True):
            # Add progress indicator
            put_html(f"""
            <div style="text-align: center; margin: 20px 0;">
                <div style="display: inline-block; padding: 10px 20px; background: #e3f2fd; border-radius: 20px; color: #1976d2; font-weight: bold;">
                    Step {s}/{S}
                </div>
            </div>
            """)
            
            # Add navigation buttons if not at the beginning
            if s > 1:
                put_row([
                    put_button("← Back", onclick=back_iter, color='secondary'),
                    put_button("To top", onclick=back_to_top, color='info')
                ])
            
            # Re-display the grid of images
            for i in range(0, len(images_div), N):
                # Display images
                put_row([
                    put_html(f'''
                    <div style="text-align: center; margin: 6px; width: 280px;">
                        <div class="image-grid-button">
                            <img src="{pil_to_base64(img, size=(280, 280))}" width="280" height="280" style="object-fit: cover;">
                        </div>
                    </div>
                    ''') for j, img in enumerate(images_div[i:i+N])
                ])
                
                # Create button pairs for each image with more space
                button_pairs = []
                for j, img in enumerate(images_div[i:i+N]):
                    img_idx = i+j
                    buttons_def = [
                        dict(label="Choose this", value='end', color='danger',
                             style="font-size: 10px; padding: 4px 8px;")
                    ] if s == S else [
                        dict(label="See more", value='continue', color='success', 
                             style="font-size: 10px; padding: 4px 8px; margin-right: 2px;"),
                        dict(label="Choose this", value='end', color='danger',
                             style="font-size: 10px; padding: 4px 8px;")
                    ]
                    button_pair = put_buttons(
                        buttons_def,
                        onclick=lambda choice, img_index=img_idx: (
                            diversity_iter_with_scroll(img_index) if choice == 'continue' 
                            else show_final_image_popup(img_index)
                        )
                    )
                    button_pairs.append(button_pair)
                put_row(button_pairs)
    print(prev_idx)


s = 0 # Current step
S = 4 # Total steps
w_d = 1 # Diversity weight
W_r = 12 # Relevance weight
decay_rate = 8
N = 4 # Grid columns
M_values = [4, 4, 4, 4] # Grid rows for each step (decreasing)
base = "gram" # "gram", "vendi", "magnitude"
progress_bar = np.exp(-decay_rate * np.linspace(0, 1, S))
weight_r_bar = []
for i in range(S):
    #weight_r_bar.append(W_r/(i+1))
    weight_r_bar.append(W_r)

# Global variables for loaded data
embeddings = None
scores = None
indices = None
reranked_indices = None
images = None
prev_idx = []

# Global variables for logging
exploration_start_time = None
selected_image_id = None
current_class_name = None


def log_exploration_session(exploration_time, image_id, query, program_type="div"):
    """Log exploration session data to a query-specific file"""
    # Create logs directory if it doesn't exist
    os.makedirs('./logs', exist_ok=True)
    
    # Create query-specific log file name with program type prefix
    safe_query = query.replace(" ", "_").replace("/", "_").lower()
    log_file = f'./logs/{program_type}_{safe_query}_exploration_log.txt'
    
    # Get the next ID number by counting existing lines
    session_id = 1
    if os.path.exists(log_file):
        with open(log_file, 'r', encoding='utf-8') as f:
            session_id = sum(1 for line in f if line.strip()) + 1
    
    # Create log entry with sequential ID
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_entry = f"{session_id}, {timestamp}, Image ID: {image_id}, Exploration Time: {exploration_time:.2f}s\n"
    
    # Append to query-specific log file
    with open(log_file, 'a', encoding='utf-8') as f:
        f.write(log_entry)


def main():
    # Global variables
    global s, S, N, M_values, embeddings, scores, indices, reranked_indices, images, prev_idx
    global exploration_start_time, selected_image_id, current_class_name
    
    # Reinitialize global variables to ensure clean start
    s = 0
    prev_idx = []
    embeddings = None
    scores = None
    indices = None
    reranked_indices = None
    images = None
    exploration_start_time = None
    selected_image_id = None
    current_class_name = None

    set_env(
        title="Image Exploration Demo",
        output_animation=False,
        output_max_width="1280px",
    )
    
    # Add global CSS styling for consistent, clean interface
    put_html("""
    <style>
    .image-grid-button {
        margin: 4px;
        padding: 6px;
        border: 2px solid #e0e0e0;
        border-radius: 8px;
        background: white;
        transition: all 0.3s ease;
        display: inline-block;
        box-shadow: 0 2px 4px rgba(0,0,0,0.05);
    }
    .image-grid-button:hover {
        border-color: #4CAF50;
        box-shadow: 0 4px 12px rgba(0,0,0,0.15);
        transform: translateY(-2px);
    }
    .image-grid-button img {
        display: block;
        border-radius: 4px;
    }
    .btn-group .btn {
        font-size: 10px !important;
        padding: 4px 8px !important;
        margin: 1px !important;
    }
    .pywebio-output {
        padding: 20px;
    }
    h1, h2, h3 {
        color: #333;
        margin-bottom: 15px;
    }
    .btn {
        margin: 5px;
        padding: 10px 20px;
        border-radius: 5px;
    }
    </style>
    """)

    put_html("""
    <div style="text-align: center; margin: 30px 0;">
        <h1 style="color: #333; font-size: 2.5em; margin-bottom: 10px;">Image Exploration Demo</h1>
        <p style="color: #666; font-size: 1.1em; margin-bottom: 30px;">
            Search for diverse relevant images
        </p>
    </div>
    <div style="max-width: 600px; margin: 0 auto; padding: 20px; background: #f9f9f9; border-radius: 10px; box-shadow: 0 2px 8px rgba(0,0,0,0.1);">
        <h2 style="color: #444; margin-bottom: 15px;">Select a task</h2>
        <p style="color: #666; margin-bottom: 15px;">
            Please select a design task below
        </p>
    </div>
    """)

    search_query = select("Select a design task", [
        "Mother's Day",
    ])
    

    # Map queries to class names and search queries
    query_to_class = {
        "Mother's Day": "mothers_day",
    }

    query_to_query = {
        "mothers_day": "a_mothers_day_celebration_for_young_moms_that_conv",
    }
    
    # Get the appropriate class name for the selected query
    class_name = query_to_class.get(search_query)
    query_name = query_to_query.get(class_name)
    
    # Store class_name globally for logging
    current_class_name = class_name
    
    if search_query:
        # Add ID to the section for scrolling
        put_html('<div id="selected-query-display"></div>')
        put_markdown("### Selected task")
        put_text(search_query)
        
        # Start exploration timer
        exploration_start_time = time.time()

    with put_loading():
        # Load data with the search query and appropriate class name
        data = load_embeddings_and_scores(
            query=query_name,
            embeddings_dir='./embeddings',
            scores_dir='./config/scores',
            class_name=class_name,
            model_type='clip')
        
        embeddings = data['embeddings']
        scores = data['scores']
        indices = list(range(len(embeddings)))  #data['sorted_indices']
        images = data['images']
        
        # DEBUG: Check data consistency
        print(f"🔍 Data consistency check:")
        print(f"   - Embeddings shape: {embeddings.shape if embeddings is not None else 'None'}")
        print(f"   - Scores length: {len(scores) if scores is not None else 'None'}")
        print(f"   - Indices length: {len(indices)}")
        print(f"   - Images length: {len(images)}")
        print(f"   - Images with None: {sum(1 for img in images if img is None)}")
        print(f"   - Valid images: {sum(1 for img in images if img is not None)}")

    # Initialize the display with diversity-ranked results
    progress = progress_bar[s] # For non-linear diversity (exponential decay)
    w_r = weight_r_bar[s]
    
    # Get initial diversity-ranked results
    initial_M = M_values[0] if len(M_values) > 0 else 5
    reranked_indices = diversity_reranking(0, base, embeddings, scores, indices, progress, w_d, w_r, N*initial_M)
    prev_idx = [reranked_indices[0]]
    
    # Convert reranked indices to images
    images_div = []
    for img_idx in reranked_indices:  # Get first N*N images
        if img_idx < len(images):
            images_div.append(images[img_idx])
        else:
            print(f"❌ WARNING (main): Index {img_idx} is out of bounds for images array (length: {len(images)})")
    
    print(f"🎯 Initial display: Created {len(images_div)} images from {len(reranked_indices)} reranked indices")
    s += 1
    
    with use_scope('my_scope'):
        # Add progress indicator
        put_html(f"""
        <div style="text-align: center; margin: 20px 0;">
            <div style="display: inline-block; padding: 10px 20px; background: #e3f2fd; border-radius: 20px; color: #1976d2; font-weight: bold;">
                Step {s}/{S}
            </div>
        </div>
        """)
        
        # Add navigation buttons if not at the beginning
        if s > 1:
            put_row([
                put_button("← Back", onclick=back_iter, color='secondary'),
                put_button("To top", onclick=back_to_top, color='info')
            ])
        
        # Re-display the grid of images
        for i in range(0, len(images_div), N):
            # Display images
            put_row([
                put_html(f'''
                <div style="text-align: center; margin: 6px; width: 280px;">
                    <div class="image-grid-button">
                        <img src="{pil_to_base64(img, size=(280, 280))}" width="280" height="280" style="object-fit: cover;">
                    </div>
                </div>
                ''') for j, img in enumerate(images_div[i:i+N])
            ])
            
            # Create button pairs for each image with more space
            button_pairs = []
            for j, img in enumerate(images_div[i:i+N]):
                img_idx = i+j
                buttons_def = [
                    dict(label="Choose this", value='end', color='danger',
                         style="font-size: 10px; padding: 4px 8px;")
                ] if s == S else [
                    dict(label="See more", value='continue', color='success', 
                         style="font-size: 10px; padding: 4px 8px; margin-right: 2px;"),
                    dict(label="Choose this", value='end', color='danger',
                         style="font-size: 10px; padding: 4px 8px;")
                ]
                button_pair = put_buttons(
                    buttons_def,
                    onclick=lambda choice, img_index=img_idx: (
                        diversity_iter_with_scroll(img_index) if choice == 'continue' 
                        else show_final_image_popup(img_index)
                    )
                )
                button_pairs.append(button_pair)
            put_row(button_pairs)
    print(prev_idx)


if __name__ == "__main__":
    start_server(main, port=8080, debug=True)