# @author lucasmiranda42
# encoding: utf-8
# module deepof

"""keras_tuner hypermodels for hyperparameter tuning of deep autoencoders in deepof.models."""

from typing import Any, NewType

import numpy as np
import tensorflow_probability as tfp
from keras_tuner import HyperModel

import deepof.model_utils
import deepof.models

tfd = tfp.distributions
tfpl = tfp.layers

# DEFINE CUSTOM ANNOTATED TYPES #
project = NewType("deepof_project", Any)
coordinates = NewType("deepof_coordinates", Any)
table_dict = NewType("deepof_table_dict", Any)


class VaDE(HyperModel):
    """Hyperparameter tuning pipeline for deepof.models.VaDE."""

    def __init__(
        self,
        input_shape: tuple,
        latent_dim: int,
        batch_size: int,
        n_components: int = 10,
        learn_rate: float = 1e-3,
        edge_feature_shape: tuple = None,
        use_gnn: bool = False,
        adjacency_matrix: np.ndarray = None,
    ):
        """Build VaDE hypermodel for hyperparameter tuning.

        Args:
            input_shape (tuple): shape of the input tensor.
            latent_dim (int): dimension of the latent space.
            batch_size (int): batch size for training.
            learn_rate (float): learning rate for the optimizer.
            n_components (int): number of components in the quantization space.
            edge_feature_shape (tuple): shape of the edge feature tensor.
            use_gnn (bool): whether to use a graph neural network to encode the input data.
            adjacency_matrix (np.ndarray): adjacency matrix of the graph.

        """
        super().__init__()
        self.input_shape = input_shape
        self.latent_dim = latent_dim
        self.batch_size = batch_size
        self.learn_rate = learn_rate
        self.n_components = n_components
        self.edge_feature_shape = edge_feature_shape
        self.use_gnn = use_gnn
        self.adjacency_matrix = adjacency_matrix

    def get_hparams(self, hp):
        """Retrieve hyperparameters to tune."""
        # Architectural hyperparameters
        encoder = hp.Choice(
            "encoder", ["recurrent", "TCN", "transformer"], default="recurrent"
        )
        kmeans_loss = hp.Float(
            "kmeans_loss", min_value=0.0, max_value=1.0, sampling="linear"
        )
        kl_annealing_mode = hp.Choice(
            "kl_annealing_mode", ["linear", "sigmoid"], default="linear"
        )
        kl_warmup_epochs = hp.Int(
            "kl_warmup_epochs", min_value=0, max_value=100, sampling="linear"
        )
        cluster_assignment_regularizer = hp.Float(
            "cluster_assignment_regularizer",
            min_value=0.0,
            max_value=1.0,
            sampling="linear",
        )

        return (
            encoder,
            kmeans_loss,
            kl_annealing_mode,
            kl_warmup_epochs,
            cluster_assignment_regularizer,
        )

    def build(self, hp):
        """Override Hypermodel's build method."""
        # Hyperparameters to tune
        (
            encoder,
            kmeans_loss,
            kl_annealing_mode,
            kl_warmup_epochs,
            cluster_assignment_regularizer,
        ) = self.get_hparams(hp)

        vade = deepof.models.VaDE(
            input_shape=self.input_shape,
            edge_feature_shape=self.edge_feature_shape,
            adjacency_matrix=self.adjacency_matrix,
            use_gnn=self.use_gnn,
            latent_dim=self.latent_dim,
            n_components=self.n_components,
            batch_size=self.batch_size,
            # hyperparameters to tune
            encoder_type=encoder,
            kmeans_loss=kmeans_loss,
            kl_annealing_mode=kl_annealing_mode,
            kl_warmup_epochs=kl_warmup_epochs,
            reg_cat_clusters=cluster_assignment_regularizer,
        )
        vade.compile()

        return vade


class VQVAE(HyperModel):
    """Hyperparameter tuning pipeline for deepof.models.VQVAE."""

    def __init__(
        self,
        input_shape: tuple,
        latent_dim: int,
        n_components: int = 10,
        learn_rate: float = 1e-3,
        edge_feature_shape: tuple = None,
        use_gnn: bool = False,
        adjacency_matrix: np.ndarray = None,
    ):
        """VQVAE hypermodel for hyperparameter tuning.

        Args:
            input_shape (tuple): shape of the input tensor.
            latent_dim (int): dimension of the latent space.
            learn_rate (float): learning rate for the optimizer.
            n_components (int): number of components in the quantization space.
            edge_feature_shape (tuple): shape of the edge feature tensor.
            use_gnn (bool): whether to use a graph neural network to encode the input data.
            adjacency_matrix (np.ndarray): adjacency matrix of the graph.

        """
        super().__init__()
        self.input_shape = input_shape
        self.latent_dim = latent_dim
        self.learn_rate = learn_rate
        self.n_components = n_components
        self.edge_feature_shape = edge_feature_shape
        self.use_gnn = use_gnn
        self.adjacency_matrix = adjacency_matrix

    def get_hparams(self, hp):
        """Retrieve hyperparameters to tune, including the encoder type and the weight of the kmeans loss."""
        # Architectural hyperparameters
        encoder = hp.Choice(
            "encoder", ["recurrent", "TCN", "transformer"], default="recurrent"
        )
        kmeans_loss = hp.Float(
            "kmeans_loss", min_value=0.0, max_value=1.0, sampling="linear"
        )

        return (encoder, kmeans_loss)

    def build(self, hp):
        """Override Hypermodel's build method."""
        # Hyperparameters to tune
        (encoder, kmeans_loss) = self.get_hparams(hp)

        vqvae = deepof.models.VQVAE(
            input_shape=self.input_shape,
            edge_feature_shape=self.edge_feature_shape,
            adjacency_matrix=self.adjacency_matrix,
            use_gnn=self.use_gnn,
            latent_dim=self.latent_dim,
            n_components=self.n_components,
            # hyperparameters to tune
            encoder_type=encoder,
            kmeans_loss=kmeans_loss,
        )
        vqvae.compile()

        return vqvae


class Contrastive(HyperModel):
    """Hyperparameter tuning pipeline for deepof.models.Contrastive."""

    def __init__(
        self,
        input_shape: tuple,
        latent_dim: int,
        learn_rate: float = 1e-3,
        edge_feature_shape: tuple = None,
        use_gnn: bool = False,
        adjacency_matrix: np.ndarray = None,
    ):
        """Contrastive hypermodel for hyperparameter tuning.

        Args:
            input_shape (tuple): shape of the input tensor.
            latent_dim (int): dimension of the latent space.
            learn_rate (float): learning rate for the optimizer.
            edge_feature_shape (tuple): shape of the edge feature tensor.
            use_gnn (bool): whether to use a graph neural network to encode the input data.
            adjacency_matrix (np.ndarray): adjacency matrix of the graph.

        """
        super().__init__()
        self.input_shape = input_shape
        self.latent_dim = latent_dim
        self.learn_rate = learn_rate
        self.edge_feature_shape = edge_feature_shape
        self.use_gnn = use_gnn
        self.adjacency_matrix = adjacency_matrix

    def get_hparams(self, hp):
        """Retrieve hyperparameters to tune, including the encoder type and the weight of the kmeans loss."""
        # Architectural hyperparameters
        encoder = hp.Choice(
            "encoder", ["recurrent", "TCN", "transformer"], default="recurrent"
        )

        return encoder

    def build(self, hp):
        """Override Hypermodel's build method."""
        # Hyperparameters to tune
        (encoder) = self.get_hparams(hp)

        contrastive = deepof.models.Contrastive(
            input_shape=self.input_shape,
            edge_feature_shape=self.edge_feature_shape,
            adjacency_matrix=self.adjacency_matrix,
            use_gnn=self.use_gnn,
            latent_dim=self.latent_dim,
            # hyperparameters to tune
            encoder_type=encoder,
        )
        contrastive.compile()

        return contrastive
