import cv2
from Prediction import predict_fight            # Assuming you have this

# Process video: read, predict, and output labeled video
def process_video(video_path):
    cap = cv2.VideoCapture(video_path)
    sequence_length = 40  # Number of frames for one prediction
    all_frames = []
    predictions = []
    
    # Step 1: Read all frames from the video
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        all_frames.append(frame)
    cap.release()

    # Step 2: Process frames in chunks of 40 to make predictions
    for i in range(0, len(all_frames), sequence_length):
        frames_buffer = all_frames[i:i + sequence_length]  # Get a batch of 40 frames

        # If the number of frames is less than 40 at the end, pad it with the last frame
        if len(frames_buffer) < sequence_length:
            frames_buffer += [frames_buffer[-1]] * (sequence_length - len(frames_buffer))

        # Perform the prediction on the current batch of frames
        fight_detected = predict_fight(frames_buffer)

        # Store the prediction for this batch
        predictions.append(fight_detected)

    # Step 3: Create output video with labels
    output_video_path = "output_labeled.mp4"
    height, width, _ = all_frames[0].shape
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video_path, fourcc, 30, (width, height))  # Adjust frame rate if needed

    frame_idx = 0
    for pred in predictions:
        label = "Violence Detected!" if pred else "No Violence"
        color = (0, 0, 255) if pred else (0, 255, 0)

        # For the next 40 frames, show the same label
        for _ in range(sequence_length):
            if frame_idx >= len(all_frames):
                break  

            frame = all_frames[frame_idx]
            cv2.putText(frame, label, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
            out.write(frame)
            frame_idx += 1

    out.release()

    return output_video_path