import gradio as gr
import random
from PIL import Image
import time
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification

# הגדרת קטגוריות
moves = ["rock", "paper", "scissors"]
computer_images = {
    "rock": "computer_rock.png",
    "paper": "computer_paper.png",
    "scissors": "computer_scissors.png"
}

# טעינת המודל
checkpoint = "facebook/deit-tiny-patch16-224"
processor = AutoImageProcessor.from_pretrained(checkpoint)
model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=3,
    id2label={0: "rock", 1: "paper", 2: "scissors"},
    label2id={"rock": 0, "paper": 1, "scissors": 2},
    ignore_mismatched_sizes=True
)

# חוקי המשחק
def game_logic(user_move, computer_move):
    if user_move == computer_move:
        return "It's a tie!"
    if (user_move == "rock" and computer_move == "scissors") or \
       (user_move == "paper" and computer_move == "rock") or \
       (user_move == "scissors" and computer_move == "paper"):
        return "You win!"
    else:
        return "You lose!"

# ספירה לאחור + משחק
def full_play(live_image):
    # ספירה לאחור בשרת (יופיע בלוגים)
    for i in ["3...", "2...", "1...", "GO!"]:
        print(i)
        time.sleep(1)

    # ניתוח תמונה
    prediction = processor(images=live_image, return_tensors="pt")
    outputs = model(**prediction)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    user_move = model.config.id2label[predicted_class_idx]

    computer_move = random.choice(moves)
    computer_img = Image.open(computer_images[computer_move])

    result = game_logic(user_move, computer_move)

    return live_image, computer_img, f"You chose {user_move}, computer chose {computer_move}. {result}"

# בניית האפליקציה
with gr.Blocks() as demo:
    gr.Markdown("# ✂️ 🪨 📄 Rock Paper Scissors - LIVE Game!")

    webcam_input = gr.Image(source="webcam", tool=None, label="Show your move!")
    play_button = gr.Button("Start Countdown and Play")

    user_output = gr.Image(label="Your Move")
    computer_output = gr.Image(label="Computer's Move")
    result_text = gr.Textbox(label="Result")

    play_button.click(
        fn=full_play,
        inputs=[webcam_input],
        outputs=[user_output, computer_output, result_text]
    )

demo.launch()