#!/usr/bin/env python3

# encoding=utf-8

"""

@Author: Bruce Shuyue Jia

@Date: Nov 5, 2021

This is a code only for reference - two-class classification

"""

import numpy as np

import pandas as pd

import tensorflow as tf

from tensorflow.keras import layers

from tensorflow import keras

print("TensorFlow version:", tf.__version__)

# Read Training Data - Target Domain

train_data = pd.read_csv('../target/T-traindata.csv', header=None)

train_data = np.array(train_data).astype('float32')

# Read Training Labels - Target Domain

train_labels = pd.read_csv('../target/T-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 - Target Domain

test_data = pd.read_csv('../target/T-testdata.csv', header=None)

test_data = np.array(test_data).astype('float32')

# Read Testing Labels - Target Domain

test_labels = pd.read_csv('../target/T-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.)

# Transformer Model

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

embed_dim = 97  # Embedding size for each token

num_heads = 8  # Number of attention heads

ff_dim = 64  # Hidden layer size in feed forward network inside transformer

def get_model():

    # Create a simple 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

# Load the pre-trained model

pretrained_model = get_model()

pretrained_model.load_weights('../transformer/model_weight')

# Make sure the original model will not be updated

# i.e., freezen the model parameters

pretrained_model.trainable = False

# Show the model architecture

pretrained_model.summary()

print('\n\n\n\n')

# Remove the last three layers

num_layer = len(pretrained_model.layers)

pretrained_part = tf.keras.models.Sequential(pretrained_model.layers[0:(num_layer - 3)])

pretrained_part.summary()

print('\n\n\n\n')

# Add some new layers for fine tuning

add_layers = tf.keras.models.Sequential([

    tf.keras.layers.Dense(256, activation='relu'),

    tf.keras.layers.Dropout(0.5),

    tf.keras.layers.Dense(128, activation='relu'),

    tf.keras.layers.Dropout(0.5),

    tf.keras.layers.Dense(2, activation='softmax')

])

# Build a new model 

rebuild_model = tf.keras.models.Sequential([

    pretrained_part,

    add_layers

])

rebuild_model.build((maxlen * embed_dim,))

rebuild_model.summary()

print('\n\n\n\n')

# Train and evaluate the new model

model = rebuild_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.evaluate(test_data, test_labels, verbose=2)