#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: imagenet-resnet-keras.py
# Author: Yuxin Wu

import numpy as np
import tensorflow as tf
import argparse

from tensorpack import InputDesc, SyncMultiGPUTrainerReplicated
from tensorpack.dataflow import FakeData, MapDataComponent
from tensorpack.utils import logger
from tensorpack.utils.gpu import get_nr_gpu
from tensorpack.contrib.keras import KerasModel
from tensorpack.callbacks import *
from tensorflow.python.keras.layers import *

from imagenet_utils import get_imagenet_dataflow, fbresnet_augmentor, ImageNetModel


TOTAL_BATCH_SIZE = 512
BASE_LR = 0.1 * (TOTAL_BATCH_SIZE // 256)


def bn(x, name, zero_init=False):
    return BatchNormalization(
        axis=1, name=name, fused=True,
        momentum=0.9, epsilon=1e-5,
        gamma_initializer='zeros' if zero_init else 'ones')(x)


def conv(x, filters, kernel, strides=1, name=None):
    return Conv2D(filters, kernel, name=name,
                  strides=strides, use_bias=False, padding='same',
                  kernel_initializer=tf.keras.initializers.VarianceScaling(
                      scale=2.0, mode='fan_out', distribution='normal'),
                  kernel_regularizer=tf.keras.regularizers.l2(5e-5))(x)


def identity_block(input_tensor, kernel_size, filters, stage, block):
    filters1, filters2, filters3 = filters
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = conv(input_tensor, filters1, 1, name=conv_name_base + '2a')
    x = bn(x, name=bn_name_base + '2a')
    x = Activation('relu')(x)

    x = conv(x, filters2, kernel_size, name=conv_name_base + '2b')
    x = bn(x, name=bn_name_base + '2b')
    x = Activation('relu')(x)

    x = conv(x, filters3, (1, 1), name=conv_name_base + '2c')
    x = bn(x, name=bn_name_base + '2c', zero_init=True)

    x = tf.keras.layers.add([x, input_tensor])
    x = Activation('relu')(x)
    return x


def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
    filters1, filters2, filters3 = filters
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = conv(input_tensor, filters1, (1, 1), name=conv_name_base + '2a')
    x = bn(x, name=bn_name_base + '2a')
    x = Activation('relu')(x)

    x = conv(x, filters2, kernel_size, strides=strides, name=conv_name_base + '2b')
    x = bn(x, name=bn_name_base + '2b')
    x = Activation('relu')(x)

    x = conv(x, filters3, (1, 1), name=conv_name_base + '2c')
    x = bn(x, name=bn_name_base + '2c', zero_init=True)

    shortcut = conv(
        input_tensor,
        filters3, (1, 1), strides=strides,
        name=conv_name_base + '1')
    shortcut = bn(shortcut, name=bn_name_base + '1')

    x = tf.keras.layers.add([x, shortcut])
    x = Activation('relu')(x)
    return x


def resnet50(image):
    input = Input(tensor=image)

    def image_preprocess(image):
        image = ImageNetModel.image_preprocess(image)
        image = tf.transpose(image, [0, 3, 1, 2])
        return image

    x = Lambda(image_preprocess)(input)

    x = conv(x, 64, (7, 7), strides=(2, 2), name='conv0')
    x = bn(x, name='bn_conv1')
    x = Activation('relu')(x)
    x = MaxPooling2D((3, 3), strides=(2, 2), padding='same')(x)

    x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')

    x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')

    x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')

    x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')

    x = GlobalAveragePooling2D(name='avg_pool')(x)
    x = Flatten()(x)
    x = Dense(1000, activation='softmax', name='fc1000',
              kernel_initializer=tf.keras.initializers.VarianceScaling(
                  scale=2.0, mode='fan_in'),
              kernel_regularizer=tf.keras.regularizers.l2(5e-5))(x)

    M = tf.keras.models.Model(input, x, name='resnet50')
    return M


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', help='ILSVRC dataset dir')
    parser.add_argument('--fake', help='use fakedata to test or benchmark this model', action='store_true')
    args = parser.parse_args()
    logger.set_logger_dir(os.path.join("train_log", "imagenet-resnet-keras"))

    tf.keras.backend.set_image_data_format('channels_first')

    nr_gpu = get_nr_gpu()
    if args.fake:
        df_train = FakeData([[64, 224, 224, 3], [64, 1000]], 5000, random=False, dtype='uint8')
        df_val = FakeData([[64, 224, 224, 3], [64, 1000]], 5000, random=False)
    else:
        batch_size = TOTAL_BATCH_SIZE // nr_gpu
        assert args.data is not None
        df_train = get_imagenet_dataflow(
            args.data, 'train', batch_size, fbresnet_augmentor(True))
        df_val = get_imagenet_dataflow(
            args.data, 'val', batch_size, fbresnet_augmentor(False))

        def one_hot(label):
            return np.eye(1000)[label]

        df_train = MapDataComponent(df_train, one_hot, 1)
        df_val = MapDataComponent(df_val, one_hot, 1)

    M = KerasModel(
        resnet50,
        inputs_desc=[InputDesc(tf.uint8, [None, 224, 224, 3], 'images')],
        targets_desc=[InputDesc(tf.float32, [None, 1000], 'labels')],
        input=df_train,
        trainer=SyncMultiGPUTrainerReplicated(nr_gpu))

    lr = tf.get_variable('learning_rate', initializer=0.1, trainable=False)
    tf.summary.scalar('lr', lr)

    M.compile(
        optimizer=tf.train.MomentumOptimizer(lr, 0.9, use_nesterov=True),
        loss='categorical_crossentropy',
        metrics='categorical_accuracy'
    )

    callbacks = [
        ModelSaver(),
        ScheduledHyperParamSetter(
            'learning_rate',
            [(0, 0.1), (3, BASE_LR)], interp='linear'),  # warmup
        ScheduledHyperParamSetter(
            'learning_rate',
            [(30, BASE_LR * 0.1), (60, BASE_LR * 1e-2), (85, BASE_LR * 1e-3)]),
        GPUUtilizationTracker()
    ]
    if not args.fake:
        callbacks.append(
            DataParallelInferenceRunner(
                df_val, ScalarStats(['categorical_accuracy']), nr_gpu))

    M.fit(
        steps_per_epoch=100 if args.fake else 1281167 // TOTAL_BATCH_SIZE,
        max_epoch=100,
        callbacks=callbacks
    )
