# featureExtraction.py

import cv2
import numpy as np
import os

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

import tensorflow as tf
from tensorflow.keras import backend as K

from globalVars import HEIGHT, LOGGER, NUM_PRESET_SPEEDS, WIDTH

def scale_crf(crf):
    return crf / 51


def scale_speed_preset(speed):
    return speed / NUM_PRESET_SPEEDS


def psnr(y_true, y_pred):
    #LOGGER.info(f"[psnr function] y_true: {y_true.shape}, y_pred: {y_pred.shape}")
    max_pixel = 1.0
    mse = K.mean(K.square(y_pred - y_true))
    mse = tf.cast(mse, tf.float32)  # Cast mse to tf.float32
    return 20.0 * K.log(max_pixel / K.sqrt(mse)) / K.log(10.0)


def ssim(y_true, y_pred):
    return (tf.image.ssim(y_true, y_pred, max_val=1.0) + 1) * 50  # Normalize SSIM from [-1, 1] to [0, 100]


def combined(y_true, y_pred):
    return (psnr(y_true, y_pred) + ssim(y_true, y_pred)) / 2


def combined_loss(y_true, y_pred):
    return -combined(y_true, y_pred)  # The goal is to maximize the combined value

# Option 1: Weight more towards PSNR
def combined_loss_weighted_psnr(y_true, y_pred):
    return -0.7 * psnr(y_true, y_pred) - 0.3 * ssim(y_true, y_pred)

# Option 2: Weight more towards SSIM
def combined_loss_weighted_ssim(y_true, y_pred):
    return -0.3 * psnr(y_true, y_pred) - 0.7 * ssim(y_true, y_pred)


def detect_noise(image, threshold=15):
    # Convert to grayscale if it's a color image
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
    # Compute the standard deviation
    std_dev = np.std(image)
    
    # If the standard deviation is higher than a threshold, it might be considered noisy
    return std_dev > threshold


def frame_difference(frame1, frame2):
    # Ensure both frames are of the same size and type
    if frame1.shape != frame2.shape:
        raise ValueError("Frames must have the same dimensions and number of channels")

    # Calculate the absolute difference between the frames
    difference = cv2.absdiff(frame1, frame2)

    return difference


def preprocess_frame(frame, resize=True, scale=True):
    
    # Check frame dimensions and resize if necessary
    if resize and frame.shape[:2] != (HEIGHT, WIDTH):
        frame = cv2.resize(frame, (WIDTH, HEIGHT), interpolation=cv2.INTER_LINEAR)
        
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    
    if scale:
    # Scale frame to [0, 1]
        frame = frame / 255.0

    return frame