import gradio as gr
from PIL import Image
from vit_model_test import CustomModel  # Ensure you import the correct class
from Vit_Traning.py import Custom_VIT_Model

# Initialize the model
model = Custom_VIT_Model()

# Variable to store the last prediction result
last_prediction = None

def predict(image: Image.Image):
    global last_prediction
    label, confidence = model.predict(image)
    result = "AI image" if label == 1 else "Real image"
    last_prediction = (image, label)  # Store the image and prediction label
    return result, f"Confidence: {confidence:.2f}%"

def report_feedback():
    global last_prediction
    if last_prediction is not None:
        image, predicted_label = last_prediction
        correct_label = 1 if predicted_label == 0 else 0  # Invert the label
        model.add_data(image, correct_label)  # Add incorrect prediction to model
        return "Feedback recorded. Thank you!"
    return "No prediction available to report."

# Define the Gradio interface for prediction and feedback
def main():
    with gr.Blocks() as demo:
        gr.Markdown("### Vision Transformer Model")
        gr.Markdown("Upload an image to classify it using the Vision Transformer model.")
        
        image_input = gr.Image(type="pil", label="Upload Image")
        prediction_output = gr.Textbox(label="Prediction", interactive=False)
        confidence_output = gr.Textbox(label="Confidence", interactive=False)
        feedback_output = gr.Textbox(label="Feedback Status", interactive=False)
        
        submit_btn = gr.Button("Submit")
        feedback_btn = gr.Button("The model was wrong")
        
        submit_btn.click(predict, inputs=image_input, outputs=[prediction_output, confidence_output])
        feedback_btn.click(report_feedback, outputs=feedback_output)

    # Launch the Gradio interface
    demo.launch(share=True)

if __name__ == "__main__":
    main()