import gradio as gr
import torch
import clip
from torchvision import transforms
from style_net import StyleNet
from dataset import FlickrStreamer
from clip_styler import style_transfer

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

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

# 1. LOAD MODELS
# A. Editing Model
#style_net = StyleNet().to(DEVICE)

# B. Retrieval Model (CLIP)
#clip_model, preprocess = clip.load("ViT-B/32", device=DEVICE)

# 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

# --- FUNCTION: EDIT ---
tf_edit = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

def edit_image(image, style_choice):
    if image is None: return None
    
    detail_prompt = {
        "Sketch": "charcoal sketch",
        "Van Gogh": "Van Gogh painting",
        "Cyberpunk": "neon cyberpunk",
    }
    WEIGHTS_MAP = {
        "Sketch": "unet_charcoal-sketch.pth",
        "Van Gogh": "unet_van-gogh-painting.pth",
        "Cyberpunk": "unet_neon-cyberpunk.pth"
    }
    '''
    try:
        style_net.load_state_dict(torch.load(weights_map.get(style_choice), map_location=DEVICE))
    except Exception:
        return image 
    
    input_tensor = tf_edit(image).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        output_tensor = style_net(input_tensor)
        
    output_img = transforms.ToPILImage()(output_tensor.squeeze().cpu())
    return output_img
    '''
    return style_transfer(image, detail_prompt[style_choice], weights_path=WEIGHTS_MAP[style_choice])

# --- 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)