#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Notice: 
[1] This code is only for reference.
Please modify the codes to fit your own data.
[2] The Code is based on TensorFlow 2.X.
Please install the TensorFlow 2.X version.
"""

import numpy as np
import pandas as pd

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow import keras

# Read Training Data
train_data = pd.read_csv('../source/S-traindata.csv', header=None)
train_data = np.array(train_data).astype('float32')

# Read Training Labels
train_labels = pd.read_csv('../source/S-trainlabel.csv', header=None)
train_labels = np.array(train_labels).astype('float32')
train_labels = tf.one_hot(indices=train_labels, depth=2)
train_labels = np.squeeze(train_labels)

# Read Testing Data
test_data = pd.read_csv('../source/S-testdata.csv', header=None)
test_data = np.array(test_data).astype('float32')

# Read Testing Labels
test_labels = pd.read_csv('../source/S-testlabel.csv', header=None)
test_labels = np.array(test_labels).astype('float32')
test_labels = tf.one_hot(indices=test_labels, depth=2)
test_labels = np.squeeze(test_labels)


class CatgoricalTP(tf.keras.metrics.Metric):
    def __init__(self, name='categorical_tp', **kwargs):
        super(CatgoricalTP, self).__init__(name=name, **kwargs)
        self.tp = self.add_weight(name='tp', initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_pred = tf.argmax(y_pred, axis=-1)
        y_true = tf.argmax(y_true, axis=-1)
        values = tf.equal(tf.cast(y_pred, 'int32'), tf.cast(y_true, 'int32'))
        values = tf.cast(values, 'float32')
        if sample_weight is not None:
            sample_weights = tf.cast(sample_weight, 'float32')
            values = tf.multiply(values, sample_weights)

        self.tp.assign_add(tf.reduce_sum(values))

    def result(self):
        return self.tp

    def reset_states(self):
        self.tp.assign(0.)


class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.5):
        super(TransformerBlock, self).__init__()
        self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = keras.Sequential([layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim), ])
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.ff_dim = ff_dim
        self.rate = rate

    def call(self, inputs, training):
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        out = self.layernorm2(out1 + ffn_output)

        return out


class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, embed_dim):
        super(TokenAndPositionEmbedding, self).__init__()
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)
        self.maxlen = maxlen
        self.embed_dim = embed_dim

    def call(self, x):
        positions = tf.range(start=0, limit=self.maxlen, delta=1)
        positions = self.pos_emb(positions)
        x = tf.reshape(x, [-1, self.maxlen, self.embed_dim])
        out = x + positions

        return out


maxlen = 3  # (Maximum) length of the signals
embed_dim = 97  # Number of features of one time point
num_heads = 8  # Number of attention heads
ff_dim = 64  # Hidden layer size in feed forward network inside transformer


def get_model():
    # Input Time-series
    inputs = layers.Input(shape=(maxlen * embed_dim,))
    embedding_layer = TokenAndPositionEmbedding(maxlen, embed_dim)
    x = embedding_layer(inputs)

    # Encoder Architecture
    transformer_block_1 = TransformerBlock(embed_dim=embed_dim, num_heads=num_heads, ff_dim=ff_dim)
    transformer_block_2 = TransformerBlock(embed_dim=embed_dim, num_heads=num_heads, ff_dim=ff_dim)
    x = transformer_block_1(x)
    x = transformer_block_2(x)

    # Output
    x = layers.GlobalMaxPooling1D()(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Dense(64, activation="relu")(x)
    x = layers.Dropout(0.5)(x)
    outputs = layers.Dense(2, activation="softmax")(x)

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

    return model


model = get_model()
model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4),
              loss="categorical_crossentropy",
              metrics=["accuracy", CatgoricalTP()])

history = model.fit(
    train_data, train_labels, batch_size=64, epochs=100, validation_data=(test_data, test_labels)
)

model.save_weights('model_weight')