#!/usr/bin/env python
'''
Wavenet model definition and generation using contrib.layers
 - Carl Quillen
'''

from __future__ import print_function

import tensorflow as tf
import tensorflow.contrib.layers as layers
from tensorflow.contrib.framework import arg_scope


def wavenet_block(xpad, x, conditioning, num_outputs,
                  num_outputs2, rate, is_training, opts, scope):
    '''
    wavenet_block: many important convolution parameters (reuse, kernel_size
    etc.) come from the arg_scope() and are set by wavenet().

    Note that convolutions below end up being causal because the
    input is padded and padding='VALID'. This causes samples to
    drop in the "right" way.
    '''

    dropout, kernel_size, skip_dimension, histogram_summaries = (
        opts.dropout, opts.kernel_size, opts.skip_dimension,
        opts.histogram_summaries)

    conv_gate = layers.conv2d(
        xpad, num_outputs=num_outputs2*2, rate=rate, kernel_size=kernel_size,
        activation_fn=None, normalizer_params=None, scope=scope + '/conv_gate')

    # Add the conditioning.
    conv_gate += layers.conv2d(
        conditioning, num_outputs=num_outputs2*2, rate=rate, kernel_size=1,
        activation_fn=None, normalizer_params=None, scope=scope + '/cur_cond')

    with tf.name_scope(scope + '/activation'):
        conv = tf.nn.tanh(conv_gate[:, :, :num_outputs2], name='conv')
        gate = tf.nn.sigmoid(conv_gate[:, :, num_outputs2:], name='gate')

    with tf.name_scope(scope + '/prod'):
        out = conv * gate

    if dropout > 0:
        out = layers.dropout(out, keep_prob=dropout,
                             is_training=is_training, scope=scope + '/dropout')

    out = layers.conv2d(out, num_outputs=num_outputs, kernel_size=1,
                        activation_fn=None,
                        scope=scope + '/output_xform')
    with tf.name_scope(scope + '/residual'):
        residual = x + out

    if skip_dimension != num_outputs:      # Upscale for more goodness.
        out = layers.conv2d(out, num_outputs=skip_dimension,
                            kernel_size=1, activation_fn=None,
                            scope=scope + '/skip_upscale')

    if histogram_summaries:
        tf.summary.histogram(name=scope + '/conv', values=conv)
        tf.summary.histogram(name=scope + '/gate', values=gate)
        tf.summary.histogram(name=scope + '/out', values=out)

    return residual, out        # out gets added to the skip connections.


def padded(new_x, pad, scope, n_chunks, reuse=False,
           reverse=False, data_format=None):
    '''
    Pad new_x, and save the rightmost window for context for the next time
    we do the same convolution.  This context carries across utterances
    during training.  Using this trick also allows us to use the same
    wavenet() routine in training as well as generation.

    reverse=True for reversing the direction of causality.
    '''

    with tf.variable_scope(scope, reuse=reuse):

        if data_format is 'NCW':
            x = tf.get_variable(
                'pad', shape=(n_chunks, new_x.get_shape()[1], pad),
                collections=[tf.GraphKeys.GLOBAL_VARIABLES, 'padding'],
                initializer=tf.constant_initializer(), trainable=False)
            if not reverse:
                y = tf.concat(values=(x, new_x), axis=2)
                x = tf.assign(x, y[:, :, -pad:])
            else:
                y = tf.concat(values=(new_x, x), axis=2)
                x = tf.assign(x, y[:, :, :pad])
        else:
            x = tf.get_variable(
                'pad', shape=(n_chunks, pad, new_x.get_shape()[2]),
                collections=[tf.GraphKeys.GLOBAL_VARIABLES, 'padding'],
                initializer=tf.constant_initializer(), trainable=False)
            if not reverse:
                y = tf.concat(values=(x, new_x), axis=1)
                x = tf.assign(x, y[:, -pad:, :])
            else:
                y = tf.concat(values=(new_x, x), axis=1)
                x = tf.assign(x, y[:, :pad, :])
        with tf.get_default_graph().control_dependencies([x]):
            return tf.identity(y)


def wavenet(inputs, opts, is_training=True, reuse=False, pad_reuse=False,
            data_format=None, extra_pad_scope=''):
    '''
    The wavenet model definition for training/generation.
    use wavenets recursively, we will want separate padding variables for
    each "layer".  So we have a separate reuse flag for padding() and
    an additional thing to add to the scope for padding() in that case.
    '''

    # Parameters for batch normalization
    normalizer_params = {}
    if opts.batch_norm:
        normalizer_params = {
            'normalizer_fn': layers.batch_norm,
            'normalizer_params': {
                'is_training': is_training,
                'trainable': False,
                'variables_collections': {
                   'gamma': ['batch_norm']},
                'reuse': reuse,
                'scale': True,  # Update Variance too.
                'scope': 'BatchNorm'
            }
        }

    l2reg = None
    if 'l2reg' in vars(opts):
        l2reg = tf.contrib.layers.l2_regularizer(opts.l2reg)
    # unpack inputs.
    inputs, user, alignment, lf0 = inputs

    with tf.variable_scope('conditioning'):
        conditioning = tf.one_hot(alignment, depth=opts.n_phones,
                                  name='align_onehot')
        conditioning = tf.reshape(
            conditioning, (opts.n_chunks, -1, opts.n_phones*opts.context))
        if user is not None:
            user = tf.one_hot(user, depth=opts.n_users, name='user_onehot')
            conditioning = tf.concat([user, conditioning], axis=2,
                                     name='cat_user')
        lf0 = tf.reshape(lf0, (opts.n_chunks, -1, 1))
        conditioning = tf.concat((conditioning, lf0), axis=2, name='cat_lf0')

    # The arg_scope below will apply to all convolutions, including the ones
    # in wavenet_block().
    with arg_scope([layers.conv2d], data_format=data_format,
                   reuse=reuse, padding='VALID', weights_regularizer=l2reg,
                   **normalizer_params):
        if opts.input_kernel_size > 1:
            inputs = padded(new_x=inputs, reuse=pad_reuse,
                            reverse=opts.reverse, pad=opts.input_kernel_size-1,
                            n_chunks=opts.n_chunks, data_format=data_format,
                            scope='input_layer/pad'+extra_pad_scope)
        x = layers.conv2d(inputs, num_outputs=opts.num_outputs,
                          kernel_size=opts.input_kernel_size, rate=1,
                          activation_fn=tf.nn.tanh, scope='input_layer')

        skip_connections = 0
        for i_block, block_dilations in enumerate(opts.dilations):
            for rate in block_dilations:
                block_rate = "block_{}/rate_{}".format(i_block, rate)
                xpad = padded(
                    new_x=x, pad=rate*(opts.kernel_size-1),
                    reuse=pad_reuse, n_chunks=opts.n_chunks,
                    reverse=opts.reverse, data_format=data_format,
                    scope=block_rate+"/pad"+extra_pad_scope)

                x, skip_connection = wavenet_block(
                    xpad, x, conditioning, opts.num_outputs,
                    opts.num_outputs2, rate, is_training,
                    opts, scope=block_rate)

                with tf.name_scope(block_rate+"_skip".format(i_block, rate)):
                    skip_connections += skip_connection

    with tf.name_scope("relu_skip"):
        skip_connections = tf.nn.relu(skip_connections)

    with arg_scope([layers.conv2d], kernel_size=1, reuse=reuse,
                   data_format=data_format):
        x = layers.conv2d(
            skip_connections, num_outputs=opts.skip_dimension,  # ?
            activation_fn=tf.nn.relu, scope='output_layer1')
        mfcc = layers.conv2d(
            x, num_outputs=opts.skip_dimension,   # ?
            activation_fn=tf.nn.relu, scope='mfcc_layer1')
        x = layers.conv2d(
            x, num_outputs=3,
            normalizer_params=None,
            activation_fn=None, scope='output_layer2')
        mfcc = layers.conv2d(
            mfcc, num_outputs=opts.n_mfcc, normalizer_params=None,
            activation_fn=None, scope='mfcc_layer2')

    with tf.name_scope("unpack_output"):
        mu = x[:, :, 0]
        r = tf.sin(x[:, :, 1]/opts.r_scale)*opts.r_scale
        r = r*r + 1.0
        q = 0.999*r*tf.sin(x[:, :, 2])
    return mu, r, q, mfcc


def compute_overlap(opts):
    total_lost = opts.input_kernel_size-1
    for block_dilations in opts.dilations:
        for rate in block_dilations:
            total_lost += rate*(opts.kernel_size-1)
    return total_lost
