import gradio as gr
import torch
import clip
from torchvision import transforms
from style_net import StyleNet
from dataset import FlickrStreamer
from util import DEVICE
from clip_styler import ClipStyler

# --- CONFIGURATION ---
#DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
INDEX_PATH = "flickr_embeddings.pt"
SEARCH_LIMIT = 5000  # Must match the indexer limit!
# ---------------------

UI_EPOCHS = 100
STYLE_PRESETS = {
    "Sketch": {
        "prompt": "charcoal sketch style",
        "source": "a Photo",
    },
    "Van Gogh": {
        "prompt": "Van Gogh painting style",
        "source": "a Photo",
    },
    "Cyberpunk": {
        "prompt": "neon cyberpunk style",
        "source": "a Photo",
    },
}

print(f"System starting on {DEVICE}...")

# 1. LOAD MODELS
# A. Retrieval Model (CLIP)
clip_model, _ = clip.load("ViT-B/32", device=DEVICE)

# B. Editing Model (CLIP-guided optimizer)
clip_styler = ClipStyler(device=DEVICE, num_epochs=UI_EPOCHS)

# 2. LOAD SEARCH INDEX
print("Loading Search Index...")
try:
    image_features = torch.load(INDEX_PATH, map_location=DEVICE).float()
    print(f"Index loaded: {image_features.shape[0]} images.")
except FileNotFoundError:
    print("WARNING: 'flickr_embeddings.pt' not found. Run indexer.py first!")
    image_features = None

# 3. CONNECT TO DATASET (For fetching the actual images)
# Note: No transform here because we want to display the original PIL image
dataset_raw = FlickrStreamer(limit=SEARCH_LIMIT, transform=None)

# --- FUNCTION: SEARCH ---
def search_flickr(text_query):
    if image_features is None:
        return []
    
    with torch.no_grad():
        # 1. Encode Text
        text_tokens = clip.tokenize([text_query]).to(DEVICE)
        text_features = clip_model.encode_text(text_tokens)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        
        # 2. Math: Cosine Similarity (Dot Product)
        # (1, 512) x (5000, 512).T  -> (1, 5000) similarity scores
        similarity = (100.0 * text_features @ image_features.T).softmax(dim=-1)
        
        # 3. Get Top 5 matches
        values, indices = similarity[0].topk(5)
        
        # 4. Retrieve actual images from dataset
        # (Deep Lake might be slightly slow here as it fetches specific indices)
        results = []
        for idx in indices:
            img_idx = idx.item()
            results.append(dataset_raw[img_idx])
            
        return results

def edit_image(image, style_choice):
    if image is None or style_choice is None:
        return None

    preset = STYLE_PRESETS.get(style_choice)
    if preset is None:
        return image

    try:
        return clip_styler.stylize(
            image=image,
            prompt=preset["prompt"],
            source=preset.get("source", "a Photo"),
            num_epochs=UI_EPOCHS,
            log_interval=0,
        )
    except Exception as exc:
        print(f"[edit_image] Failed to stylize image: {exc}")
        return image 

# --- THE UI ---
with gr.Blocks() as demo:
    gr.Markdown("# Group 19: CLIP Search & Edit Engine")
    
    with gr.Tab("Retrieval"):
        gr.Markdown(f"Search through {SEARCH_LIMIT} Flickr images by concept.")
        with gr.Row():
            txt_query = gr.Textbox(label="Type a query (e.g., 'a dog on a boat', 'neon lights')")
            btn_search = gr.Button("Search Library")
        
        # Output is a Gallery (Grid of images)
        gallery_out = gr.Gallery(label="Top Results", columns=5, height=300)
        
        btn_search.click(search_flickr, inputs=txt_query, outputs=gallery_out)
        
    with gr.Tab("Image Editing"):
        gr.Markdown("Transform any image using our trained CLIPstyler models.")
        with gr.Row():
            im_in = gr.Image(type="pil", label="Input Image")
            style = gr.Dropdown(["Sketch", "Van Gogh", "Cyberpunk"], label="Choose Style")
        
        btn_edit = gr.Button("Apply")
        im_out = gr.Image(label="Result")
        
        btn_edit.click(edit_image, inputs=[im_in, style], outputs=im_out)

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