import gradio as gr

import torch

from fastai.text.all import *

import warnings

warnings.filterwarnings('ignore')



# Load the trained model

def load_model():

    try:

        learn = load_learner('js_concept_tagger_simple_numbers.pkl')

        return learn

    except Exception as e:

        print(f"Error loading model: {e}")

        return None



# Initialize model

model = load_model()



def predict_js_concept(code_snippet):

    """Predict JavaScript concept for given code snippet"""

    if model is None:

        return "Error: Model not loaded", 0.0, "Model loading failed"

    

    try:

        # Make prediction

        pred, pred_idx, probs = model.predict(code_snippet)

        

        # Get the predicted concept and confidence

        predicted_concept = str(pred)

        confidence = float(probs.max())

        

        # Get all predictions with probabilities

        concept_names = model.dls.vocab[1]  # label vocabulary

        all_predictions = []

        

        for i, prob in enumerate(probs):

            if prob > 0.1:  # Show predictions above 10%

                all_predictions.append(f"{concept_names[i]}: {float(prob):.3f}")

        

        all_predictions_str = "\n".join(all_predictions)

        

        return predicted_concept, confidence, all_predictions_str

        

    except Exception as e:

        return f"Error: {str(e)}", 0.0, "Prediction failed"



# Create Gradio interface

def create_interface():

    # Example code snippets for users to try

    examples = [

        ["const age = 25;"],

        ["const name = 'John Doe';"],

        ["const isActive = true;"],

        ["const items = [1, 2, 3];"],

        ["function greet() { return 'Hello'; }"],

        ["if (x > 0) { console.log('positive'); }"],

        ["for (let i = 0; i < 10; i++) { }"],

        ["arr.push(newItem);"],

        ["const sum = 10 + 20;"],

        ["let message = 'Welcome to our app';"]

    ]

    

    interface = gr.Interface(

        fn=predict_js_concept,

        inputs=[

            gr.Textbox(

                label="JavaScript Code Snippet",

                placeholder="Enter your JavaScript code here...",

                lines=3,

                max_lines=10

            )

        ],

        outputs=[

            gr.Textbox(label="Predicted Concept", interactive=False),

            gr.Number(label="Confidence Score", precision=3),

            gr.Textbox(label="All Predictions", lines=6, interactive=False)

        ],

        title="🔍 JavaScript Concept Classifier",

        description="""

        This AI model classifies JavaScript code snippets into different programming concepts:

        

        **Supported Concepts:**

        - 🔢 **Numbers**: Numeric values and assignments

        - 📝 **Strings**: Text and string literals  

        - ✅ **Booleans**: True/false values

        - 📋 **Arrays**: Array operations and methods

        - ⚡ **Functions**: Function definitions and calls

        - 🔄 **Control Flow**: If statements, loops, switches

        

        **How to use:** Simply paste your JavaScript code in the text box below and click Submit!

        """,

        examples=examples,

        theme=gr.themes.Soft(),

        allow_flagging="never"

    )

    

    return interface



# Launch the interface

if __name__ == "__main__":

    demo = create_interface()

    demo.launch()