import numpy as np
import tensorflow as tf
import itertools
import copy


class Bunch(object):
    def __init__(self, adict):
        self.__dict__.update(adict)


def dict_2_namespace(dic):
    return Bunch(dic)


def next_batch(args):
    t0 = np.random.randn(args.batch_size, 1, 1)
    mixed_noise = np.random.randn(args.batch_size, args.seq_length, 1) * 0.1

    t = (np.arange(args.seq_length)[np.newaxis, :, np.newaxis] / 10. + t0)
    A = np.random.randn(args.batch_size, args.seq_length, 1)
    x = np.sin(2 * np.pi * t) + A * 0.1 + mixed_noise * 0.1

    t = np.arange(1, args.seq_length + 1)[np.newaxis, :, np.newaxis] / 10. + t0
    A = np.random.randn(args.batch_size, args.seq_length, 1)
    y = np.sin(2 * np.pi * (t) + A * 0.1 + mixed_noise * 0.1)

    return x, y


def tf_normal(y, mu, s, args):
    # TODO deleted rho
    ss = tf.maximum(1e-10, tf.square(s))
    norm = tf.subtract(y, mu)
    z = tf.div(tf.square(norm), ss)
    denom_log = tf.log(2 * np.pi * ss, name='denom_log')
    result = tf.reduce_sum(z + denom_log, 1) / 2  # -
    return result


def tf_kl_gaussgauss(mu_1, sigma_1, mu_2, sigma_2):
    return tf.reduce_sum(0.5 * (
        2 * tf.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
        - 2 * tf.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
        + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) /
        tf.maximum(1e-9, (tf.square(sigma_2))) - 1
    ), 1)
    # return tf.reduce_sum(0.5 * (
    #     2 * tf.log(sigma_2, name='log_sigma_2')
    #     - 2 * tf.log(sigma_1, name='log_sigma_1')
    #     + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) /
    #     (tf.square(sigma_2)) - 1
    # ), 1)


def get_lossfunc(enc_mu, enc_sigma, dec_mu, dec_sigma, prior_mu,
                 prior_sigma, y, args):
    kl_loss = tf_kl_gaussgauss(
        enc_mu, enc_sigma, prior_mu, prior_sigma)
    likelihood_loss = tf_normal(y, dec_mu, dec_sigma, args)
    return tf.reduce_mean(kl_loss), tf.reduce_mean(likelihood_loss)


def expand_cv_parameters(params):
    param_pool = [[(key, v) for v in value] for key,
                  value in params.items() if isinstance(value, list)]

    cv_params = []
    for i in list(itertools.product(*param_pool)):
        d = copy.deepcopy(params)
        name = d['name']
        for j in i:
            d[j[0]] = j[1]
            name += "_" + j[0] + "_" + str(j[1])
            d['name'] = name
        cv_params.append(d)
    if not cv_params:
        return [params]*params.num_runs

    gs_params = []
    for p in cv_params:
        gs_params += [p]*p['num_runs']
    return gs_params
