"""
Example: Encoding an image using a pre-trained convolutional autoencoder.

This script:
1. Loads a 256x256 grayscale image ('frame.png').
2. Loads the previously trained autoencoder model..
3. Encodes the image and prints the feature vector.
"""

import tensorflow as tf
import numpy as np
from PIL import Image

# ======================================================
# Configuration
# ======================================================
MODEL_PATH = 'frame_encoder_session1_128_arch7b'   # Path to the saved model
IMAGE_PATH = 'frame.png'                            # Input image path
ENCODE_LAYER_NAME = 'encoded_vector'                 # Name of encoding layer from model


# ======================================================
# Load image and preprocess
# ======================================================
def load_grayscale_image(image_path, target_size=(256, 256)):
    """
    Loads a grayscale image, resizes it, and normalizes pixel values.

    Args:
        image_path (str): Path to the image file.
        target_size (tuple): Target (width, height) of the image.

    Returns:
        np.ndarray: Preprocessed image of shape (1, 256, 256, 1).
    """
    img = Image.open(image_path).convert('L')  # Convert to grayscale
    img = img.resize(target_size)
    img_array = np.array(img, dtype=np.float32) / 255.0  # Normalize to [0, 1]
    img_array = img_array.reshape(1, target_size[0], target_size[1], 1)  # Add batch and channel dims
    return img_array


# ======================================================
# Define function to encode image using autoencoder
# ======================================================
def encode_image(autoencoder, image, layer_name=ENCODE_LAYER_NAME):
    """
    Encodes a preprocessed image using the bottleneck layer of the autoencoder.

    Args:
        autoencoder (tf.keras.Model): Trained autoencoder model.
        image (np.ndarray): Preprocessed image (1, 256, 256, 1).
        layer_name (str): Name of the encoding layer in the model.

    Returns:
        np.ndarray: Encoded feature vector.
    """
    # Extract the encoding layer output
    encoder = tf.keras.Model(inputs=autoencoder.input,
                             outputs=autoencoder.get_layer(layer_name).output)
    encoded_vector = encoder.predict(image)
    return encoded_vector


# ======================================================
# Main execution
# ======================================================
if __name__ == "__main__":
    # Load the trained autoencoder
    print("Loading trained autoencoder...")
    autoencoder = tf.keras.models.load_model(MODEL_PATH)

    # Load and preprocess image
    print(f"Loading image from {IMAGE_PATH}...")
    input_image = load_grayscale_image(IMAGE_PATH)

    # Encode the image
    print("Encoding image...")
    encoded_output = encode_image(autoencoder, input_image)

    # Display results
    print("\nEncoded feature vector values:")
    print(encoded_output.flatten())  # Print all elements as a 1D array