import tensorflow as tf
from dlomix.constants import ALPHABET_UNMOD
from tensorflow.keras.layers.experimental import preprocessing


class DeepLCRetentionTimePredictor(tf.keras.Model):
    def __init__(
        self, seq_length=60, vocab_dict=ALPHABET_UNMOD, use_global_features=False
    ):
        super(DeepLCRetentionTimePredictor, self).__init__()
        self.seq_length = seq_length
        self._use_global_features = use_global_features

        self.leaky_relu = tf.keras.layers.ReLU(max_value=20, negative_slope=0.1)
        self.string_lookup = preprocessing.StringLookup(
            vocabulary=list(vocab_dict.keys())
        )

        self._build_aminoacid_branch()
        self._build_diaminoacid_branch()
        self._build_onehot_encoding_branch()
        self._build_regressor()
        self.output_layer = tf.keras.layers.Dense(1)

        if self._use_global_features:
            self._build_global_features_branch()

    def _build_aminoacid_branch(self):
        self.aminoacid_branch = tf.keras.Sequential(
            [
                self._build_conv_pool_block(n_filters=256, kernel=8, padding="same"),
                self._build_conv_pool_block(n_filters=128, kernel=8, padding="same"),
                self._build_conv_pool_block(
                    n_filters=64, kernel=8, padding="same", pool=False
                ),
                tf.keras.layers.Flatten(),
            ]
        )

    def _build_diaminoacid_branch(self):

        self.diaminoacid_branch = tf.keras.Sequential(
            [
                self._build_conv_pool_block(n_filters=128, kernel=2, padding="same"),
                self._build_conv_pool_block(n_filters=64, kernel=2, padding="same"),
                tf.keras.layers.Flatten(),
            ]
        )

    def _build_global_features_branch(self):
        self.global_features_branch = tf.keras.Sequential(
            [
                tf.keras.layers.Dense(16, activation=self.leaky_relu),
                tf.keras.layers.Dense(16, activation=self.leaky_relu),
                tf.keras.layers.Dense(16, activation=self.leaky_relu),
            ]
        )

    def _build_onehot_encoding_branch(self):
        self.onehot_encoding_branch = tf.keras.Sequential(
            [
                self._build_conv_pool_block(
                    n_filters=2,
                    kernel=2,
                    padding="same",
                    activation="tanh",
                    pool_strides=10,
                    pool_size=10,
                ),
                tf.keras.layers.Flatten(),
            ]
        )

    def _build_regressor(self):
        self.regressor = tf.keras.Sequential(
            [tf.keras.layers.Dense(128, activation=self.leaky_relu) for _ in range(5)]
        )

    def _build_conv_pool_block(
        self,
        n_conv_layers=2,
        n_filters=256,
        kernel=8,
        padding="same",
        activation="leaky_relu",
        pool=True,
        pool_strides=2,
        pool_size=2,
    ):

        # leaky relu by default
        activation_fn = self.leaky_relu

        if activation in ["tanh", "relu"]:
            activation_fn = activation

        block = tf.keras.Sequential(
            [
                tf.keras.layers.Conv1D(
                    filters=n_filters,
                    kernel_size=kernel,
                    padding=padding,
                    activation=activation_fn,
                )
                for _ in range(n_conv_layers)
            ]
        )
        if pool:
            block.add(tf.keras.layers.MaxPooling1D(pool_size=2, strides=2))

        return block

    def call(self, inputs, **kwargs):
        outputs = {}

        integer_encoded = self.string_lookup(inputs["seq"])
        onehot_encoded = tf.one_hot(integer_encoded, depth=self.seq_length)

        if self._use_global_features:
            outputs["global_features_output"] = self.global_features_branch(
                inputs["global_features"]
            )

        outputs["onehot_branch_output"] = self.onehot_encoding_branch(onehot_encoded)
        outputs["aminoacids_branch_output"] = self.aminoacid_branch(inputs["counts"])
        outputs["diaminoacids_branch_output"] = self.diaminoacid_branch(
            inputs["di_counts"]
        )

        concatenated_output = tf.concat(outputs.values(), axis=1)
        concatenated_output = self.regressor(concatenated_output)
        return self.output_layer(concatenated_output)
