import gradio as gr
import cv2
import numpy as np
import tensorflow as tf
import os
from tempfile import NamedTemporaryFile
from tqdm import tqdm

# Load model
model = tf.keras.models.load_model('rain_detection_model.h5')
IMG_SIZE = (256, 256)

def batch_predict(frames):
    resized = [cv2.resize(f, IMG_SIZE) for f in frames]
    batch = np.array(resized) / 255.0
    preds = model.predict(batch)
    results = []
    for p in preds:
        label = "Hujan" if p[0] > 0.7 else "Terang"
        confidence = float(p[0]) if label == "Hujan" else 1 - float(p[0])
        results.append((label, round(confidence, 2)))
    return results

def process_video(video_file):
    cap = cv2.VideoCapture(video_file)
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_interval = int(fps * 1)

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    extra_space = 50

    out_path = NamedTemporaryFile(delete=False, suffix=".mp4").name
    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height + extra_space))

    frames = []
    original_frames = []
    frame_idxs = []
    i = 0

    with tqdm(total=total_frames) as pbar:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            if i % frame_interval == 0:
                frames.append(frame.copy())
                frame_idxs.append(i)
            original_frames.append(frame)
            i += 1
            pbar.update(1)

    print("[INFO] Melakukan prediksi cuaca setiap 1 detik...")

    label_conf_list = batch_predict(frames)
    label_dict = dict(zip(frame_idxs, label_conf_list))

    current_label = "Tidak diketahui"
    current_conf = 0.0
    current_color = (255, 255, 255)

    for i, frame in enumerate(original_frames):
        if i in label_dict:
            current_label, current_conf = label_dict[i]
            current_color = (0, 255, 0) if current_label == "Terang" else (0, 0, 255)

        # Tambah canvas putih di atas frame
        label_area = np.ones((extra_space, width, 3), dtype=np.uint8) * 255
        combined_frame = np.vstack((label_area, frame))

        text = f"{current_label} ({current_conf:.2f})"
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 1
        thickness = 2

        (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
        text_x = (width - text_width) // 2
        text_y = (extra_space + text_height) // 2

        # Tambah teks di canvas (tengah atas)
        cv2.putText(combined_frame, text, (text_x, text_y), font, font_scale, current_color, thickness)


        out.write(combined_frame)

    cap.release()
    out.release()

    print("[INFO] Proses selesai. Video disimpan di:", out_path)
    return out_path, out_path

# Gradio Interface
interface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Upload Video"),
    outputs=[
        gr.Video(label="Hasil Prediksi (Putar Video)"),
        gr.File(label="Unduh Video Hasil Prediksi")
    ],
    title="Prediksi Cuaca",
    description="Upload video. Sistem akan memprediksi 'Hujan' atau 'Terang'."
)

interface.launch(debug=True)