import os
from time import sleep
from typing import Any, Generator

import numpy as np
import tensorflow as tf
import tensorflow.lite as tflite
from tensorflow import keras
from tensorflow.keras import layers

assert float(tf.__version__[:3]) >= 2.3
keras.backend.clear_session()


def kaiming(shape, dtype, partition_info=None):
    """Kaiming initialization as described in https://arxiv.org/pdf/1502.01852.pdf"""
    return tf.random.truncated_normal(shape) * tf.sqrt(2 / float(shape[0]))


def kaiming_and_clip(shape, dtype, partition_info=None):
    """Kaiming initialization but cliping by norm 1"""
    return tf.clip_by_norm(kaiming(shape, dtype, partition_info), 1)


class ModelHandler:
    """Handler for training proposed model in https://arxiv.org/pdf/1502.01852.pdf"""

    MODEL_VERSION = '1'
    CUSTOM_OBJS = {
        'kaiming': kaiming,
        'kaiming_and_clip': kaiming_and_clip
    }

    def __init__(self, num_2d_kpts: int = 16,
                 num_3d_kpts: int = 16,
                 save_dir: str = 'training',
                 load: bool = True,
                 linear_size: int = 1024,
                 num_bi_layers: int = 2,
                 dropout_keep_prob: float = 1.0,
                 clip_weights: bool = False,
                 add_batch_norm: bool = False,
                 add_residual: bool = False) -> None:
        """Create a handler for proposed model in https://arxiv.org/pdf/1502.01852.pdf.

        Args:
            num_2d_kpts (int, optional): Number of input 2D keypoints. Defaults to 16.
            num_3d_kpts (int, optional): Number of input 3D keypoints. Defaults to 16.
            save_dir (str, optional): Directory where store trained models and checkpoints. Defaults to 'training'.
            load (bool, optional): Whether to load model from checkpoint file before training. Defaults to True.
            linear_size (int, optional): Number of units in each layer of the model. Defaults to 1024.
            num_bi_layers (int, optional): Number of bilinear blocks in the model. Defaults to 2.
            dropout_keep_prob (float, optional): Dropout keep probability. Value '1.0' means no dropout. Defaults to 1.0.
            clip_weights (bool, optional): Whether to clip weights to a norm of 1. Defaults to False.
            add_batch_norm (bool, optional): Whether to use batch normalization. Defaults to False.
            add_residual (bool, optional): Whether to add residual connections between bilinear blocks. Defaults to False.

        Raises:
            ValueError: if save_dir does not exist or it is not a directory.
        """
        if not os.path.exists(save_dir):
            raise ValueError(f"Path '{save_dir}' does not exist.")
        elif not os.path.isdir(save_dir):
            raise ValueError(f"Path '{save_dir}' is not a directory.")

        self.input_size = num_2d_kpts * 2
        self.output_size = num_3d_kpts * 3
        self.save_dir = os.path.join(
            save_dir, '{}_{}bi_{}drop_{}clip_{}norm_{}residual'.format(
                linear_size, num_bi_layers, dropout_keep_prob,
                'no_' if not clip_weights else '',
                'no_' if not add_batch_norm else '',
                'no_' if not add_residual else '',
            ))

        if not os.path.exists(self.save_dir):
            print(f'Creating {self.save_dir} directory')
            os.system(f'mkdir -p {self.save_dir}')

        if load:
            if os.path.exists(self.checkpoint_path):
                self.model = keras.models.load_model(
                    self.checkpoint_path,
                    custom_objects=self.CUSTOM_OBJS)
                return
            else:
                print("Checkpoint file '{}' does not exist.".format(
                    self.checkpoint_path))
        self.model = self._create_model(
            linear_size=linear_size,
            num_bi_layers=num_bi_layers,
            dropout_keep_prob=dropout_keep_prob,
            clip_weights=clip_weights,
            add_batch_norm=add_batch_norm,
            add_residual=add_residual)

    @property
    def checkpoint_path(self) -> str:
        return os.path.join(self.save_dir, 'checkpoint.h5')

    @property
    def model_path(self) -> str:
        return os.path.join(self.save_dir, 'model')  # SavedModel format

    @property
    def tflite_model_path(self) -> str:
        return self.model_path + '.tflite'

    @property
    def tflite_quant_model_path(self) -> str:
        return os.path.join(self.save_dir, 'quant_model') + '.tflite'

    def _create_model(
            self, linear_size: int = 1024,
            num_bi_layers: int = 2,
            dropout_keep_prob: float = 1.0,
            clip_weights: bool = False,
            add_batch_norm: bool = False,
            add_residual: bool = False) -> keras.Model:
        inputs = keras.Input(shape=(self.input_size,), dtype=tf.float32)

        x = layers.Dense(
            units=linear_size, activation=None,
            kernel_initializer=kaiming_and_clip if clip_weights else kaiming,
            use_bias=True, bias_initializer=kaiming)(inputs)
        if add_batch_norm:
            x = layers.BatchNormalization()(x, training=True)
        x = layers.ReLU()(x)
        x = layers.Dropout(1 - dropout_keep_prob)(x)

        for i in range(num_bi_layers):
            if add_residual:
                residual = x
            for j in range(2):
                x = layers.Dense(
                    units=linear_size, activation=None,
                    kernel_initializer=kaiming_and_clip if clip_weights else kaiming,
                    use_bias=True, bias_initializer=kaiming)(x)
                if add_batch_norm:
                    x = layers.BatchNormalization()(x, training=True)
                x = layers.ReLU()(x)
                x = layers.Dropout(1 - dropout_keep_prob)(x)
            if add_residual:
                x = layers.Add()([residual, x])

        outputs = layers.Dense(
            units=self.output_size, activation=None,
            kernel_initializer=kaiming_and_clip if clip_weights else kaiming,
            use_bias=True, bias_initializer=kaiming)(x)

        return keras.Model(inputs=inputs, outputs=outputs)

    def _update_learning_rate(self, epoch: float, lr: float) -> float:
        # From tf official docs -> NOT from original repo
        #   As we cannot change lr each step with standard callbacks (only each epoch)
        # if epoch < 10:
        #     return lr
        # else:
        return lr * tf.math.exp(-0.1)

    def train(self, x_train: np.ndarray, y_train: np.ndarray,
              epochs: int = 1, learning_rate: float = 1e-3) -> None:
        self.model.compile(
            optimizer=keras.optimizers.Adam(learning_rate),
            loss=keras.losses.MeanSquaredError(),
            metrics=[
                keras.metrics.Accuracy(),
                keras.metrics.FalsePositives(),
                keras.metrics.FalseNegatives(),
                keras.metrics.TruePositives(),
                keras.metrics.TrueNegatives()
            ])

        batch_size = x_train.shape[0]
        for i in range(2, int(batch_size / 2), -1):
            if batch_size % i == 0:
                batch_size = i
                break

        num_val = int(x_train.shape[0] * 0.1)
        x_val = x_train[-num_val:]
        y_val = y_train[-num_val:]
        x_train = x_train[:-num_val]
        y_train = y_train[:-num_val]

        self.model.fit(
            x_train, y_train,
            epochs=epochs, batch_size=batch_size,
            callbacks=[
                keras.callbacks.ModelCheckpoint(
                    self.checkpoint_path, save_best_only=True,
                    monitor='accuracy', verbose=1),
                keras.callbacks.LearningRateScheduler(
                    self._update_learning_rate, verbose=1)
            ],
            validation_data=(x_val, y_val))

    def test(self, x_test: np.ndarray, y_test: np.ndarray) -> list:
        batch_size = x_test.shape[0]
        for i in range(2, int(batch_size / 2), -1):
            if batch_size % i == 0:
                batch_size = i
                break

        return self.model.evaluate(
            x_test, y_test,
            batch_size=batch_size,
            verbose=1)

    def save(self) -> None:
        """Save model."""
        tf.saved_model.save(self.model, self.model_path)

    def _generate_representative(self) -> Generator:
        # FIXME Move to data_utils module
        # FIXME Use real dataset -> 'Random. From training or evaluation set' --> Length of 100 for example
        for i in list(np.arange(100 * 32, dtype='float32').reshape((100, 32))):
            yield [i]

    def tflite_convert(self, integer_quant: bool = False) -> None:
        if not os.path.exists(self.model_path):
            raise RuntimeError(
                f"Saved model path '{self.model_path}' does not exist.")

        converter = tflite.TFLiteConverter.from_saved_model(
            self.model_path)
        converter.optimizations = [tflite.Optimize.DEFAULT]
        converter.experimental_enable_resource_variables = True

        if integer_quant:
            converter.inference_input_type = tf.uint8
            converter.inference_output_type = tf.uint8
            converter.representative_dataset = self._generate_representative
            converter.target_spec.supported_ops = [
                tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

            with open(self.tflite_quant_model_path, 'wb') as f:
                f.write(converter.convert())
        else:
            with open(self.tflite_model_path, 'wb') as f:
                f.write(converter.convert())


# handler = ModelHandler(dropout_keep_prob=0.5, clip_weights=True,
#                        add_batch_norm=True, add_residual=True,
#                        load=True)
# print(handler.model.summary())
# handler.train(epochs=4)
# handler.save()
# handler.tflite_convert(True)

import os
from time import sleep
from typing import Any, Generator

import numpy as np
import tensorflow as tf
import tensorflow.lite as tflite
from tensorflow import keras
from tensorflow.keras import layers

assert float(tf.__version__[:3]) >= 2.3
keras.backend.clear_session()


def kaiming(shape, dtype, partition_info=None):
    """Kaiming initialization as described in https://arxiv.org/pdf/1502.01852.pdf"""
    return tf.random.truncated_normal(shape) * tf.sqrt(2 / float(shape[0]))


def kaiming_and_clip(shape, dtype, partition_info=None):
    """Kaiming initialization but cliping by norm 1"""
    return tf.clip_by_norm(kaiming(shape, dtype, partition_info), 1)


class ModelHandler:
    """Handler for training proposed model in https://arxiv.org/pdf/1502.01852.pdf"""

    MODEL_VERSION = '1'
    CUSTOM_OBJS = {
        'kaiming': kaiming,
        'kaiming_and_clip': kaiming_and_clip
    }

    def __init__(self, num_2d_kpts: int = 16,
                 num_3d_kpts: int = 16,
                 save_dir: str = 'training',
                 load: bool = True,
                 linear_size: int = 1024,
                 num_bi_layers: int = 2,
                 dropout_keep_prob: float = 1.0,
                 clip_weights: bool = False,
                 add_batch_norm: bool = False,
                 add_residual: bool = False) -> None:
        """Create a handler for proposed model in https://arxiv.org/pdf/1502.01852.pdf.

        Args:
            num_2d_kpts (int, optional): Number of input 2D keypoints. Defaults to 16.
            num_3d_kpts (int, optional): Number of input 3D keypoints. Defaults to 16.
            save_dir (str, optional): Directory where store trained models and checkpoints. Defaults to 'training'.
            load (bool, optional): Whether to load model from checkpoint file before training. Defaults to True.
            linear_size (int, optional): Number of units in each layer of the model. Defaults to 1024.
            num_bi_layers (int, optional): Number of bilinear blocks in the model. Defaults to 2.
            dropout_keep_prob (float, optional): Dropout keep probability. Value '1.0' means no dropout. Defaults to 1.0.
            clip_weights (bool, optional): Whether to clip weights to a norm of 1. Defaults to False.
            add_batch_norm (bool, optional): Whether to use batch normalization. Defaults to False.
            add_residual (bool, optional): Whether to add residual connections between bilinear blocks. Defaults to False.

        Raises:
            ValueError: if save_dir does not exist or it is not a directory.
        """
        if not os.path.exists(save_dir):
            raise ValueError(f"Path '{save_dir}' does not exist.")
        elif not os.path.isdir(save_dir):
            raise ValueError(f"Path '{save_dir}' is not a directory.")

        self.input_size = num_2d_kpts * 2
        self.output_size = num_3d_kpts * 3
        self.save_dir = os.path.join(
            save_dir, '{}_{}bi_{}drop_{}clip_{}norm_{}residual'.format(
                linear_size, num_bi_layers, dropout_keep_prob,
                'no_' if not clip_weights else '',
                'no_' if not add_batch_norm else '',
                'no_' if not add_residual else '',
            ))

        if not os.path.exists(self.save_dir):
            print(f'Creating {self.save_dir} directory')
            os.system(f'mkdir -p {self.save_dir}')

        if load:
            if os.path.exists(self.checkpoint_path):
                self.model = keras.models.load_model(
                    self.checkpoint_path,
                    custom_objects=self.CUSTOM_OBJS)
                return
            else:
                print("Checkpoint file '{}' does not exist.".format(
                    self.checkpoint_path))
        self.model = self._create_model(
            linear_size=linear_size,
            num_bi_layers=num_bi_layers,
            dropout_keep_prob=dropout_keep_prob,
            clip_weights=clip_weights,
            add_batch_norm=add_batch_norm,
            add_residual=add_residual)

    @property
    def checkpoint_path(self) -> str:
        return os.path.join(self.save_dir, 'checkpoint.h5')

    @property
    def model_path(self) -> str:
        return os.path.join(self.save_dir, 'model')  # SavedModel format

    @property
    def tflite_model_path(self) -> str:
        return self.model_path + '.tflite'

    @property
    def tflite_quant_model_path(self) -> str:
        return os.path.join(self.save_dir, 'quant_model') + '.tflite'

    def _create_model(
            self, linear_size: int = 1024,
            num_bi_layers: int = 2,
            dropout_keep_prob: float = 1.0,
            clip_weights: bool = False,
            add_batch_norm: bool = False,
            add_residual: bool = False) -> keras.Model:
        inputs = keras.Input(shape=(self.input_size,), dtype=tf.float32)

        x = layers.Dense(
            units=linear_size, activation=None,
            kernel_initializer=kaiming_and_clip if clip_weights else kaiming,
            use_bias=True, bias_initializer=kaiming)(inputs)
        if add_batch_norm:
            x = layers.BatchNormalization()(x, training=True)
        x = layers.ReLU()(x)
        x = layers.Dropout(1 - dropout_keep_prob)(x)

        for i in range(num_bi_layers):
            if add_residual:
                residual = x
            for j in range(2):
                x = layers.Dense(
                    units=linear_size, activation=None,
                    kernel_initializer=kaiming_and_clip if clip_weights else kaiming,
                    use_bias=True, bias_initializer=kaiming)(x)
                if add_batch_norm:
                    x = layers.BatchNormalization()(x, training=True)
                x = layers.ReLU()(x)
                x = layers.Dropout(1 - dropout_keep_prob)(x)
            if add_residual:
                x = layers.Add()([residual, x])

        outputs = layers.Dense(
            units=self.output_size, activation=None,
            kernel_initializer=kaiming_and_clip if clip_weights else kaiming,
            use_bias=True, bias_initializer=kaiming)(x)

        return keras.Model(inputs=inputs, outputs=outputs)

    def _update_learning_rate(self, epoch: float, lr: float) -> float:
        # From tf official docs -> NOT from original repo
        #   As we cannot change lr each step with standard callbacks (only each epoch)
        # if epoch < 10:
        #     return lr
        # else:
        return lr * tf.math.exp(-0.1)

    def train(self, x_train: np.ndarray, y_train: np.ndarray,
              epochs: int = 1, learning_rate: float = 1e-3) -> None:
        self.model.compile(
            optimizer=keras.optimizers.Adam(learning_rate),
            loss=keras.losses.MeanSquaredError(),
            metrics=[
                keras.metrics.Accuracy(),
                keras.metrics.FalsePositives(),
                keras.metrics.FalseNegatives(),
                keras.metrics.TruePositives(),
                keras.metrics.TrueNegatives()
            ])

        batch_size = x_train.shape[0]
        for i in range(2, int(batch_size / 2), -1):
            if batch_size % i == 0:
                batch_size = i
                break

        num_val = int(x_train.shape[0] * 0.1)
        x_val = x_train[-num_val:]
        y_val = y_train[-num_val:]
        x_train = x_train[:-num_val]
        y_train = y_train[:-num_val]

        self.model.fit(
            x_train, y_train,
            epochs=epochs, batch_size=batch_size,
            callbacks=[
                keras.callbacks.ModelCheckpoint(
                    self.checkpoint_path, save_best_only=True,
                    monitor='accuracy', verbose=1),
                keras.callbacks.LearningRateScheduler(
                    self._update_learning_rate, verbose=1)
            ],
            validation_data=(x_val, y_val))

    def test(self, x_test: np.ndarray, y_test: np.ndarray) -> list:
        batch_size = x_test.shape[0]
        for i in range(2, int(batch_size / 2), -1):
            if batch_size % i == 0:
                batch_size = i
                break

        return self.model.evaluate(
            x_test, y_test,
            batch_size=batch_size,
            verbose=1)

    def save(self) -> None:
        """Save model."""
        tf.saved_model.save(self.model, self.model_path)

    def _generate_representative(self) -> Generator:
        # FIXME Move to data_utils module
        # FIXME Use real dataset -> 'Random. From training or evaluation set' --> Length of 100 for example
        for i in list(np.arange(100 * 32, dtype='float32').reshape((100, 32))):
            yield [i]

    def tflite_convert(self, integer_quant: bool = False) -> None:
        if not os.path.exists(self.model_path):
            raise RuntimeError(
                f"Saved model path '{self.model_path}' does not exist.")

        converter = tflite.TFLiteConverter.from_saved_model(
            self.model_path)
        converter.optimizations = [tflite.Optimize.DEFAULT]
        converter.experimental_enable_resource_variables = True

        if integer_quant:
            converter.inference_input_type = tf.uint8
            converter.inference_output_type = tf.uint8
            converter.representative_dataset = self._generate_representative
            converter.target_spec.supported_ops = [
                tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

            with open(self.tflite_quant_model_path, 'wb') as f:
                f.write(converter.convert())
        else:
            with open(self.tflite_model_path, 'wb') as f:
                f.write(converter.convert())


# handler = ModelHandler(dropout_keep_prob=0.5, clip_weights=True,
#                        add_batch_norm=True, add_residual=True,
#                        load=True)
# print(handler.model.summary())
# handler.train(epochs=4)
# handler.save()
# handler.tflite_convert(True)

Kaiming initialization as described in https://arxiv.org/pdf/1502.01852.pdf Kaiming initialization but cliping by norm 1 Handler for training proposed model in https://arxiv.org/pdf/1502.01852.pdf Create a handler for proposed model in https://arxiv.org/pdf/1502.01852.pdf.          Args:             num_2d_kpts (int, optional): Number of input 2D keypoints. Defaults to 16.             num_3d_kpts (int, optional): Number of input 3D keypoints. Defaults to 16.             save_dir (str, optional): Directory where store trained models and checkpoints. Defaults to 'training'.             load (bool, optional): Whether to load model from checkpoint file before training. Defaults to True.             linear_size (int, optional): Number of units in each layer of the model. Defaults to 1024.             num_bi_layers (int, optional): Number of bilinear blocks in the model. Defaults to 2.             dropout_keep_prob (float, optional): Dropout keep probability. Value '1.0' means no dropout. Defaults to 1.0.             clip_weights (bool, optional): Whether to clip weights to a norm of 1. Defaults to False.             add_batch_norm (bool, optional): Whether to use batch normalization. Defaults to False.             add_residual (bool, optional): Whether to add residual connections between bilinear blocks. Defaults to False.          Raises:             ValueError: if save_dir does not exist or it is not a directory. # SavedModel format # From tf official docs -> NOT from original repo #   As we cannot change lr each step with standard callbacks (only each epoch) # if epoch < 10: #     return lr # else: Save model. # FIXME Move to data_utils module # FIXME Use real dataset -> 'Random. From training or evaluation set' --> Length of 100 for example # handler = ModelHandler(dropout_keep_prob=0.5, clip_weights=True, #                        add_batch_norm=True, add_residual=True, #                        load=True) # print(handler.model.summary()) # handler.train(epochs=4) # handler.save() # handler.tflite_convert(True)