# inference.py

import tensorflow as tf
import logging
from preprocessing import VideoPreprocessor

# Configure Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class LipReadingModel:
    def __init__(self, model_path='best_model_1_WER.keras'):
        # Initialize character mappings before loading the model
        vocab = [x for x in "aăâbcdđeêghiklmnoôơpqrstuưvxyáàảãạấầẩẫậắằẳẵặéèẻẽẹếềểễệíìỉĩịóòỏõọốồổỗộớờởỡợúùủũụứừửữựýỳỷỹỵ'?!123456789 "]
        self.char_to_num = tf.keras.layers.StringLookup(vocabulary=vocab, oov_token="")
        self.num_to_char = tf.keras.layers.StringLookup(
            vocabulary=self.char_to_num.get_vocabulary(), oov_token="", invert=True
        )

        try:
            self.model = tf.keras.models.load_model(
                model_path,
                custom_objects={'CTCLoss': self.CTCLoss}  # Include custom loss if needed
            )
            logger.info("Model loaded successfully.")
        except Exception as e:
            logger.error(f"Error loading model: {e}")
            self.model = self.build_model()  # Fallback to building the model if loading fails

    @staticmethod
    def CTCLoss(y_true, y_pred):
        batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
        input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")
        label_length = tf.cast(tf.shape(y_true)[1], dtype="int64")
        input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")
        label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64")
        return tf.keras.backend.ctc_batch_cost(y_true, y_pred, input_length, label_length)

    def build_model(self):
        model = tf.keras.Sequential()

        # First convolutional layer with BatchNormalization
        model.add(tf.keras.layers.Conv3D(64, (3, 3, 3), strides=(1, 2, 2), input_shape=(None, 85, 85, 1), padding='same'))
        model.add(tf.keras.layers.BatchNormalization())
        model.add(tf.keras.layers.Activation('relu'))
        model.add(tf.keras.layers.MaxPool3D((1, 2, 2), padding='same'))

        # Second convolutional layer
        model.add(tf.keras.layers.Conv3D(128, (3, 3, 3), strides=(1, 2, 2), padding='same'))
        model.add(tf.keras.layers.BatchNormalization())
        model.add(tf.keras.layers.Activation('relu'))
        model.add(tf.keras.layers.MaxPool3D((1, 2, 2), padding='same'))

        # Third convolutional layer
        model.add(tf.keras.layers.Conv3D(256, (3, 3, 3), strides=(1, 2, 2), padding='same'))
        model.add(tf.keras.layers.LayerNormalization())
        model.add(tf.keras.layers.Activation('relu'))
        model.add(tf.keras.layers.MaxPool3D((1, 2, 2), padding='same'))

        # Fourth convolutional layer
        model.add(tf.keras.layers.Conv3D(256, (3, 3, 3), padding='same'))
        model.add(tf.keras.layers.BatchNormalization())
        model.add(tf.keras.layers.Activation('relu'))
        model.add(tf.keras.layers.MaxPool3D((1, 2, 2), padding='same'))

        # Flatten and pass through TimeDistributed
        model.add(tf.keras.layers.TimeDistributed(tf.keras.layers.Flatten()))

        # Bidirectional LSTM layers
        model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(512, kernel_initializer='Orthogonal', return_sequences=True)))
        model.add(tf.keras.layers.Dropout(0.4))

        model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(256, kernel_initializer='Orthogonal', return_sequences=True)))
        model.add(tf.keras.layers.Dropout(0.4))

        # Final dense layer
        model.add(tf.keras.layers.Dense(self.char_to_num.vocabulary_size() + 1, kernel_initializer='he_normal', activation='softmax'))

        logger.info("Built the model architecture successfully.")
        return model

    def predict(self, normalized_frames):
        if self.model is None:
            return "❌ Model not loaded. Please check the model path and ensure the model file is accessible."

        if normalized_frames is None or int(tf.size(normalized_frames)) == 0:
            return "❌ No frames extracted from the video. Please ensure the video contains a clear view of the face and lips."

        try:
            # Add batch dimension
            frames = tf.expand_dims(normalized_frames, axis=0)  # Shape: (1, num_frames, 85, 85, 1)

            # Perform prediction
            yhat = self.model.predict(frames, verbose=0)

            input_length = [yhat.shape[1]]  # batch size of 1
            # Perform CTC decoding
            decoded_tf = tf.keras.backend.ctc_decode(yhat, input_length=input_length, greedy=True)[0][0]
            decoded = decoded_tf.numpy().flatten()  # Convert to numpy array and flatten

            logger.debug(f"Decoded prediction: {decoded}")

            # Convert numerical predictions to characters
            prediction = ''.join(
                [
                    self.num_to_char(int(num)).numpy().decode('utf-8')
                    for num in decoded
                    if int(num) != -1
                ]
            )

            return prediction
        except Exception as e:
            logger.error(f"Error during prediction: {e}")
            return f"❌ An error occurred during prediction: {e}"

def predict_from_video(video_path=None, frames=None, model=None, preprocessor=None):
    """
    Predicts the text from a video file or webcam frames using the provided model.
    Args:
        video_path (str, optional): Path to the video file. Defaults to None.
        frames (List[np.ndarray], optional): List of frames from webcam. Defaults to None.
        model (LipReadingModel, optional): An instance of the LipReadingModel. Defaults to None.
        preprocessor (VideoPreprocessor, optional): An instance of the VideoPreprocessor. Defaults to None.
    Returns:
        str: Predicted text.
    """
    if model is None:
        model = LipReadingModel()

    if preprocessor is None:
        preprocessor = VideoPreprocessor()

    if video_path:
        # Preprocess video from file
        normalized_frames = preprocessor.preprocess_video(video_path)
    elif frames is not None:
        # Preprocess frames from webcam
        normalized_frames = preprocessor.preprocess_frames(frames)
    else:
        return "❌ No video or frames provided for prediction."

    prediction = model.predict(normalized_frames)
    return prediction