import cv2
import numpy as np
import torch
from torch import nn
from torchvision import transforms, models
import gradio as gr

class BISINDOClassifier(nn.Module):
    def __init__(self, hidden_dim, num_classes, num_layers=1):
        super(BISINDOClassifier, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_classes = num_classes
        self.num_layers = num_layers

        # Load pre-trained VGG-19
        self.vgg19 = models.vgg19(pretrained=True)
        self.vgg19.classifier = nn.Sequential()

        # GRU
        self.gru = nn.GRU(input_size=512*7*7, hidden_size=hidden_dim, num_layers=num_layers, batch_first=True)

        self.fc = nn.Linear(hidden_dim, num_classes)

    def forward(self, x):
        batch_size, seq_length, c, h, w = x.size()
        x = x.view(batch_size * seq_length, c, h, w)

        with torch.no_grad():
            x = self.vgg19.features(x)

        x = x.view(batch_size, seq_length, -1)

        h0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(x.device)
        x, _ = self.gru(x, h0)

        x = self.fc(x[:, -1, :])

        return x

transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = BISINDOClassifier(hidden_dim=512, num_classes=40, num_layers=1).to(device)
model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))

classes = {'Bagaimana', 'Berapa', 'Baik', 'Selamat Malam', 'Halo', 'Pendek', 'Tidur', 'Selamat Sore', 'Membaca', 'Senang', 'Kemana', 'Dia', 'Apa Kabar', 'Saya', 'Apa', 'Kita', 'Sabar', 'Selamat Siang', 'Kalian', 'Dimana', 'Duduk', 'Kapan', 'Mereka', 'Kamu', 'Ramah', 'Makan', 'Tinggi', 'Marah', 'Berdiri', 'Melihat', 'Minum', 'Siapa', 'Selamat Pagi', 'Kami', 'Mandi', 'Menulis', 'Terima Kasih', 'Sedih', 'Bingung', 'Belajar'}
classes = sorted(list(classes))

def extract_frames(video_path, sequence_length, interval=10, img_size=(224, 224)):
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        print(f"Error opening video file {video_path}")
        return []

    frames = []
    frame_count = 0

    while len(frames) < sequence_length:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_count % interval == 0:
            frame = cv2.resize(frame, img_size)
            frames.append(frame)
        frame_count += 1

    cap.release()

    if len(frames) < sequence_length:
        while len(frames) < sequence_length:
            if frames:
                frames.append(frames[-1])
            else:
                frames.append(cv2.resize(np.zeros((img_size[1], img_size[0], 3), dtype=np.uint8), img_size))
    else:
        frames = frames[:sequence_length]

    return frames

def preprocess_frames(frames):
    frames = [transform(frame) for frame in frames]
    frames = torch.stack(frames)
    return frames.unsqueeze(0)

def predict(video_path, sequence_length):
    model.eval()

    frames = extract_frames(video_path, sequence_length)
    input_tensor = preprocess_frames(frames)
    input_tensor = input_tensor.to(device)

    with torch.no_grad():
        output = model(input_tensor)
        probabilities = nn.functional.softmax(output, dim=1)
        confidence, predicted = torch.max(probabilities, 1)

    confidence_score = confidence.item()
    predicted_label = classes[predicted.item()]

    return {"label": predicted_label, "confidence": confidence_score}

with gr.Blocks() as demo:
    with gr.Row():
        input_video = gr.Video(label="Input")
        output_json = gr.JSON(label="Output")
        process_video_btn = gr.Button("Process Video")

    def process_video(video):
        result = predict(video, sequence_length=10)
        result["confidence"] = f"{result['confidence']:.2f}"
        return result

    process_video_btn.click(process_video, input_video, output_json)

demo.launch()