import tensorflow as tf
import numpy as np
import cv2
from video_compression_model import VideoCompressionModel

# Constants
NUM_CHANNELS = 3

# Step 2: Load the trained model
model = tf.keras.models.load_model('models/model.keras', custom_objects={'VideoCompressionModel': VideoCompressionModel})

# Step 3: Load the uncompressed video
UNCOMPRESSED_VIDEO_FILE = 'test_data/test_video.mkv'

def load_frames_from_video(video_file, num_frames = 0):
    print("Extracting video frames...")
    cap = cv2.VideoCapture(video_file)
    frames = []
    count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            print("Max frames from file reached")
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame)
        count += 1
        if num_frames == 0 or count >= num_frames:
            print("Max Frames wanted reached: ", num_frames)
            break
    cap.release()
    print("Extraction Complete")
    return frames

uncompressed_frames = load_frames_from_video(UNCOMPRESSED_VIDEO_FILE, 200)
if len(uncompressed_frames) == 0 or None:
    print("IO ERROR!")
    exit()
    
uncompressed_frames = np.array(uncompressed_frames) / 255.0

if len(uncompressed_frames) == 0 or None:
    print("np.array ERROR!")
    exit()

# Step 4: Compress the video frames using the loaded model
compressed_frames = model.predict(uncompressed_frames)

# Step 5: Save the compressed video frames
COMPRESSED_VIDEO_FILE = 'compressed_video.mkv'

def save_frames_as_video(frames, video_file):
    print("Saving video frames...")
    height, width = frames[0].shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(video_file, fourcc, 24.0, (width, height))
    for frame in frames:
        frame = np.clip(frame * 255.0, 0, 255).astype(np.uint8)
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        out.write(frame)
    out.release()

save_frames_as_video(compressed_frames, COMPRESSED_VIDEO_FILE)
print("Compression completed.")