"""Time-series Generative Adversarial Networks (TimeGAN) Codebase.

Reference: Jinsung Yoon, Daniel Jarrett, Mihaela van der Schaar,
"Time-series Generative Adversarial Networks,"
Neural Information Processing Systems (NeurIPS), 2019.

Paper link: https://papers.nips.cc/paper/8789-time-series-generative-adversarial-networks

Last updated Date: April 24th 2020
Code author: Jinsung Yoon (jsyoon0823@gmail.com)

-----------------------------

predictive_metrics.py

Note: Use post-hoc RNN to classify original data and synthetic data

Output: discriminative score (np.abs(classification accuracy - 0.5))
"""

# mypy: ignore-errors
# Necessary Packages
import numpy as np
import tensorflow as tf
from sklearn.metrics import accuracy_score
from tf_slim.layers import layers as _layers

from timegan.utils import batch_generator, extract_time, train_test_divide


def discriminative_score_metrics(ori_data, generated_data):
    """Use post-hoc RNN to classify original data and synthetic data

    Args:
      - ori_data: original data
      - generated_data: generated synthetic data

    Returns:
      - discriminative_score: np.abs(classification accuracy - 0.5)
    """
    # Initialization on the Graph
    tf.compat.v1.reset_default_graph()

    # Basic Parameters
    no, seq_len, dim = np.asarray(ori_data).shape

    # Set maximum sequence length and each sequence length
    ori_time, ori_max_seq_len = extract_time(ori_data)
    generated_time, generated_max_seq_len = extract_time(ori_data)
    max_seq_len = max([ori_max_seq_len, generated_max_seq_len])

    ## Builde a post-hoc RNN discriminator network
    # Network parameters
    hidden_dim = int(dim / 2)
    iterations = 2000
    batch_size = 128

    # Input place holders
    # Feature
    X = tf.compat.v1.placeholder(tf.float32, [None, max_seq_len, dim], name="myinput_x")
    X_hat = tf.compat.v1.placeholder(tf.float32, [None, max_seq_len, dim], name="myinput_x_hat")

    T = tf.compat.v1.placeholder(tf.int32, [None], name="myinput_t")
    T_hat = tf.compat.v1.placeholder(tf.int32, [None], name="myinput_t_hat")

    # discriminator function
    def discriminator(x, t):
        """Simple discriminator function.

        Args:
          - x: time-series data
          - t: time information

        Returns:
          - y_hat_logit: logits of the discriminator output
          - y_hat: discriminator output
          - d_vars: discriminator variables
        """
        with tf.compat.v1.variable_scope("discriminator", reuse=tf.compat.v1.AUTO_REUSE) as vs:
            d_cell = tf.compat.v1.nn.rnn_cell.GRUCell(num_units=hidden_dim, activation=tf.nn.tanh, name="d_cell")
            d_outputs, d_last_states = tf.compat.v1.nn.dynamic_rnn(d_cell, x, dtype=tf.float32, sequence_length=t)
            y_hat_logit = _layers.fully_connected(d_last_states, 1, activation_fn=None)
            y_hat = tf.nn.sigmoid(y_hat_logit)
            d_vars = [v for v in tf.compat.v1.all_variables() if v.name.startswith(vs.name)]

        return y_hat_logit, y_hat, d_vars

    y_logit_real, y_pred_real, d_vars = discriminator(X, T)
    y_logit_fake, y_pred_fake, _ = discriminator(X_hat, T_hat)

    # Loss for the discriminator
    d_loss_real = tf.reduce_mean(
        input_tensor=tf.nn.sigmoid_cross_entropy_with_logits(logits=y_logit_real, labels=tf.ones_like(y_logit_real))
    )
    d_loss_fake = tf.reduce_mean(
        input_tensor=tf.nn.sigmoid_cross_entropy_with_logits(logits=y_logit_fake, labels=tf.zeros_like(y_logit_fake))
    )
    d_loss = d_loss_real + d_loss_fake

    # optimizer
    d_solver = tf.compat.v1.train.AdamOptimizer().minimize(d_loss, var_list=d_vars)

    ## Train the discriminator
    # Start session and initialize
    sess = tf.compat.v1.Session()
    sess.run(tf.compat.v1.global_variables_initializer())

    # Train/test division for both original and generated data
    train_x, train_x_hat, test_x, test_x_hat, train_t, train_t_hat, test_t, test_t_hat = train_test_divide(
        ori_data, generated_data, ori_time, generated_time
    )

    # Training step
    for itt in range(iterations):
        # Batch setting
        X_mb, T_mb = batch_generator(train_x, train_t, batch_size)
        X_hat_mb, T_hat_mb = batch_generator(train_x_hat, train_t_hat, batch_size)

        # Train discriminator
        _, step_d_loss = sess.run([d_solver, d_loss], feed_dict={X: X_mb, T: T_mb, X_hat: X_hat_mb, T_hat: T_hat_mb})

    ## Test the performance on the testing set
    y_pred_real_curr, y_pred_fake_curr = sess.run(
        [y_pred_real, y_pred_fake], feed_dict={X: test_x, T: test_t, X_hat: test_x_hat, T_hat: test_t_hat}
    )

    y_pred_final = np.squeeze(np.concatenate((y_pred_real_curr, y_pred_fake_curr), axis=0))
    y_label_final = np.concatenate(
        (
            np.ones(
                [
                    len(y_pred_real_curr),
                ]
            ),
            np.zeros(
                [
                    len(y_pred_real_curr),
                ]
            ),
        ),
        axis=0,
    )

    # Compute the accuracy
    acc = accuracy_score(y_label_final, (y_pred_final > 0.5))
    discriminative_score = np.abs(0.5 - acc)

    return discriminative_score
