#!/usr/bin/env python
import argparse
from tensorflow.compat import v1 as tf
import tqdm

def split_grad_list(grad_list):
    g = []
    v = []
    for tower in grad_list:
        g.append([x[0] for x in tower])
        v.append([x[1] for x in tower])
    return g, v

def allreduce_grads(all_grads):
    # reduce gradients for N variables on K devices
    from tensorflow.python.ops import nccl_ops as nccl
    nr_tower = len(all_grads)
    assert nr_tower > 1
    new_all_grads = []  # N x K
    for grads in zip(*all_grads):
        # k grads
        summed = nccl.all_sum(grads)

        grads_for_devices = []  # K
        true_sum = tf.add_n(grads)
        for g in summed:
            diff = tf.abs(true_sum - g)
            eql = diff < 1e-4
            nccl_res_correct = tf.reduce_all(eql, name="corr_" + grads[0].op.name)

            def flat(x):
                x = tf.reshape(x, [-1])
                x = tf.slice(x, [0], [tf.minimum(tf.size(x), 200)])
                return x

            assert_op = tf.debugging.Assert(nccl_res_correct, [
                tf.reduce_max(diff), flat(true_sum), flat(g)], summarize=1000,
                name='assert_' + grads[0].op.name)
            with tf.control_dependencies([assert_op]):
                g = tf.identity(g)
            grads_for_devices.append(g)
        new_all_grads.append(grads_for_devices)
    # transpose to K x N
    ret = list(zip(*new_all_grads))
    return ret

def build_graph(image, label, idx):
    v1 = tf.get_variable('aaa/W', shape=[3, 3, 3, 64], trainable=True)
    v2 = tf.get_variable('bbb/W', shape=[3, 3, 3, 64], trainable=True)
    v = v1 if idx == 0 else v2
    image = tf.nn.conv2d(image, v, 1, padding='SAME', data_format='NCHW')

    def conv(name, x, chan, stride=1):
        with tf.variable_scope(name):
            in_chan = x.shape[1]
            W = tf.get_variable('W', [3, 3, in_chan, chan])
            ret = tf.nn.conv2d(x, W, strides=stride, padding="SAME", data_format="NCHW")
            return tf.nn.relu(ret)

    x = conv('conv1', image, 64)
    x = conv('conv2', x, 64)
    x = conv('conv3', x, 1280, stride=2)
    x = conv('conv4', x, 1280, stride=2)
    x = conv('conv5', x, 10)
    logits = tf.reduce_mean(x, axis=[2, 3])
    cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=label)
    cost = tf.reduce_mean(cost, name='cross_entropy_loss')
    return cost


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', type=int)
    args = parser.parse_args()
    num_gpu = args.gpu

    with tf.Graph().as_default():
        opt = tf.train.GradientDescentOptimizer(0.001)

        grad_list = []
        for k in range(num_gpu):
            with tf.device("/gpu:{}".format(k)), tf.variable_scope("tower{}".format(k)):
                print("Building {} ...".format(k))
                image = tf.random.uniform([32, 3, 30, 30])
                label = tf.random.uniform([32], maxval=9, dtype=tf.int32)
                cost = build_graph(image, label, k)
                varlist = [x for x in tf.trainable_variables() if x.name.startswith("tower{}".format(k))]
                print("Varlist for tower {}: ".format(k), [x.name for x in varlist])
                wd_cost = [tf.reduce_sum(x) * 1e-3 for x in varlist]
                cost = tf.add_n([cost] + wd_cost)
                grads = opt.compute_gradients(cost, var_list=varlist)
                grad_list.append(grads)

        all_grads, all_vars = split_grad_list(grad_list)
        all_grads = allreduce_grads(all_grads)
        grad_list = [list(zip(gs, vs)) for gs, vs in zip(all_grads, all_vars)]

        train_ops = []
        for idx, grad_and_vars in enumerate(grad_list):
            with tf.device('/gpu:{}'.format(idx)):
                train_ops.append(opt.apply_gradients(
                    grad_and_vars, name='apply_grad_{}'.format(idx)))
        train_op = tf.group(*train_ops)

        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
        print("Training ...")
        for k in tqdm.trange(5000):
            sess.run(train_op)