from __future__ import annotations

import os

import tensorflow as tf
from tensorflow.keras.layers import (Activation, Bidirectional, Conv3D, Dense,
                                     Dropout, Flatten, LSTM, MaxPool3D,
                                     TimeDistributed)
from tensorflow.keras.models import Sequential 


# Disable all GPUS
tf.config.set_visible_devices([], 'GPU')

def load_model() -> Sequential:
    model = Sequential()

    model.add(Conv3D(128, 3, input_shape=(75,46,140,1), padding='same'))
    model.add(Activation('relu'))
    model.add(MaxPool3D((1,2,2)))

    model.add(Conv3D(256, 3, padding='same'))
    model.add(Activation('relu'))
    model.add(MaxPool3D((1,2,2)))

    model.add(Conv3D(75, 3, padding='same'))
    model.add(Activation('relu'))
    model.add(MaxPool3D((1,2,2)))

    model.add(TimeDistributed(Flatten()))

    model.add(Bidirectional(LSTM(128, kernel_initializer='Orthogonal', return_sequences=True)))
    model.add(Dropout(.5))

    model.add(Bidirectional(LSTM(128, kernel_initializer='Orthogonal', return_sequences=True)))
    model.add(Dropout(.5))

    model.add(Dense(41, kernel_initializer='he_normal', activation='softmax'))

    base_dir = os.path.dirname(os.path.abspath(__file__))
    weights_path = os.path.abspath(os.path.join(base_dir, '..', 'models', 'checkpoint'))
    if not os.path.exists(weights_path):
        raise FileNotFoundError(f"Model weights not found at {weights_path}")
    model.load_weights(weights_path)

    return model