import gradio as gr
import pandas as pd
import numpy as np
import re
import string
from sentence_transformers import SentenceTransformer, util

# Load model
model = SentenceTransformer("all-mpnet-base-v2")

# Preprocessing
def preprocess(text):
    text = text.lower()
    text = re.sub(f"[{re.escape(string.punctuation)}]", "", text)
    return text

# Load dataset
df = pd.read_csv("leetcode_dataset - lc.csv")
df = df.drop(columns=['id', 'is_premium', 'difficulty', 'solution_link', 'acceptance_rate', 
                      'frequency', 'discuss_count', 'accepted', 'submissions', 
                      'companies', 'related_topics', 'likes', 'dislikes', 'asked_by_faang'])
descriptions = [preprocess(d) for d in df['description']]
embeddings = model.encode(descriptions, convert_to_tensor=True)

# Search logic
def find_similar_problems(problem_statement):
    if not problem_statement.strip():
        return "Please enter a problem statement.", None
    try:
        query_embedding = model.encode(preprocess(problem_statement), convert_to_tensor=True)
        scores = util.cos_sim(query_embedding, embeddings)[0].cpu().numpy()
        top_k_idx = np.argsort(scores)[-5:][::-1]

        similar = df.iloc[top_k_idx][['title', 'description', 'url']]
        formatted = "\n\n".join(
            f"**{row['title']}**\n{row['description']}\n[View on LeetCode]({row['url']})"
            for _, row in similar.iterrows()
        )
        return "", formatted
    except Exception as e:
        return str(e), None

# Gradio interface
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray", radius_size="lg")) as demo:
    gr.Markdown(
        """
        # 🔍 LeetCode Problem Finder
        Find LeetCode problems that match your coding challenge description.
        """,
        elem_classes=["header"]
    )
    
    with gr.Tabs():
        with gr.Tab("Search Problems"):
            with gr.Row(variant="panel"):
                input_box = gr.Textbox(
                    placeholder="Enter your coding problem description (e.g., 'Find the longest substring without repeating characters')",
                    lines=4,
                    label="Problem Statement",
                    show_label=True,
                    elem_classes=["input-text"]
                )
            with gr.Row():
                search_button = gr.Button("Search Problems", variant="primary", size="lg")
                clear_button = gr.Button("Clear", variant="secondary", size="sm")
            
            error_box = gr.Markdown("", elem_classes=["error-box"], visible=False)
            with gr.Row(variant="panel"):
                results_box = gr.Markdown(
                    "Enter a problem statement above and click 'Search Problems' to see matching LeetCode challenges.",
                    elem_classes=["results-box"]
                )
            
            def wrapper(query):
                err, res = find_similar_problems(query)
                return err, res if res else "No results found.", gr.update(visible=bool(err))

            def clear_input():
                return "", "", "", gr.update(visible=False)

            search_button.click(
                wrapper,
                inputs=input_box,
                outputs=[error_box, results_box, error_box]
            )
            clear_button.click(
                clear_input,
                inputs=None,
                outputs=[input_box, results_box, error_box, error_box]
            )
        
        with gr.Tab("About"):
            gr.Markdown(
                """
                ## About LeetCode Finder
                LeetCode Finder helps you discover relevant LeetCode problems based on your problem description.  
                Powered by the `all-mpnet-base-v2` model, it uses semantic search to match your query with problem descriptions.  
                Ideal for coding interview preparation, algorithm practice, or finding similar challenges.

                ### How It Works
                1. Enter a problem statement in the search tab.
                2. Click "Search Problems" to get the top 5 matching LeetCode problems.
                3. Explore problem details and visit LeetCode links for solutions.

                Built with ❤️ using Gradio and Sentence Transformers.
                """,
                elem_classes=["about-section"]
            )
    
    gr.Markdown(
        "© 2025 LeetCode Finder | Crafted for the coding community",
        elem_classes=["footer"]
    )

demo.launch()