cnn_aae_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FpESl5A2sgIQ",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "References:  \n",
    "1. [Time series classification baseline with CNN](https://github.com/cauchyturing/UCR_Time_Series_Classification_Deep_Learning_Baseline/blob/master/FCN.py)\n",
    "2. GAN implementation examples:  \n",
    "    1. [Image Captioning with Attention](https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/contrib/eager/python/examples/generative_examples/image_captioning_with_attention.ipynb)\n",
    "    2. [Convolutional VAE](https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Mount Google Drive as remote drive\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 127
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 23435,
     "status": "ok",
     "timestamp": 1558215065460,
     "user": {
      "displayName": "Dmitry Balabka",
      "photoUrl": "https://lh3.googleusercontent.com/-JJUbezxdVCY/AAAAAAAAAAI/AAAAAAAAAI4/SFNVd_aY2Ro/s64/photo.jpg",
      "userId": "17224709416263358599"
     },
     "user_tz": -180
    },
    "id": "BZ2KDGlbsgIU",
    "outputId": "63093979-e1ad-4008-9a48-b47b9b6b992d",
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/gdrive')\n",
    "import os\n",
    "os.chdir('/content/gdrive/My Drive/TSI/Master\\'s thesis/space-type-recognition')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "if 'COLAB_TPU_ADDR' in os.environ:\n",
    "  TPU_ENABLED = True\n",
    "  TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
    "  print('TPU address is', TPU_WORKER)\n",
    "\n",
    "  with tf.Session(TPU_WORKER) as session:\n",
    "    devices = session.list_devices()\n",
    "    \n",
    "  print('TPU devices:')\n",
    "  display(devices)\n",
    "else:\n",
    "  TPU_ENABLED = False\n",
    "  \n",
    "nb_epochs = 2000\n",
    "\n",
    "%connect_info\n"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Install required libraries"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install tensorflow-gpu==1.13.1 tensorboard==1.13.1"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Import base libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SZUvxhT8sgIb",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras.utils as np_utils\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.callbacks import ReduceLROnPlateau, TensorBoard\n",
    "from tensorflow.keras import layers, regularizers, optimizers\n",
    "from tensorflow.contrib.gan.python.losses.python.losses_impl import wasserstein_discriminator_loss, wasserstein_generator_loss\n",
    "from timeseries.helper import plot_model_recursive\n",
    "from timeseries.data import load_ucr_dataset\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score, f1_score\n",
    "from timeseries.data import Dataset\n",
    "from timeseries.transformers import AaeClustering, AdversarialAutoencoder\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "from functools import reduce\n",
    "\n",
    "from sklearn.preprocessing import minmax_scale, scale\n",
    "\n",
    "import sys\n",
    "\n",
    "display(sys.version)\n",
    "display(tf.VERSION)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Define Adversarial Autoencoder (AAE) Model and Blocks\n",
    "\n",
    "Originally, AAE is described by [https://arxiv.org/abs/1511.05644](https://arxiv.org/abs/1511.05644).\n",
    "More detailed description can be found [https://pdfs.semanticscholar.org/7989/533d5169904998271ecfa5377c21404c7348.pdf](https://pdfs.semanticscholar.org/7989/533d5169904998271ecfa5377c21404c7348.pdf)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "See Tensorflow documentation to [create custom Models](https://www.tensorflow.org/alpha/guide/keras/custom_layers_and_models#putting_it_all_together_an_end-to-end_example) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def create_encoder(\n",
    "        input_shape,\n",
    "        #classes_num,\n",
    "        code_len,\n",
    "        name='encoder',\n",
    "        reg=regularizers.l1_l2(1e-7, 0),\n",
    "        units=256,\n",
    "        renorm=True):\n",
    "    \n",
    "    if input_shape[0] > units:\n",
    "        raise Exception('Encoder input length is %s but expecting less then %s units size. Otherwise autoencoder will not convergence by reconstruction loss' % (input_shape[0], units))\n",
    "    inputs = layers.Input(input_shape, name=name + \"_input\")\n",
    "    conv = inputs\n",
    "    \n",
    "    pad_units = units - input_shape[0]\n",
    "    if pad_units > 0:\n",
    "        conv = layers.ZeroPadding1D(padding=(0, (pad_units)))(conv)\n",
    "    \n",
    "    # conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "        \n",
    "    conv = layers.Conv1D(int(units / 4), 8, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # conv = layers.Activation('relu')(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(int(units / 4), 8, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # conv = layers.Activation('relu')(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    \n",
    "    conv = layers.MaxPooling1D()(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(int(units / 2), 5, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # conv = layers.Activation('relu')(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(int(units / 2), 5, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # conv = layers.Activation('relu')(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    \n",
    "    conv = layers.MaxPooling1D()(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(units, 3, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # conv = layers.Activation('relu')(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    \n",
    "    conv = layers.MaxPooling1D()(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(units, 3, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    \n",
    "    if False:\n",
    "        conv = layers.GlobalAveragePooling2D()(conv)\n",
    "        \n",
    "        conv = layers.Flatten()(conv)\n",
    "        \n",
    "        conv = layers.Dense(code_len)(conv)\n",
    "        # conv = layers.Dense(code_len, kernel_regularizer=reg)(conv)\n",
    "        # TODO: Dropout?\n",
    "        # TODO: Activation sigmoid?\n",
    "        conv = layers.Activation('relu')(conv)\n",
    "        # conv = layers.Activation('sigmoid')(conv)\n",
    "    \n",
    "    output = conv\n",
    "\n",
    "    return tf.keras.Model(inputs, output, name=name)\n",
    "            \n",
    "        \n",
    "        \n",
    "def create_decoder(\n",
    "        input_shape,\n",
    "        output_shape,\n",
    "        #classes_num,\n",
    "        code_len,\n",
    "        name='decoder',\n",
    "        units=256,\n",
    "        renorm=True):\n",
    "    \n",
    "    inputs = layers.Input(input_shape, name=name + \"_input\")\n",
    "    conv = inputs\n",
    "    if False:\n",
    "        inputs = layers.Input((code_len, ), name=name + \"_input\")\n",
    "        conv = inputs\n",
    "        \n",
    "        conv = layers.Dense(units=output_shape[0], activation=tf.nn.relu)(conv)\n",
    "        conv = layers.Reshape(target_shape=output_shape)(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(units, 3, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # conv = layers.Activation('relu')(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    conv = layers.UpSampling1D(size=2)(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(int(units / 2), 5, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # conv = layers.Activation('relu')(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    conv = layers.UpSampling1D(size=2)(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(int(units / 4), 8, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    conv = layers.UpSampling1D(size=2)(conv)\n",
    "    \n",
    "    # conv = layers.Conv2DTranspose(1, 251, 1, padding='same')(conv)\n",
    "    \n",
    "    conv = layers.Conv1D(output_shape[1], 8, padding='same')(conv)\n",
    "    conv = layers.BatchNormalization(renorm=renorm)(conv)\n",
    "    # last layer with linear or sigmoid activation\n",
    "    conv = layers.Activation('linear')(conv)\n",
    "    # conv = layers.Activation('sigmoid')(conv)\n",
    "    \n",
    "    if conv.shape[1] > output_shape[0]: \n",
    "        conv = layers.Cropping1D(cropping=(0, conv.shape[1].value - output_shape[0]))(conv)\n",
    "    else:\n",
    "        conv = layers.ZeroPadding1D(padding=(0, output_shape[0] - conv.shape[1].value))(conv)\n",
    "    \n",
    "    output = conv\n",
    "    return tf.keras.Model(inputs, output, name=name)\n",
    "\n",
    "    \n",
    "def create_generator_softmax(\n",
    "        inputs,\n",
    "        generator,\n",
    "        classes_num,\n",
    "        name='generator_softmax',\n",
    "        trainable=True,\n",
    "        renorm=True):\n",
    "    \n",
    "        generator = layers.Dense(512, trainable=trainable)(generator(inputs))\n",
    "        generator = layers.BatchNormalization(renorm=renorm)(generator)\n",
    "        generator = layers.LeakyReLU(0.2)(generator)\n",
    "        generator = layers.Dense(classes_num, trainable=trainable)(generator)\n",
    "        generator = layers.BatchNormalization(momentum=0.8, renorm=renorm)(generator)\n",
    "        output = layers.Softmax(trainable=trainable)(generator)\n",
    "        \n",
    "        return tf.keras.Model(inputs, output, name=name)\n",
    "\n",
    "\n",
    "def create_generator_linear(\n",
    "        inputs,\n",
    "        generator,\n",
    "        code_len,\n",
    "        name='generator_linear',\n",
    "        trainable=True,\n",
    "        renorm=True,\n",
    "        **kwargs):\n",
    "    \n",
    "        generator = layers.Dense(code_len, trainable=trainable)(generator(inputs))\n",
    "        generator = layers.BatchNormalization(momentum=0.8, renorm=renorm)(generator)\n",
    "\n",
    "        return tf.keras.Model(inputs, generator, name=name)\n",
    "\n",
    "def create_generator_linear_softmax(\n",
    "        inputs,\n",
    "        generator,\n",
    "        code_len,\n",
    "        classes_num,\n",
    "        dense_size=1024,\n",
    "        name='generator',\n",
    "        trainable=True,\n",
    "        renorm=True,\n",
    "        reg=regularizers.l1_l2(1e-7, 1e-7),\n",
    "):\n",
    "    generator = layers.Flatten()(generator(inputs))\n",
    "    generator = layers.Dense(dense_size, trainable=trainable, kernel_regularizer=reg)(generator)\n",
    "    generator = layers.BatchNormalization(renorm=renorm)(generator)\n",
    "    generator = layers.LeakyReLU(0.2)(generator)\n",
    "    generator = tf.keras.Model(inputs, generator, name=name)\n",
    "    \n",
    "    generator_linear = create_generator_linear(inputs, generator, code_len)\n",
    "    \n",
    "    generator_softmax = create_generator_softmax(inputs, generator, classes_num)\n",
    "    \n",
    "    return generator_linear, generator_softmax\n",
    "\n",
    "def create_generator_concatenate(\n",
    "        inputs,\n",
    "        generator_softmax,\n",
    "        generator_linear,\n",
    "        decoder,\n",
    "        classes_num,\n",
    "        code_len,\n",
    "        name='generator_concatinate'):\n",
    "    \n",
    "        autoencoder = layers.Concatenate()([generator_linear(inputs), generator_softmax(inputs)])\n",
    "        \n",
    "        autoencoder = layers.Dense(code_len + classes_num)(autoencoder)\n",
    "        autoencoder = layers.LeakyReLU(0.2)(autoencoder)\n",
    "        \n",
    "        autoencoder = layers.Dense(code_len + classes_num)(autoencoder)\n",
    "        autoencoder = layers.LeakyReLU(0.2)(autoencoder)\n",
    "        \n",
    "        autoencoder = layers.Dense(reduce(lambda a, b: a*b, decoder.input_shape[1:]))(autoencoder)\n",
    "        # TODO: Do we need ReLU here? In original implementation it has linear activation\n",
    "        # autoencoder = layers.LeakyReLU(0.2)(autoencoder)\n",
    "        # autoencoder = layers.Dense(code_len)(autoencoder)\n",
    "        # autoencoder = layers.LeakyReLU(0.2)(autoencoder)\n",
    "        autoencoder = layers.Reshape(decoder.input_shape[1:])(autoencoder)\n",
    "        output = decoder(autoencoder)\n",
    "        \n",
    "        return tf.keras.Model(inputs=inputs, outputs=output, name=name)\n",
    "    \n",
    "    \n",
    "def create_discriminator(\n",
    "        latent_dim,\n",
    "        output_dim=1,\n",
    "        hidden_dim=256,\n",
    "        reg=regularizers.l1_l2(1e-7, 1e-7),\n",
    "        name='discriminator',\n",
    "        **kwargs):\n",
    "    \n",
    "        with K.name_scope('Discriminator'):\n",
    "            inputs = layers.Input((latent_dim,))\n",
    "            h = layers.Dense(hidden_dim, kernel_regularizer=reg)(inputs)\n",
    "            h = layers.LeakyReLU(0.2)(h)\n",
    "            h = layers.Dense(hidden_dim, kernel_regularizer=reg)(h)\n",
    "            h = layers.LeakyReLU(0.2)(h)\n",
    "            \n",
    "            output = layers.Dense(output_dim, activation=\"sigmoid\", kernel_regularizer=reg)(h)\n",
    "            \n",
    "            return tf.keras.Model(inputs, output, name=name)\n",
    "    \n",
    "# class AAE(tf.keras.Model):\n",
    "#     def __init__(self, latent_dim):\n",
    "#         super(AAE, self).__init__()\n",
    "#         self.latent_dim = latent_dim\n",
    "#         \n",
    "#     def call(self, inputs, training=None, mask=None):\n",
    "#         sequence_encoder = generator = Encoder()(inputs)\n",
    "#         sequence_decoder = Decoder()(sequence_encoder)\n",
    "#         return sequence_decoder\n",
    "\n",
    "class Output(layers.Layer):\n",
    "\n",
    "    def __init__(self, output_dim, **kwargs):\n",
    "        self.output_dim = output_dim\n",
    "        super(Output, self).__init__(**kwargs)\n",
    "        \n",
    "def set_trainable(model: tf.keras.Model, is_trainable):\n",
    "    model.trainable = is_trainable\n",
    "    for layer in model.layers:\n",
    "        if type(layer) is tf.keras.Model:\n",
    "            set_trainable(layer, is_trainable)\n",
    "        else:\n",
    "            layer.trainable = is_trainable\n",
    "                \n",
    "def create_adversarial_autoencoder(\n",
    "        input_shape,\n",
    "        classes_num=10, \n",
    "        code_len=36,\n",
    "        autoencoder_optimizer=None,\n",
    "        adversarial_model_optimizer=None,\n",
    "        generator_optimizer=None,\n",
    "        sequence_autoencoder_optimizer=None,\n",
    "        autoencoder_loss='mean_squared_error',\n",
    "        adversarial_model_loss='binary_crossentropy',\n",
    "        generator_loss='binary_crossentropy',\n",
    "        sequence_autoencoder_loss='mean_squared_error',\n",
    "        generator_softmax_optimizer=None,\n",
    "        generator_softmax_loss=None,\n",
    "        units=256\n",
    "        ):\n",
    "    \n",
    "    if autoencoder_optimizer is None:\n",
    "        # For LSTM autoencoder use RMSprop: https://github.com/keras-team/keras/issues/1401#issuecomment-168773845\n",
    "        # or Adam https://github.com/keras-team/keras/issues/1401#issuecomment-169295237\n",
    "        autoencoder_optimizer = optimizers.Adam(lr=1e-5)\n",
    "    if sequence_autoencoder_optimizer is None:\n",
    "        sequence_autoencoder_optimizer = optimizers.Adam(lr=0.001)\n",
    "    if adversarial_model_optimizer is None:\n",
    "        # Use SGD for descriminator training: https://github.com/soumith/ganhacks#9-use-the-adam-optimizer\n",
    "        # adversarial_model_optimizer = SGD(lr=0.01)\n",
    "        adversarial_model_optimizer = optimizers.Adam(lr=1e-3, beta_1=0.1)\n",
    "    if generator_optimizer is None:\n",
    "        # For LSTM autoencoder use RMSprop: https://github.com/keras-team/keras/issues/1401#issuecomment-168773845\n",
    "        # or Adam https://github.com/keras-team/keras/issues/1401#issuecomment-169295237\n",
    "        generator_optimizer = optimizers.Adam(lr=1e-2, beta_1=0.1)\n",
    "    if generator_softmax_optimizer is None:\n",
    "        generator_softmax_optimizer = optimizers.Adam(lr=1e-3)\n",
    "    if generator_softmax_loss is None:\n",
    "        generator_softmax_loss = 'categorical_crossentropy'\n",
    "    \n",
    "    output_shape = input_shape\n",
    "    timeseries_count=input_shape[1]\n",
    "    \n",
    "    sequence_encoder = generator = create_encoder(input_shape=input_shape, code_len=code_len, units=units)\n",
    "    sequence_decoder = decoder = create_decoder(input_shape=generator.output_shape[1:], output_shape=output_shape, code_len=code_len+classes_num, units=units)\n",
    "    inputs = generator_input = sequence_encoder.inputs\n",
    "    \n",
    "    # generator_softmax = create_generator_softmax(inputs, generator, classes_num=classes_num)\n",
    "    \n",
    "    generator_linear, generator_softmax = create_generator_linear_softmax(inputs, generator, classes_num=classes_num, code_len=code_len)\n",
    "    \n",
    "    generator_softmax.compile(optimizer=generator_softmax_optimizer, loss=generator_softmax_loss, metrics=['accuracy'])\n",
    "    \n",
    "    # generator_linear = create_generator_linear(inputs, generator, code_len=code_len)\n",
    "    \n",
    "    #\n",
    "    # Model for reconstruction phase\n",
    "    #\n",
    "    \n",
    "    semisupervised_autoencoder = autoencoder = create_generator_concatenate(\n",
    "        inputs,\n",
    "        generator_softmax,\n",
    "        generator_linear,\n",
    "        decoder,\n",
    "        classes_num=classes_num, \n",
    "        code_len=code_len)\n",
    "    autoencoder = autoencoder(inputs)\n",
    "    autoencoder = tf.keras.Model(inputs=inputs, outputs=autoencoder, name='autoencoder')\n",
    "    autoencoder.compile(optimizer=autoencoder_optimizer, loss=autoencoder_loss)\n",
    "    semisupervised_autoencoder.compile(optimizer=autoencoder_optimizer, loss=autoencoder_loss)\n",
    "    \n",
    "    print(\"----------\\nAutoencoder\\n----------\")\n",
    "    autoencoder.summary()\n",
    "    print(\"----------\\nSemi-supervised decoder\\n----------\")\n",
    "    semisupervised_autoencoder.summary()\n",
    "    print(\"----------\\nGenerator softmax\\n----------\")\n",
    "    generator_softmax.summary()\n",
    "    print(\"----------\\nGenerator linear\\n----------\")\n",
    "    generator_linear.summary()\n",
    "    \n",
    "    #\n",
    "    # Prepare some reusable parts\n",
    "    #\n",
    "    \n",
    "    y_onehot_fake = generator_softmax\n",
    "    z_fake = generator_linear\n",
    "    \n",
    "    discriminator_z = create_discriminator(code_len, name='discriminator_z_output')\n",
    "    discriminator_y = create_discriminator(classes_num, name='discriminator_y_output')\n",
    "    \n",
    "    #\n",
    "    # Model for adversarial phase - update only discriminator\n",
    "    #\n",
    "    \n",
    "    set_trainable(y_onehot_fake, False)\n",
    "    set_trainable(z_fake, False)\n",
    "    \n",
    "    # TODO: move to separate build function\n",
    "    y_onehot_true = layers.Input(y_onehot_fake.output_shape[-1:], name='y_onehot_true_input')\n",
    "    # (12, )\n",
    "    z_true = layers.Input(z_fake.output_shape[-1:], name='z_true_input')\n",
    "\n",
    "    dz_fake = discriminator_z(z_fake(inputs))\n",
    "    dz_fake = Output(1, name='dz_fake_output')(dz_fake)\n",
    "\n",
    "    dy_fake = discriminator_y(y_onehot_fake(inputs))\n",
    "    dy_fake = Output(1, name='dy_fake_output')(dy_fake)\n",
    "\n",
    "    dz_true = discriminator_z(z_true)\n",
    "    dz_true = Output(1, name='dz_true_output')(dz_true)\n",
    "\n",
    "    dy_true = discriminator_y(y_onehot_true)\n",
    "    dy_true = Output(1, name='dy_true_output')(dy_true)\n",
    "\n",
    "    outputs = [dz_fake,\n",
    "               dy_fake,\n",
    "               dz_true,\n",
    "               dy_true]\n",
    "\n",
    "    adversarial_model = tf.keras.Model(inputs=inputs + [y_onehot_true, z_true], outputs=outputs, name='adversarial_model')\n",
    "\n",
    "    adversarial_model.compile(optimizer=adversarial_model_optimizer, loss=adversarial_model_loss)\n",
    "    \n",
    "    print(\"----------\\nAdversarial Model\\n----------\")\n",
    "    adversarial_model.summary()\n",
    "    \n",
    "    #\n",
    "    # Generator phase - update only encoder\n",
    "    #\n",
    "    \n",
    "    # We will reuse discriminators in adversarial model, but we do not want to train it\n",
    "    # Important that discriminator it self will be trainable, but not in scope of adversarial model\n",
    "    set_trainable(discriminator_z, False)\n",
    "    set_trainable(discriminator_y, False)\n",
    "    set_trainable(y_onehot_fake, True)\n",
    "    set_trainable(z_fake, True)\n",
    "    \n",
    "    # TODO: move separate build function\n",
    "    \n",
    "    outputs = [dz_fake, dy_fake]\n",
    "\n",
    "    generator_model = tf.keras.Model(inputs=inputs, outputs=outputs, name='generator_model')\n",
    "    generator_model.compile(optimizer=generator_optimizer, loss=generator_loss)\n",
    "\n",
    "    print(\"----------\\nGenerator Model\\n----------\")\n",
    "    generator_model.summary()\n",
    "    \n",
    "    sequence_autoencoder = autoencoder\n",
    "    \n",
    "    return autoencoder, adversarial_model, generator_model, generator_softmax, generator_linear, sequence_autoencoder, semisupervised_autoencoder\n",
    "\n",
    "# create_adversarial_autoencoder(input_shape=(251, 1, 1,))\n",
    "\n",
    "def readucr(filename, type):\n",
    "    data = np.loadtxt(os.path.join('data', 'UCR_TS_Archive_2015', filename, filename + '_' + type.upper()), delimiter = ',')\n",
    "    Y = data[:,0]\n",
    "    X = data[:,1:]\n",
    "    return X, Y\n",
    "\n",
    "def readsussex(type, position, sensors=None):\n",
    "    if sensors is None:\n",
    "        sensors = ['Acc', 'Gyr']\n",
    "    data = None\n",
    "    for sensor in sensors:\n",
    "        for axis in ['x', 'y', 'z']:\n",
    "            path = os.path.join('data', type, position, sensor + '_' + axis + '.txt')\n",
    "            print('Read file %s' % (path))\n",
    "            tmp = np.loadtxt(path, delimiter = ' ')\n",
    "            print('Reshape')\n",
    "            tmp = np.reshape(tmp, tmp.shape + (1,))\n",
    "            print('Append')\n",
    "            if data is None:\n",
    "                data = tmp\n",
    "            else:\n",
    "                data = np.concatenate((data, tmp), axis=2)\n",
    "    print(\"Size in memory %d MB\" % ((data.size * data.itemsize) / 1024 / 1024))\n",
    "    return data\n",
    "\n",
    "def readsussex_labels(type, position, path = None):\n",
    "    path = os.path.join('data', type, position, 'Label.txt') if path is None else path\n",
    "    data = np.loadtxt(path, delimiter = ' ')\n",
    "    print(\"Size in memory is %d MB\" % ((data.size * data.itemsize) / 1024 / 1024))\n",
    "    return data \n",
    "    \n",
    "\n",
    "from sklearn.preprocessing import StandardScaler, MinMaxScaler\n",
    "from sklearn.pipeline import Pipeline\n",
    "\n",
    "\n",
    "def prepare_data(dataset_name, is_conv2d=False, feature_range=(-1, 1)):\n",
    "    X_train, y_train = readucr(dataset_name, 'train')\n",
    "    X_test, y_test = readucr(dataset_name, 'test')\n",
    "    \n",
    "    nb_classes = len(np.unique(y_test))\n",
    "    # batch_size = int(min(X_train.shape[0]/10, 16))\n",
    "\n",
    "    # y_train = np.repeat(np.reshape(y_train, (y_train.shape[0], 1, 1)), X_train.shape[1], 1)    \n",
    "    # y_test = np.repeat(np.reshape(y_test, (y_test.shape[0], 1, 1)), X_train.shape[1], 1) \n",
    "    \n",
    "    # y_train = (y_train - y_train.min())/(y_train.max()-y_train.min())*(nb_classes-1)\n",
    "    # y_test = (y_test - y_test.min())/(y_test.max()-y_test.min())*(nb_classes-1)\n",
    "    # y_train = np_utils.to_categorical(y_train, nb_classes)\n",
    "    # y_test = np_utils.to_categorical(y_test, nb_classes)\n",
    "    \n",
    "    # X_train_mean = X_train.mean()\n",
    "    # X_train_std = X_train.std()\n",
    "    # X_train = (X_train - X_train_mean)/(X_train_std)\n",
    "    #  \n",
    "    # X_test = (X_test - X_train_mean)/(X_train_std)\n",
    "    \n",
    "    pipeline = Pipeline([\n",
    "        # ('standard_scale', StandardScaler()),\n",
    "        ('minmax_scale', MinMaxScaler(feature_range=feature_range)),\n",
    "    ])\n",
    "    X_train, X_test, _ = np.split(\n",
    "        pipeline.fit_transform(np.concatenate((X_train, X_test), axis=0).transpose()).transpose(),\n",
    "        [X_train.shape[0], X_train.shape[0] + X_test.shape[0]]\n",
    "    )\n",
    "    \n",
    "    X_train_new_shape = X_train.shape + ((1,1,) if is_conv2d else (1,))\n",
    "    X_test_new_shape = X_test.shape + ((1,1,) if is_conv2d else (1,))\n",
    "    \n",
    "    X_train = X_train.reshape(X_train_new_shape)\n",
    "    X_test = X_test.reshape(X_test_new_shape)\n",
    "    \n",
    "    input_shape = X_train.shape[1:]\n",
    "    \n",
    "    return input_shape, X_train, y_train, X_test, y_test\n",
    "\n",
    "def prepare_sussex_data(\n",
    "        batch_size,\n",
    "        feature_range=(-1, 1),\n",
    "        test_size=0.66,\n",
    "        include_test=False,\n",
    "):\n",
    "    x_validate_hand = np.load('data/validate_Hand_AG.npy')#[:, :, :1]\n",
    "    y_validate_hand = np.load('data/validate_Hand_Label.npy')\n",
    "    y_validate_hand = AaeClustering.get_labels_for_segments(\n",
    "        np.reshape(y_validate_hand, y_validate_hand.shape + (1,))\n",
    "    ).astype(int)\n",
    "    # make class numbers starting from 0\n",
    "    y_validate_hand -= 1\n",
    "    x_test_hand = np.load('data/test_Hand_AG.npy')#[:, :, :1]\n",
    "    \n",
    "    x_train, x_test, y_train, y_test = train_test_split(\n",
    "        x_validate_hand,\n",
    "        y_validate_hand,\n",
    "        test_size=test_size,\n",
    "        random_state=0,\n",
    "        stratify=y_validate_hand\n",
    "    )\n",
    "    \n",
    "    # normalize time series\n",
    "    pipeline = Pipeline([\n",
    "        #('minmax_scale', MinMaxScaler(feature_range=feature_range)),\n",
    "        ('standard_scale', StandardScaler()),\n",
    "    ])\n",
    "    for axis in range(x_train.shape[2]):\n",
    "        x_train[:, :, axis] = pipeline.fit_transform(x_train[:, :, axis])\n",
    "        x_test[:, :, axis] = pipeline.fit_transform(x_test[:, :, axis])\n",
    "        x_test_hand[:, :, axis] = pipeline.fit_transform(x_test_hand[:, :, axis])\n",
    "    \n",
    "    classes_num = np.unique(y_validate_hand).shape[0]\n",
    "    \n",
    "    test_data = (x_test, y_test)\n",
    "    if include_test:\n",
    "        test_data = (\n",
    "            np.concatenate((x_test, x_test_hand), axis=0), \n",
    "            np.concatenate((y_test, np.zeros(x_test_hand.shape[0], int)))\n",
    "        )\n",
    "    \n",
    "    dataset = Dataset(\n",
    "        (x_train, y_train),\n",
    "        test_data,\n",
    "        num_labeled_data=x_train.shape[0], \n",
    "        num_classes=classes_num,\n",
    "        use_test_as_unlabeled=True,\n",
    "        batch_size=batch_size)\n",
    "    dataset.summary()\n",
    "    \n",
    "    return dataset, x_train, y_train, x_test, y_test, classes_num\n",
    "\n",
    "def prepare_sussex_test_data():\n",
    "    x_validate_hand = np.load('data/validate_Hand_AG.npy')#[:, :, :1]\n",
    "    y_validate_hand = np.load('data/validate_Hand_Label.npy')\n",
    "    y_validate_hand = AaeClustering.get_labels_for_segments(\n",
    "        np.reshape(y_validate_hand, y_validate_hand.shape + (1,))\n",
    "    ).astype(int)\n",
    "    # make class numbers starting from 0\n",
    "    y_validate_hand -= 1\n",
    "    x_test = np.load('data/test_Hand_AG.npy')#[:, :, :1]\n",
    "    \n",
    "    x_train, y_train = x_validate_hand, y_validate_hand\n",
    "    \n",
    "    # normalize time series\n",
    "    pipeline = Pipeline([\n",
    "        #('minmax_scale', MinMaxScaler(feature_range=feature_range)),\n",
    "        ('standard_scale', StandardScaler()),\n",
    "    ])\n",
    "    for axis in range(x_train.shape[2]):\n",
    "        x_train[:, :, axis] = pipeline.fit_transform(x_train[:, :, axis])\n",
    "        x_test[:, :, axis] = pipeline.fit_transform(x_test[:, :, axis])\n",
    "    \n",
    "    classes_num = np.unique(y_validate_hand).shape[0]\n",
    "        \n",
    "    return x_train, y_train, x_test, classes_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "dataset, x_train, y_train, x_test, y_test, classes_num = prepare_sussex_data(1000)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "is_executing": false,
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "from math import ceil\n",
    "from tensorflow.python.keras.callbacks import CallbackList\n",
    "from timeseries.profiler import Stopwatch\n",
    "\n",
    "def train_aae(\n",
    "        classes_num=3,\n",
    "        code_len=36,\n",
    "        labeled_data_num=11,\n",
    "        batchsize=None,\n",
    "        epochs=500,\n",
    "        verbose=True,\n",
    "        wgan=True,\n",
    "        adversarial_model_loss='binary_crossentropy',\n",
    "        generator_loss='binary_crossentropy',\n",
    "        evaluation_step=5,\n",
    "        profiling=False\n",
    "):\n",
    "    stopwatch = Stopwatch(disable=(not profiling))\n",
    "    input_shape, X_train, y_train, X_test, y_test = prepare_data('ArrowHead')\n",
    "    \n",
    "    if wgan:\n",
    "        # adversarial_model_loss = self.wasserstein_loss\n",
    "        # generator_loss = self.wasserstein_loss\n",
    "        adversarial_model_loss = wasserstein_discriminator_loss\n",
    "        generator_loss = wasserstein_generator_loss\n",
    "        \n",
    "    # Create and compile model\n",
    "    autoencoder, adversarial_model, generator_model, generator_softmax, generator_linear, sequence_autoencoder, semisupervised_autoencoder = create_adversarial_autoencoder(\n",
    "        input_shape,\n",
    "        classes_num=classes_num,\n",
    "        code_len=code_len,\n",
    "        adversarial_model_loss=adversarial_model_loss,\n",
    "        generator_loss=generator_loss,\n",
    "    )\n",
    "    \n",
    "    exclude_models_by_name = []\n",
    "    plot_model_recursive(adversarial_model, exclude_models_by_name=exclude_models_by_name)\n",
    "    plot_model_recursive(autoencoder, exclude_models_by_name=exclude_models_by_name)\n",
    "            \n",
    "    dataset = Dataset(\n",
    "        (X_train, AaeClustering.get_labels_for_segments(y_train).astype(int)),\n",
    "        (X_test, AaeClustering.get_labels_for_segments(y_test).astype(int)),\n",
    "        num_labeled_data=labeled_data_num, \n",
    "        num_classes=classes_num,\n",
    "        use_test_as_unlabeled=True,\n",
    "        batch_size=batchsize)\n",
    "    dataset.summary()\n",
    "    batchsize = dataset.batch_size\n",
    "    steps = dataset.steps\n",
    "\n",
    "    progbar = tf.keras.callbacks.ProgbarLogger(count_mode='steps')\n",
    "    tensorboard = tf.keras.callbacks.TensorBoard(\n",
    "        log_dir='./tb_logs/adversarial-autoencoder-' + datetime.datetime.today().strftime('%Y%m%dT%H%M'),\n",
    "        histogram_freq=False,\n",
    "        # histogram_freq=20,\n",
    "        write_graph=True,\n",
    "        # TODO: this option might affect logs size\n",
    "        write_grads=False\n",
    "    )\n",
    "    \n",
    "    callbacks = CallbackList([\n",
    "        progbar,\n",
    "        tensorboard])\n",
    "    callbacks.set_params({\n",
    "        'epochs': epochs,\n",
    "        'verbose': verbose,\n",
    "        'steps': steps,\n",
    "        'metrics': []})\n",
    "    callbacks.set_model(autoencoder)\n",
    "    \n",
    "    model_checkpoint = tf.keras.callbacks.ModelCheckpoint('saved_models/aae_generator_softmax-{epoch:02d}.hdf5', monitor='accuracy', verbose=1)\n",
    "    model_checkpoint.set_model(generator_softmax)\n",
    "    callbacks.append(model_checkpoint)\n",
    "    \n",
    "    autoencoder_out_labels = ['mse']\n",
    "    sequence_autoencoder_out_labels = ['mse']\n",
    "    \n",
    "    callbacks.on_train_begin()\n",
    "    \n",
    "    for epoch in np.arange(0, epochs):\n",
    "        epoch_logs = {}\n",
    "        \n",
    "        callbacks.on_epoch_begin(epoch)\n",
    "        \n",
    "        # dataset should be shuffled on each epoch\n",
    "        dataset.shuffle()\n",
    "        \n",
    "        # TODO: add learning rate adjustments\n",
    "        \n",
    "        for index in range(0, steps):\n",
    "            \n",
    "            X_train_batch = dataset.sample_unlabeled_minibatch(batchsize)\n",
    "            if labeled_data_num > 0:\n",
    "                X_train_batch_labeled, y_l, y_onehot_l = dataset.sample_labeled_minibatch(batchsize)\n",
    "            actual_batch_size = X_train_batch.shape[0]\n",
    "            \n",
    "            batch_logs = {'batch': index, 'size': actual_batch_size}\n",
    "            \n",
    "            callbacks.on_batch_begin(index)\n",
    "            \n",
    "            # reconstruction phase\n",
    "            stopwatch.start('reconstruction phase')\n",
    "            autoencoder_outs = autoencoder.train_on_batch(X_train_batch, X_train_batch)\n",
    "            batch_logs = AdversarialAutoencoder.get_logs(batch_logs, autoencoder_outs, autoencoder_out_labels)\n",
    "            stopwatch.stop('reconstruction phase')\n",
    "            \n",
    "            # if labeled_data_num > 0:\n",
    "            #     stopwatch.start('semisupervised phase')\n",
    "            #     autoencoder_outs = semisupervised_autoencoder.train_on_batch([X_train_batch_labeled, y_onehot_l], X_train_batch_labeled)\n",
    "            #     batch_logs = AdversarialAutoencoder.get_logs(batch_logs, autoencoder_outs, autoencoder_out_labels)\n",
    "            #     stopwatch.stop('semisupervised phase')\n",
    "            \n",
    "            # adversarial phase\n",
    "            stopwatch.start('adversarial phase')\n",
    "            # TODO: might be vise versa https://github.com/eriklindernoren/Keras-GAN/blob/master/aae/aae.py#L120\n",
    "            if wgan:\n",
    "                class_true = np.ones(actual_batch_size, dtype=np.int32)\n",
    "                class_fake = -class_true\n",
    "            else:\n",
    "                class_fake = np.zeros(actual_batch_size, dtype=np.int32)\n",
    "                class_true = np.ones(actual_batch_size, dtype=np.int32)\n",
    "                \n",
    "            adversarial_input = {'encoder_input': X_train_batch,\n",
    "                                 'z_true_input': AdversarialAutoencoder.gaussian(actual_batch_size, code_len),\n",
    "                                 'y_onehot_true_input': AdversarialAutoencoder.onehot_categorical(actual_batch_size, classes_num)}\n",
    "            adversarial_output = {'dz_fake_output': class_fake,\n",
    "                                  'dy_fake_output': class_fake,\n",
    "                                  'dz_true_output': class_true,\n",
    "                                  'dy_true_output': class_true}\n",
    "            \n",
    "            adversarial_model.train_on_batch(adversarial_input, adversarial_output)\n",
    "            \n",
    "            # generator phase\n",
    "            generator_input = {'encoder_input': X_train_batch}\n",
    "            generator_output = {'dz_fake_output': class_true,\n",
    "                                'dy_fake_output': class_true}\n",
    "\n",
    "            generator_model.train_on_batch(generator_input, generator_output)\n",
    "            stopwatch.stop('adversarial phase')\n",
    "            \n",
    "            # supervised phase\n",
    "            stopwatch.start('supervised phase')\n",
    "            if labeled_data_num > 0:\n",
    "                autoencoder_outs = generator_softmax.train_on_batch(X_train_batch_labeled, y_onehot_l)\n",
    "                batch_logs = AdversarialAutoencoder.get_logs(batch_logs, autoencoder_outs, autoencoder_out_labels)\n",
    "            stopwatch.stop('supervised phase')\n",
    "            \n",
    "            # TODO: there was progbar instead callbacks. It might affect logs size\n",
    "            callbacks.on_batch_end(index, batch_logs)\n",
    "            \n",
    "        #\n",
    "        # Evaluate autoencoder\n",
    "        #\n",
    "\n",
    "        if epoch % evaluation_step == 0:\n",
    "            stopwatch.start('evaluate phase')\n",
    "            def evaluate_autoencoder(epoch_logs, X_train, logs_name):\n",
    "                autoencoder_outs = autoencoder.evaluate(\n",
    "                    X_train, X_train,\n",
    "                    batch_size=batchsize,\n",
    "                    verbose=0)\n",
    "                return AdversarialAutoencoder.get_logs(epoch_logs, autoencoder_outs, [logs_name])\n",
    "    \n",
    "            epoch_logs = evaluate_autoencoder(epoch_logs, X_train, 'autoencoder_loss')\n",
    "            epoch_logs = evaluate_autoencoder(epoch_logs, X_test, 'autoencoder_loss_test')\n",
    "    \n",
    "            #\n",
    "            # Evaluate adversarial model\n",
    "            #\n",
    "    \n",
    "            class_true = np.zeros(X_train.shape[0], dtype=np.int32)\n",
    "            class_fake = np.ones(X_train.shape[0], dtype=np.int32)\n",
    "            class_true_test = np.zeros(X_test.shape[0], dtype=np.int32)\n",
    "            class_fake_test = np.ones(X_test.shape[0], dtype=np.int32)\n",
    "    \n",
    "            def evaluate_adversarial_model(epoch_logs, X_train, code_len, class_true, class_fake, logs_name):\n",
    "                adversarial_input = {'encoder_input': X_train,\n",
    "                                     'z_true_input': AdversarialAutoencoder.gaussian(X_train.shape[0], code_len),\n",
    "                                     # TODO: fix errors in some cases. You might use `code_len` instead of `self.classes_num`\n",
    "                                     'y_onehot_true_input': AdversarialAutoencoder.onehot_categorical(X_train.shape[0], classes_num)}\n",
    "                adversarial_output = {'dz_fake_output': class_fake,\n",
    "                                      'dy_fake_output': class_fake,\n",
    "                                      'dz_true_output': class_true,\n",
    "                                      'dy_true_output': class_true\n",
    "                                      }\n",
    "                adversarial_model_outs = adversarial_model.evaluate(\n",
    "                    adversarial_input, adversarial_output,\n",
    "                    batch_size=batchsize,\n",
    "                    verbose=0)\n",
    "                return AdversarialAutoencoder.get_logs(epoch_logs, adversarial_model_outs, [logs_name])\n",
    "    \n",
    "            epoch_logs = evaluate_adversarial_model(epoch_logs, X_train, code_len, class_true, class_fake, 'adversarial_loss')\n",
    "            epoch_logs = evaluate_adversarial_model(epoch_logs, X_test, code_len, class_true_test, class_fake_test, 'adversarial_loss_test')\n",
    "    \n",
    "            #\n",
    "            # Generator phase\n",
    "            #\n",
    "    \n",
    "            # train\n",
    "            def evaluate_generator(epoch_logs, X, class_true, logs_name):\n",
    "                generator_input = {'encoder_input': X}\n",
    "                generator_output = {'dz_fake_output': class_true,\n",
    "                                    'dy_fake_output': class_true}\n",
    "                generator_model_outs = generator_model.evaluate(\n",
    "                    generator_input, generator_output,\n",
    "                    batch_size=batchsize,\n",
    "                    verbose=0)\n",
    "                return AdversarialAutoencoder.get_logs(epoch_logs, generator_model_outs, [logs_name])\n",
    "    \n",
    "            epoch_logs = evaluate_generator(epoch_logs, X_train, class_true, 'generator_loss')\n",
    "            epoch_logs = evaluate_generator(epoch_logs, X_test, class_true_test, 'generator_loss_test')\n",
    "    \n",
    "            # eval clustering\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_train, y_train, generator_softmax)], ['adjusted_rand_score'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_test, y_test, generator_softmax)], ['adjusted_rand_score_test'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_train, y_train, generator_softmax, metric_func=accuracy_score)], ['accuracy_score'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_test, y_test, generator_softmax, metric_func=accuracy_score)], ['accuracy_score_test'])\n",
    "            metric_func = lambda y_true, y_pred: f1_score(y_true, y_pred, average='macro')\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_train, y_train, generator_softmax, metric_func=metric_func)], ['f1_score'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_test, y_test, generator_softmax, metric_func=metric_func)], ['f1_score_test'])\n",
    "\n",
    "            stopwatch.stop('evaluate phase')\n",
    "            \n",
    "        callbacks.on_epoch_end(epoch, epoch_logs)\n",
    "            \n",
    "    callbacks.on_train_end()\n",
    "    \n",
    "    sequence_autoencoder = autoencoder\n",
    "    \n",
    "    # TODO: implement plot prediction\n",
    "    # self.plotPrediction(data, sequence_autoencoder)\n",
    "    \n",
    "    predictions = autoencoder.predict(X_test[:1])\n",
    "    for axis in range(predictions.shape[2]):\n",
    "        plt.plot(np.reshape(X_test[:1, :, axis:axis+1], (X_test.shape[1])))\n",
    "        plt.plot(np.reshape(predictions[:, :, axis:axis+1], (predictions.shape[1])))\n",
    "        plt.show()\n",
    "    \n",
    "    return generator_softmax, sequence_autoencoder, autoencoder, adversarial_model, generator_model, generator_linear\n",
    "\n",
    "train_aae()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "from math import ceil\n",
    "from tensorflow.python.keras.callbacks import CallbackList\n",
    "from timeseries.data import Dataset\n",
    "from timeseries.transformers import AaeClustering, AdversarialAutoencoder\n",
    "from tensorflow.contrib.gan.python.losses.python.losses_impl import wasserstein_discriminator_loss, wasserstein_generator_loss\n",
    "from timeseries.profiler import Stopwatch\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "def train_sussex_aae(\n",
    "        code_len=50,\n",
    "        batchsize=50,\n",
    "        epochs=50,\n",
    "        verbose=True,\n",
    "        wgan=True,\n",
    "        adversarial_model_loss='binary_crossentropy',\n",
    "        generator_loss='binary_crossentropy',\n",
    "        evaluation_step=1,\n",
    "        profiling=False,\n",
    "        units=512,\n",
    "        test_size=0.66,\n",
    "        include_test=True,\n",
    "        experiment_name=None,\n",
    "):\n",
    "    if experiment_name is None:\n",
    "        experiment_name = '{:.5}'.format(str(abs(hash(frozenset(\n",
    "            {'code_len':code_len, 'batchsize': batchsize, 'units': \n",
    "                units, 'test_size': test_size, 'include_test': include_test}.items()\n",
    "        )))))\n",
    "        \n",
    "    print('Experiment name: %s' % (experiment_name,))\n",
    "    \n",
    "    stopwatch = Stopwatch(disable=(not profiling))\n",
    "    \n",
    "    dataset, X_train, y_train, X_test, y_test, classes_num = prepare_sussex_data(batchsize, test_size=test_size, include_test=include_test)\n",
    "    batchsize = dataset.batch_size\n",
    "    steps = dataset.steps\n",
    "    labeled_data_num = dataset.get_num_labeled_data()\n",
    "    \n",
    "    input_shape = X_train.shape[1:]\n",
    "    \n",
    "    if wgan:\n",
    "        # adversarial_model_loss = self.wasserstein_loss\n",
    "        # generator_loss = self.wasserstein_loss\n",
    "        adversarial_model_loss = wasserstein_discriminator_loss\n",
    "        generator_loss = wasserstein_generator_loss\n",
    "        \n",
    "    # Create and compile model\n",
    "    autoencoder, adversarial_model, generator_model, generator_softmax, generator_linear, sequence_autoencoder, semisupervised_autoencoder = create_adversarial_autoencoder(\n",
    "        input_shape,\n",
    "        classes_num=classes_num,\n",
    "        code_len=code_len,\n",
    "        adversarial_model_loss=adversarial_model_loss,\n",
    "        generator_loss=generator_loss,\n",
    "        units=units\n",
    "    )\n",
    "    \n",
    "    exclude_models_by_name = []\n",
    "    plot_model_recursive(adversarial_model, exclude_models_by_name=exclude_models_by_name)\n",
    "    plot_model_recursive(autoencoder, exclude_models_by_name=exclude_models_by_name)\n",
    "\n",
    "    progbar = tf.keras.callbacks.ProgbarLogger(count_mode='steps')\n",
    "    tensorboard = tf.keras.callbacks.TensorBoard(\n",
    "        log_dir='./tb_logs/adversarial-autoencoder-' + experiment_name + '-' + datetime.datetime.today().strftime('%Y%m%dT%H%M'),\n",
    "        histogram_freq=False,\n",
    "        # histogram_freq=20,\n",
    "        write_graph=True,\n",
    "        # TODO: this option might affect logs size\n",
    "        write_grads=False\n",
    "    )\n",
    "    \n",
    "    callbacks = CallbackList([\n",
    "        progbar,\n",
    "        tensorboard])\n",
    "    callbacks.set_params({\n",
    "        'epochs': epochs,\n",
    "        'verbose': verbose,\n",
    "        'steps': steps,\n",
    "        'metrics': []})\n",
    "    callbacks.set_model(autoencoder)\n",
    "    \n",
    "    model_checkpoint = tf.keras.callbacks.ModelCheckpoint('saved_models/sussex_aae_generator_softmax-' + experiment_name + '-{epoch:02d}.hdf5', monitor='accuracy', verbose=1)\n",
    "    model_checkpoint.set_model(generator_softmax)\n",
    "    callbacks.append(model_checkpoint)\n",
    "    \n",
    "    autoencoder_out_labels = ['mse']\n",
    "    sequence_autoencoder_out_labels = ['mse']\n",
    "    \n",
    "    callbacks.on_train_begin()\n",
    "    \n",
    "    for epoch in np.arange(0, epochs):\n",
    "        epoch_logs = {}\n",
    "        \n",
    "        callbacks.on_epoch_begin(epoch)\n",
    "        \n",
    "        # dataset should be shuffled on each epoch\n",
    "        dataset.shuffle()\n",
    "        \n",
    "        # TODO: add learning rate adjustments\n",
    "        \n",
    "        for index in range(0, steps):\n",
    "            \n",
    "            X_train_batch = dataset.sample_unlabeled_minibatch(batchsize)\n",
    "            if labeled_data_num > 0:\n",
    "                X_train_batch_labeled, y_l, y_onehot_l = dataset.sample_labeled_minibatch(batchsize)\n",
    "            actual_batch_size = X_train_batch.shape[0]\n",
    "            \n",
    "            batch_logs = {'batch': index, 'size': actual_batch_size}\n",
    "            \n",
    "            callbacks.on_batch_begin(index)\n",
    "            \n",
    "            # reconstruction phase\n",
    "            stopwatch.start('reconstruction phase')\n",
    "            autoencoder_outs = autoencoder.train_on_batch(X_train_batch, X_train_batch)\n",
    "            batch_logs = AdversarialAutoencoder.get_logs(batch_logs, autoencoder_outs, autoencoder_out_labels)\n",
    "            stopwatch.stop('reconstruction phase')\n",
    "            \n",
    "            # if labeled_data_num > 0:\n",
    "            #     stopwatch.start('semisupervised phase')\n",
    "            #     autoencoder_outs = semisupervised_autoencoder.train_on_batch([X_train_batch_labeled, y_onehot_l], X_train_batch_labeled)\n",
    "            #     batch_logs = AdversarialAutoencoder.get_logs(batch_logs, autoencoder_outs, autoencoder_out_labels)\n",
    "            #     stopwatch.stop('semisupervised phase')\n",
    "            \n",
    "            # adversarial phase\n",
    "            stopwatch.start('adversarial phase')\n",
    "            # TODO: might be vise versa https://github.com/eriklindernoren/Keras-GAN/blob/master/aae/aae.py#L120\n",
    "            if wgan:\n",
    "                class_true = np.ones(actual_batch_size, dtype=np.int32)\n",
    "                class_fake = -class_true\n",
    "            else:\n",
    "                class_fake = np.zeros(actual_batch_size, dtype=np.int32)\n",
    "                class_true = np.ones(actual_batch_size, dtype=np.int32)\n",
    "                \n",
    "            adversarial_input = {'encoder_input': X_train_batch,\n",
    "                                 'z_true_input': AdversarialAutoencoder.gaussian(actual_batch_size, code_len),\n",
    "                                 'y_onehot_true_input': AdversarialAutoencoder.onehot_categorical(actual_batch_size, classes_num)}\n",
    "            adversarial_output = {'dz_fake_output': class_fake,\n",
    "                                  'dy_fake_output': class_fake,\n",
    "                                  'dz_true_output': class_true,\n",
    "                                  'dy_true_output': class_true}\n",
    "            \n",
    "            adversarial_model.train_on_batch(adversarial_input, adversarial_output)\n",
    "            \n",
    "            # generator phase\n",
    "            generator_input = {'encoder_input': X_train_batch}\n",
    "            generator_output = {'dz_fake_output': class_true,\n",
    "                                'dy_fake_output': class_true}\n",
    "\n",
    "            generator_model.train_on_batch(generator_input, generator_output)\n",
    "            stopwatch.stop('adversarial phase')\n",
    "            \n",
    "            # supervised phase\n",
    "            stopwatch.start('supervised phase')\n",
    "            if labeled_data_num > 0:\n",
    "                autoencoder_outs = generator_softmax.train_on_batch(X_train_batch_labeled, y_onehot_l)\n",
    "                batch_logs = AdversarialAutoencoder.get_logs(batch_logs, autoencoder_outs, autoencoder_out_labels)\n",
    "            stopwatch.stop('supervised phase')\n",
    "            \n",
    "            # TODO: there was progbar instead callbacks. It might affect logs size\n",
    "            callbacks.on_batch_end(index, batch_logs)\n",
    "            \n",
    "        #\n",
    "        # Evaluate autoencoder\n",
    "        #\n",
    "\n",
    "        if epoch % evaluation_step == 0:\n",
    "            stopwatch.start('evaluate phase')\n",
    "            def evaluate_autoencoder(epoch_logs, X_train, logs_name):\n",
    "                autoencoder_outs = autoencoder.evaluate(\n",
    "                    X_train, X_train,\n",
    "                    batch_size=batchsize,\n",
    "                    verbose=0)\n",
    "                return AdversarialAutoencoder.get_logs(epoch_logs, autoencoder_outs, [logs_name])\n",
    "    \n",
    "            epoch_logs = evaluate_autoencoder(epoch_logs, X_train, 'autoencoder_loss')\n",
    "            epoch_logs = evaluate_autoencoder(epoch_logs, X_test, 'autoencoder_loss_test')\n",
    "    \n",
    "            #\n",
    "            # Evaluate adversarial model\n",
    "            #\n",
    "    \n",
    "            class_true = np.zeros(X_train.shape[0], dtype=np.int32)\n",
    "            class_fake = np.ones(X_train.shape[0], dtype=np.int32)\n",
    "            class_true_test = np.zeros(X_test.shape[0], dtype=np.int32)\n",
    "            class_fake_test = np.ones(X_test.shape[0], dtype=np.int32)\n",
    "    \n",
    "            def evaluate_adversarial_model(epoch_logs, X_train, code_len, class_true, class_fake, logs_name):\n",
    "                adversarial_input = {'encoder_input': X_train,\n",
    "                                     'z_true_input': AdversarialAutoencoder.gaussian(X_train.shape[0], code_len),\n",
    "                                     # TODO: fix errors in some cases. You might use `code_len` instead of `self.classes_num`\n",
    "                                     'y_onehot_true_input': AdversarialAutoencoder.onehot_categorical(X_train.shape[0], classes_num)}\n",
    "                adversarial_output = {'dz_fake_output': class_fake,\n",
    "                                      'dy_fake_output': class_fake,\n",
    "                                      'dz_true_output': class_true,\n",
    "                                      'dy_true_output': class_true\n",
    "                                      }\n",
    "                adversarial_model_outs = adversarial_model.evaluate(\n",
    "                    adversarial_input, adversarial_output,\n",
    "                    batch_size=batchsize,\n",
    "                    verbose=0)\n",
    "                return AdversarialAutoencoder.get_logs(epoch_logs, adversarial_model_outs, [logs_name])\n",
    "    \n",
    "            epoch_logs = evaluate_adversarial_model(epoch_logs, X_train, code_len, class_true, class_fake, 'adversarial_loss')\n",
    "            epoch_logs = evaluate_adversarial_model(epoch_logs, X_test, code_len, class_true_test, class_fake_test, 'adversarial_loss_test')\n",
    "    \n",
    "            #\n",
    "            # Generator phase\n",
    "            #\n",
    "    \n",
    "            # train\n",
    "            def evaluate_generator(epoch_logs, X, class_true, logs_name):\n",
    "                generator_input = {'encoder_input': X}\n",
    "                generator_output = {'dz_fake_output': class_true,\n",
    "                                    'dy_fake_output': class_true}\n",
    "                generator_model_outs = generator_model.evaluate(\n",
    "                    generator_input, generator_output,\n",
    "                    batch_size=batchsize,\n",
    "                    verbose=0)\n",
    "                return AdversarialAutoencoder.get_logs(epoch_logs, generator_model_outs, [logs_name])\n",
    "    \n",
    "            epoch_logs = evaluate_generator(epoch_logs, X_train, class_true, 'generator_loss')\n",
    "            epoch_logs = evaluate_generator(epoch_logs, X_test, class_true_test, 'generator_loss_test')\n",
    "    \n",
    "            # eval clustering\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_train, y_train, generator_softmax)], ['adjusted_rand_score'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_test, y_test, generator_softmax)], ['adjusted_rand_score_test'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_train, y_train, generator_softmax, metric_func=accuracy_score)], ['accuracy_score'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_test, y_test, generator_softmax, metric_func=accuracy_score)], ['accuracy_score_test'])\n",
    "            metric_func = lambda y_pred, y_true: f1_score(y_pred, y_true, average='macro')\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_train, y_train, generator_softmax, metric_func=metric_func)], ['f1_score'])\n",
    "            epoch_logs = AdversarialAutoencoder.get_logs(epoch_logs, [AaeClustering.evaluate(X_test, y_test, generator_softmax, metric_func=metric_func)], ['f1_score_test'])\n",
    "\n",
    "            stopwatch.stop('evaluate phase')\n",
    "            \n",
    "        callbacks.on_epoch_end(epoch, epoch_logs)\n",
    "            \n",
    "    callbacks.on_train_end()\n",
    "    \n",
    "    sequence_autoencoder = autoencoder\n",
    "    \n",
    "    # TODO: implement plot prediction\n",
    "    # self.plotPrediction(data, sequence_autoencoder)\n",
    "    \n",
    "    predictions = autoencoder.predict(X_test[:1])\n",
    "    for axis in range(predictions.shape[2]):\n",
    "        plt.plot(np.reshape(X_test[:1, :, axis:axis+1], (X_test.shape[1])))\n",
    "        plt.plot(np.reshape(predictions[:, :, axis:axis+1], (predictions.shape[1])))\n",
    "        plt.show()\n",
    "    \n",
    "    return generator_softmax, sequence_autoencoder, autoencoder, adversarial_model, generator_model, generator_linear\n",
    "\n",
    "# Experiment w/o test data and large batchsize\n",
    "train_sussex_aae(batchsize=500, include_test=False)\n",
    "# Experiment w/o test data\n",
    "train_sussex_aae(batchsize=50, include_test=False)\n",
    "# Final experiment\n",
    "train_sussex_aae(batchsize=50, include_test=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def create_autoencoder(\n",
    "        input_shape,\n",
    "        code_len,\n",
    "        name='autoencoder'\n",
    "):\n",
    "    output_shape = input_shape\n",
    "    encoder = create_encoder(input_shape=input_shape, code_len=code_len)\n",
    "    display(encoder.output_shape)\n",
    "    decoder = create_decoder(\n",
    "        input_shape=encoder.output_shape[1:], \n",
    "        output_shape=output_shape,\n",
    "        code_len=code_len)\n",
    "    \n",
    "    autoencoder_optimizer = optimizers.Adam(lr=0.001)\n",
    "    autoencoder_loss='mean_squared_error'\n",
    "    \n",
    "    autoencoder = tf.keras.Model(encoder.inputs, decoder(encoder(encoder.inputs)), name=name)\n",
    "    autoencoder.compile(\n",
    "        optimizer=autoencoder_optimizer, \n",
    "        loss=autoencoder_loss,)\n",
    "    \n",
    "    return autoencoder\n",
    "    \n",
    "def train_autoencoder(epochs=500, verbose=True, batch_size=36):\n",
    "    input_shape, X_train, y_train, X_test, y_test = prepare_data('ArrowHead')\n",
    "    \n",
    "    display(input_shape)\n",
    "    \n",
    "    autoencoder = create_autoencoder(input_shape, code_len=36)\n",
    "    plot_model_recursive(autoencoder)\n",
    "    \n",
    "    tensorboard = tf.keras.callbacks.TensorBoard(\n",
    "        log_dir='./tb_logs/cnn_autoencoder-' + datetime.datetime.today().strftime('%Y%m%dT%H%M'),\n",
    "        histogram_freq=False,\n",
    "        # histogram_freq=20,\n",
    "        write_graph=True,\n",
    "        # TODO: this option might affect logs size\n",
    "        write_grads=False\n",
    "    )\n",
    "\n",
    "    history = autoencoder.fit(\n",
    "        X_train, \n",
    "        X_train,\n",
    "        epochs=epochs,\n",
    "        verbose=verbose,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        validation_data=(X_test, X_test),\n",
    "        callbacks=[\n",
    "            tensorboard,\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "    predictions = autoencoder.predict(X_test[:1])\n",
    "    plt.plot(np.reshape(X_test[:1], (X_test.shape[1])))\n",
    "    plt.plot(np.reshape(predictions, (predictions.shape[1])))\n",
    "    \n",
    "train_autoencoder()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def create_autoencoder_multivariate(\n",
    "        input_shape,\n",
    "        code_len,\n",
    "        classes_num,\n",
    "        name='autoencoder_multivariate',\n",
    "        units=512,\n",
    "):\n",
    "    adversarial_model_loss = wasserstein_discriminator_loss\n",
    "    generator_loss = wasserstein_generator_loss\n",
    "        \n",
    "    # Create and compile model\n",
    "    autoencoder, _, _, _, _, _, _ = create_adversarial_autoencoder(\n",
    "        input_shape,\n",
    "        classes_num=classes_num,\n",
    "        code_len=code_len,\n",
    "        adversarial_model_loss=adversarial_model_loss,\n",
    "        generator_loss=generator_loss,\n",
    "        units=units,\n",
    "        autoencoder_optimizer=optimizers.Adam(lr=1e-4),\n",
    "    )\n",
    "    \n",
    "    return autoencoder\n",
    "    \n",
    "    \n",
    "def train_autoencoder_multivariate(epochs=10, verbose=True, batchsize=50, code_len=50, test_size=0.33, include_test=False):\n",
    "    \n",
    "    dataset, X_train, y_train, X_test, y_test, classes_num = prepare_sussex_data(batchsize, test_size=test_size, include_test=include_test)\n",
    "    batchsize = dataset.batch_size\n",
    "    steps = dataset.steps\n",
    "    labeled_data_num = dataset.get_num_labeled_data()\n",
    "    input_shape = X_train.shape[1:]\n",
    "    \n",
    "    display(input_shape)\n",
    "    \n",
    "    autoencoder = create_autoencoder_multivariate(input_shape, code_len=code_len, classes_num=classes_num)\n",
    "    plot_model_recursive(autoencoder)\n",
    "    \n",
    "    tensorboard = tf.keras.callbacks.TensorBoard(\n",
    "        log_dir='./tb_logs/cnn_autoencoder_multivariate-' + datetime.datetime.today().strftime('%Y%m%dT%H%M'),\n",
    "        histogram_freq=False,\n",
    "        # histogram_freq=20,\n",
    "        write_graph=True,\n",
    "        # TODO: this option might affect logs size\n",
    "        write_grads=False\n",
    "    )\n",
    "\n",
    "    history = autoencoder.fit(\n",
    "        X_train, \n",
    "        X_train,\n",
    "        epochs=epochs,\n",
    "        verbose=verbose,\n",
    "        batch_size=batchsize,\n",
    "        shuffle=True,\n",
    "        validation_data=(X_test, X_test),\n",
    "        callbacks=[\n",
    "            tensorboard,\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "    predictions = autoencoder.predict(X_test[:1])\n",
    "    for axis in range(predictions.shape[2]):\n",
    "        plt.plot(np.reshape(X_test[:1, :, axis:axis+1], (X_test.shape[1])))\n",
    "        plt.plot(np.reshape(predictions[:, :, axis:axis+1], (predictions.shape[1])))\n",
    "        plt.show()\n",
    "        \n",
    "    return autoencoder\n",
    "    \n",
    "train_autoencoder_multivariate()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def create_supervised_cnn_classifier(\n",
    "        input_shape,\n",
    "        code_len,\n",
    "        classes_num,\n",
    "        name='supervised_cnn_classifier',\n",
    "):\n",
    "    generator = create_encoder(input_shape=input_shape, code_len=code_len)\n",
    "    inputs = generator.inputs\n",
    "        \n",
    "    _, classifier = create_generator_linear_softmax(inputs, generator, classes_num=classes_num, code_len=code_len, name=name)\n",
    "        \n",
    "    classifier.compile(loss='categorical_crossentropy',\n",
    "              optimizer=optimizers.Adam(),\n",
    "              metrics=['accuracy'])\n",
    "    \n",
    "    return classifier\n",
    "\n",
    "def train_supervised_cnn_classifier():\n",
    "    nb_epochs = 1000\n",
    "    nb_classes = 3\n",
    "    \n",
    "    fname = 'ArrowHead'\n",
    "    input_shape, X_train, y_train, X_test, y_test = prepare_data(fname)\n",
    "    \n",
    "    y_train = np_utils.to_categorical(y_train, nb_classes)\n",
    "    y_test = np_utils.to_categorical(y_test, nb_classes)\n",
    "    \n",
    "    batch_size = int(min(X_train.shape[0]/10, 16))\n",
    "    \n",
    "    model = create_supervised_cnn_classifier(input_shape, code_len=36, classes_num=nb_classes)\n",
    "    plot_model_recursive(model)\n",
    "    \n",
    "    reduce_lr = ReduceLROnPlateau(monitor = 'loss', factor=0.5,\n",
    "                  patience=50, min_lr=0.0001) \n",
    "    tb_callback = TensorBoard(\n",
    "                log_dir='./tb_logs/supervised_cnn_classifier-' + fname + '-' + datetime.datetime.today().strftime('%Y%m%dT%H%M'),\n",
    "                histogram_freq=100,\n",
    "                write_graph=True,\n",
    "                # TODO: this option might affect logs size\n",
    "                write_grads=False)\n",
    "    hist = model.fit(X_train, y_train, batch_size=batch_size, epochs=nb_epochs,\n",
    "              verbose=1, validation_data=(X_test, y_test), callbacks=[reduce_lr, tb_callback])\n",
    "    #Print the testing results which has the lowest training loss.\n",
    "    log = pd.DataFrame(hist.history)\n",
    "    print(log.loc[log['loss'].idxmin]['loss'], log.loc[log['loss'].idxmin]['val_acc'])\n",
    "    K.clear_session()\n",
    "    \n",
    "train_supervised_cnn_classifier()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def create_supervised_cnn_classifier2(\n",
    "        input_shape,\n",
    "        code_len,\n",
    "        classes_num,\n",
    "        name='supervised_cnn_classifier',\n",
    "):\n",
    "    adversarial_model_loss = wasserstein_discriminator_loss\n",
    "    generator_loss = wasserstein_generator_loss\n",
    "        \n",
    "    # Create and compile model\n",
    "    _, _, _, classifier, _, _, _ = create_adversarial_autoencoder(\n",
    "        input_shape,\n",
    "        classes_num=classes_num,\n",
    "        code_len=code_len,\n",
    "        adversarial_model_loss=adversarial_model_loss,\n",
    "        generator_loss=generator_loss,\n",
    "    )\n",
    "    \n",
    "    return classifier\n",
    "\n",
    "def train_supervised_cnn_classifier2():\n",
    "    nb_epochs = 1000\n",
    "    nb_classes = 3\n",
    "    \n",
    "    fname = 'ArrowHead'\n",
    "    input_shape, X_train, y_train, X_test, y_test = prepare_data(fname)\n",
    "    \n",
    "    y_train = np_utils.to_categorical(y_train, nb_classes)\n",
    "    y_test = np_utils.to_categorical(y_test, nb_classes)\n",
    "    \n",
    "    batch_size = int(min(X_train.shape[0]/10, 16))\n",
    "    \n",
    "    model = create_supervised_cnn_classifier2(input_shape, code_len=36, classes_num=nb_classes)\n",
    "    plot_model_recursive(model)\n",
    "    \n",
    "    reduce_lr = ReduceLROnPlateau(monitor = 'loss', factor=0.5,\n",
    "                  patience=50, min_lr=0.0001) \n",
    "    tb_callback = TensorBoard(\n",
    "                log_dir='./tb_logs/supervised_cnn_classifier2-' + fname + '-' + datetime.datetime.today().strftime('%Y%m%dT%H%M'),\n",
    "                histogram_freq=100,\n",
    "                write_graph=True,\n",
    "                # TODO: this option might affect logs size\n",
    "                write_grads=False)\n",
    "    hist = model.fit(X_train, y_train, batch_size=batch_size, epochs=nb_epochs,\n",
    "              verbose=1, validation_data=(X_test, y_test), callbacks=[reduce_lr, tb_callback])\n",
    "    #Print the testing results which has the lowest training loss.\n",
    "    log = pd.DataFrame(hist.history)\n",
    "    print(log.loc[log['loss'].idxmin]['loss'], log.loc[log['loss'].idxmin]['val_acc'])\n",
    "    K.clear_session()\n",
    "    \n",
    "train_supervised_cnn_classifier2()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "def test_sussex(\n",
    "        generator_softmax_path='saved_models/sussex_aae_generator_softmax-47.hdf5',\n",
    "        epochs=10, verbose=True, batchsize=50, code_len=50, test_size=0.66, include_test=True\n",
    "):\n",
    "    x_train, y_train, x_test, classes_num = prepare_sussex_test_data()\n",
    "\n",
    "    input_shape = x_train.shape[1:]\n",
    "    adversarial_model_loss = wasserstein_discriminator_loss\n",
    "    generator_loss = wasserstein_generator_loss\n",
    "        \n",
    "    # Create and compile model\n",
    "#     _, _, _, classifier, _, _, _ = create_adversarial_autoencoder(\n",
    "#         input_shape,\n",
    "#         classes_num=classes_num,\n",
    "#         code_len=code_len,\n",
    "#         adversarial_model_loss=adversarial_model_loss,\n",
    "#         generator_loss=generator_loss,\n",
    "#         units=512\n",
    "#     )\n",
    "    \n",
    "#     classifier.load_weights(generator_softmax_path)\n",
    "    # FIXME: https://github.com/tensorflow/tensorflow/issues/27769\n",
    "    classifier = tf.keras.models.load_model(generator_softmax_path)\n",
    "    \n",
    "    metric_func = lambda y_pred, y_true: f1_score(y_pred, y_true, average='macro')\n",
    "#     display(AaeClustering.evaluate(x_train, y_train, classifier, metric_func=metric_func))\n",
    "    #AaeClustering.evaluate(x_test, y_test, classifier, metric_func=metric_func)\n",
    "    \n",
    "    w_pred = classifier.predict(x_test)\n",
    "    y_pred = AaeClustering.get_labels_for_softmax(w_pred)\n",
    "    display(x_test.shape)\n",
    "    display(y_pred.shape)\n",
    "    np.savetxt('data/DB_prediction.txt',\n",
    "               np.repeat(np.reshape(y_pred, (y_pred.shape[0],1)), x_test.shape[1], axis=1),\n",
    "               fmt='%i',\n",
    "               delimiter=\" \")\n",
    "    \n",
    "test_sussex()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "%%time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from sklearn import svm, datasets\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.utils.multiclass import unique_labels\n",
    "import matplotlib.colors as mcolors\n",
    "\n",
    "def plot_confusion_matrix(y_true, y_pred, classes,\n",
    "                          labels, dpi=300):\n",
    "    \"\"\"\n",
    "    This function prints and plots the confusion matrix.\n",
    "    Normalization can be applied by setting `normalize=True`.\n",
    "    \"\"\"\n",
    "            \n",
    "    # combine a new colormap\n",
    "    colors = np.vstack(([[1, 1, 1, 1,]], plt.cm.Reds(np.linspace(0,1.,127)), plt.cm.Greens(np.linspace(0,1.,128))))\n",
    "    cmap = mcolors.LinearSegmentedColormap.from_list('confusion_matrix_colormap', colors)\n",
    "\n",
    "    # Compute confusion matrix\n",
    "    cm = confusion_matrix(y_true, y_pred)\n",
    "    # Only use the labels that appear in the data\n",
    "    classes[np.where(classes == unique_labels(y_true, y_pred))]\n",
    "    \n",
    "    cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] * 100.0\n",
    "    print(\"Normalized confusion matrix\")\n",
    "\n",
    "    print(cm)\n",
    "\n",
    "    fig, ax = plt.subplots(dpi=dpi)\n",
    "    im = ax.imshow(cm, interpolation='nearest', cmap=cmap)\n",
    "    #ax.figure.colorbar(im, ax=ax)\n",
    "    # We want to show all ticks...\n",
    "    ax.set(xticks=np.arange(cm.shape[1]),\n",
    "           yticks=np.arange(cm.shape[0]),\n",
    "           # ... and label them with the respective list entries\n",
    "           xticklabels=labels, yticklabels=labels,\n",
    "           ylabel='True label',\n",
    "           xlabel='Predicted label')\n",
    "\n",
    "    # Rotate the tick labels and set their alignment.\n",
    "    plt.setp(ax.get_xticklabels(), rotation=45, ha=\"right\",\n",
    "             rotation_mode=\"anchor\")\n",
    "\n",
    "    # Loop over data dimensions and create text annotations.\n",
    "    fmt = '{0:.0f}%'\n",
    "    # TODO: fix the thresh\n",
    "    thresh = cm.max() / 2.\n",
    "    display(cm.shape, thresh)\n",
    "    for i in range(cm.shape[0]):\n",
    "        for j in range(cm.shape[1]):\n",
    "            if cm[i, j] >= 1.0:\n",
    "                text = fmt.format(cm[i, j])\n",
    "            elif cm[i, j] == 0.0:\n",
    "                text = ''\n",
    "            else:\n",
    "                text = '<1%'\n",
    "            ax.text(j, i, text,\n",
    "                    ha=\"center\", va=\"center\",\n",
    "                    color=\"white\" if cm[i, j] > thresh else \"black\")\n",
    "    fig.tight_layout()\n",
    "    \n",
    "    print('F1-score = %s' % f1_score(y_true, y_pred, average='macro'))\n",
    "    \n",
    "    plt.savefig('plots/confusion_matrix.png', dpi=dpi)\n",
    "    \n",
    "    plt.show()\n",
    "    return ax\n",
    "\n",
    "def validate_sussex(\n",
    "        generator_softmax_path='saved_models/sussex_aae_generator_softmax-47.hdf5',\n",
    "        test_size=0.66,\n",
    "        batchsize=50,\n",
    "):\n",
    "    # x_train, y_train, x_test, classes_num = prepare_sussex_test_data()\n",
    "    _, _, _, x_test, y_test, _ = prepare_sussex_data(batchsize, test_size=test_size, include_test=False)\n",
    "\n",
    "    classifier = tf.keras.models.load_model(generator_softmax_path)\n",
    "        \n",
    "    w_pred = classifier.predict(x_test)\n",
    "    y_pred = AaeClustering.get_labels_for_softmax(w_pred).astype(int)\n",
    "    \n",
    "    y_true = y_test\n",
    "    \n",
    "    display(y_pred.shape)\n",
    "    display(y_true.shape)\n",
    "\n",
    "    classes = np.array([0, 1, 2, 3, 4, 5, 6, 7])\n",
    "    plot_confusion_matrix(y_true, y_pred, classes, labels=['1', '2', '3', '4', '5', '6', '7', '8']) \n",
    "    \n",
    "validate_sussex()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Best models\n",
    "46 => 0.88\n",
    "\n",
    "50 => 0.9239 \n",
    "\n",
    "49 => 0.9208 \n",
    "\n",
    "48 => 0.9170 \n",
    "\n",
    "47 => 0.9320 (best model !!!)\n",
    "\n",
    "45 => 0.9255 \n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "cnn_baseline.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}