# featureExtraction.py

import cv2
import numpy as np

from globalVars import HEIGHT, NUM_PRESET_SPEEDS, WIDTH


def extract_edge_features(frame):
    """
    Extract edge features using Canny edge detection.

    Args:
    - frame (ndarray): Image frame.

    Returns:
    - ndarray: Edge feature map.
    """
    edges = cv2.Canny(frame, threshold1=100, threshold2=200)
    return edges.astype(np.float32) / 255.0

def extract_histogram_features(frame, bins=64):
    """
    Extract histogram features from a frame with 3 channels.

    Args:
    - frame (ndarray): Image frame with shape (height, width, 3).
    - bins (int): Number of bins for the histogram.

    Returns:
    - ndarray: Normalized histogram feature vector.
    """
    feature_vector = []
    for channel in range(3):
        histogram, _ = np.histogram(frame[:,:,channel].flatten(), bins=bins, range=[0, 255])
        normalized_histogram = histogram.astype(np.float32) / frame[:,:,channel].size
        feature_vector.extend(normalized_histogram)

    return np.array(feature_vector)


def preprocess_frame(frame, crf, speed):
    # Check frame dimensions and resize if necessary
    if frame.shape[:2] != (HEIGHT, WIDTH):
        frame = cv2.resize(frame, (WIDTH, HEIGHT), interpolation=cv2.INTER_NEAREST)
    
    # Scale frame to [0, 1]
    compressed_frame = frame / 255.0

    # Scale CRF and SPEED to [0, 1] (assuming they are within known bounds)
    crf_scaled = crf / 51
    speed_scaled = speed / NUM_PRESET_SPEEDS

    # Create images with the CRF and SPEED values, filling extra channels
    crf_image = np.full((HEIGHT, WIDTH, 1), crf_scaled) # Note the added dimension
    speed_image = np.full((HEIGHT, WIDTH, 1), speed_scaled) # Note the added dimension

    # Combine the frames with the CRF and SPEED images
    combined_frame = np.concatenate([compressed_frame, crf_image, speed_image], axis=-1)

    return combined_frame