# DeepEncode.py

import os

from featureExtraction import preprocess_frame, psnr, scale_crf, scale_speed_preset
from globalVars import PRESET_SPEED_CATEGORIES

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

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

# Constants
COMPRESSED_VIDEO_FILE = 'compressed_video.avi'
MAX_FRAMES = 0  # Limit the number of frames processed
CRF = 0
SPEED = PRESET_SPEED_CATEGORIES.index("ultrafast")

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

# Load the uncompressed video
UNCOMPRESSED_VIDEO_FILE = 'test_data/B4_t02.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
    cap.release()
    
    return frame

def predict_frame(uncompressed_frame):

    #display_frame = np.clip(cv2.cvtColor(uncompressed_frame, cv2.COLOR_BGR2RGB) * 255.0, 0, 255).astype(np.uint8)
    #cv2.imshow("uncomp", uncompressed_frame)
    scaled_crf = scale_crf(CRF)
    scaled_speed = scale_speed_preset(SPEED)
    
    frame = combine_batch(uncompressed_frame, scaled_crf, scaled_speed, resize=False)
    
    compressed_frame = MODEL.predict([np.expand_dims(frame, axis=0)])[0]
    
    compressed_frame = compressed_frame[:, :, :3]  # Keep only the first 3 channels (BGR)
    
    compressed_frame = np.clip(compressed_frame * 255.0, 0, 255).astype(np.uint8)
    
    cv2.imshow("comp", compressed_frame)
    cv2.waitKey(1)
    
    return compressed_frame

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), True)

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)
    
    compressed_frame = cv2.resize(compressed_frame, (width, height))
    
    #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)

    #if i % 10 == 0:  # Print progress every 10 frames
    #    print(f"Processed {i} / {total_frames} frames")

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