#!/usr/bin/python

""" tests.py: Quantitative and qualitative tests on the effectiveness of the
        VAE implementation

    TODO:
        - Use nosetests-like framework
        - Auto-save to graphs/
        - Click CLI
"""

import tensorflow as tf
import click as cl
import numpy as np
import yaml

import vae
import vis

FCONFIG = 'config/nn_config.yaml'

try:
    config = yaml.load(file(FCONFIG, 'r'))
except yaml.YAMLError, exc:
    cl.secho("Error in configuration file: {}".format(exc), fg='red')

ARCH   = config["architecture"]
DIMS   = config["dims"]
OPT    = config["optimization"]
PARAMS = config["AEVB"]
TRAIN  = config["training"]

# TF backend initializations
tf.set_random_seed(0)


def simple_test():
    '''
    Train VAE using default parameters specified in config/nn_config.yaml and
    plot progress over training iterations
    '''
    # Instantiate and train vanilla autoencoder
    nn = vae.VAE(ARCH, DIMS, OPT, PARAMS, TRAIN)
    results = nn.train()
    iters = results["iters"]
    ELBOs = results["ELBO"]
    KLs = results["KL"]
    LLs = results["LL"]

    # Graph training results
    titles = ["$\mathcal{L}(\phi,\\theta;x)$",
              "$KL(q_{\phi}(z|x)||p_{\\theta}(z))$",
              "$\log(p_{\\theta}(x|z))$"]
    params = {"$\eta_{%s}$" % OPT["type"]: OPT["{}_rate".format(OPT["type"])],
              "$Activation$": ARCH["activation"],
              "$Batch$ $Size$": TRAIN["batch_size"],
              "$MCE$ $Samples$": PARAMS["L"],
              "$Latent$ $Dim$": DIMS["data"]}
    vis.basic_multiplot([iters]*3, [[ELBOs], [KLs], [LLs]], titles,
        show_legend=False, params=params) 

def bivariate_latent_space():
    '''
    Train VAE using default parameters specified in config/nn_config.yaml and
    plot examples of latent space distribution generated by each data point
    '''
    # Instantiate and train vanilla autoencoder
    nn = vae.VAE(ARCH, DIMS, OPT, PARAMS, TRAIN)
    nn.train()

    # Do single forward pass on the encoder
    new_data = vae.mnist.test.next_batch(TRAIN["batch_size"])[0]
    mu_q, stddev_q, rho_p = nn.full_fp(new_data)

    vis.full_pass_vis(
        [im.reshape(28,28) for im in new_data[:3]],
        [im.reshape(28,28) for im in rho_p[:3]],
        mu_q[:3], stddev_q[:3])


def multivariate_latent_space_PCA():
    '''
    Do same as bivariate latent space visualization except perform PCA
    to project onto lower dimension before visualizing
    '''
    pass


def reconstruction_test():
    '''
    Reconstruct samples after training VAE
    '''
    # Instantiate and train vanilla autoencoder
    nn = vae.VAE(ARCH, DIMS, OPT, PARAMS, TRAIN)
    nn.train()

    # Do single forward pass on the encoder
    new_data = vae.mnist.test.next_batch(TRAIN["batch_size"])[0]
    mu_q, stddev_q, rho_p = nn.full_fp(new_data)

    vis.juxtapose_images(
        [im.reshape(28,28) for im in new_data[:5]],
        [im.reshape(28,28) for im in rho_p[:5]])


if __name__ == "__main__":
    #simple_test()
    #bivariate_latent_space()
    reconstruction_test()
