# DeepEncode.py

import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'

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

# Constants
CHUNK_SIZE = 24  # Adjust based on available memory and video resolution
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
MAX_FRAMES = 0  # Limit the number of frames processed
CRF = 24.0  # Example CRF value
PRESET_SPEED = "veryslow"  # Index for "fast" in our defined list

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

# Load the uncompressed video
UNCOMPRESSED_VIDEO_FILE = 'test_data/training_video.mkv'

def load_frame_from_video(video_file, frame_num):
    cap = cv2.VideoCapture(video_file)
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
    ret, frame = cap.read()
    if not ret:
        return None
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0  # Normalize and convert to float32
    cap.release()
    
    return frame

def predict_frame(uncompressed_frame, model, crf_value, preset_speed_value):
    crf_array = np.array([crf_value])
    preset_speed_array = np.array([preset_speed_value])
    
    crf_array = np.expand_dims(np.array([crf_value]), axis=-1) # Shape: (1, 1)
    preset_speed_array = np.expand_dims(np.array([preset_speed_value]), axis=-1) # Shape: (1, 1)

    
    # Expand dimensions to include batch size
    uncompressed_frame = np.expand_dims(uncompressed_frame, 0)

    #display_frame = np.clip(cv2.cvtColor(uncompressed_frame[0], cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
    #cv2.imshow("uncomp", display_frame)
    #cv2.waitKey(0)
    
    compressed_frame = model.predict({
        "compressed_frame": uncompressed_frame,
        "uncompressed_frame": uncompressed_frame,
        "crf": crf_array, 
        "preset_speed": preset_speed_array
    })
    
    display_frame = np.clip(cv2.cvtColor(compressed_frame[0], cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
    
    cv2.imshow("comp", display_frame)
    cv2.waitKey(1)
    
    return compressed_frame[0]

cap = cv2.VideoCapture(UNCOMPRESSED_VIDEO_FILE)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
height, width = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
cap.release()

fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(COMPRESSED_VIDEO_FILE, fourcc, 24.0, (width, height))

if not out.isOpened():
    print("Error: VideoWriter could not be opened.")
    exit()

if MAX_FRAMES != 0 and total_frames > MAX_FRAMES:
    total_frames = MAX_FRAMES

for i in range(total_frames):
    uncompressed_frame = load_frame_from_video(UNCOMPRESSED_VIDEO_FILE, frame_num=i)
    compressed_frame = predict_frame(uncompressed_frame, model, CRF, PRESET_SPEED_CATEGORIES.index(PRESET_SPEED))
    
    compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
    compressed_frame = cv2.cvtColor(compressed_frame, cv2.COLOR_RGB2BGR)
    out.write(compressed_frame)
    #cv2.imshow("output", compressed_frame)

out.release()
print("Compression completed.")