# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Generate MXNet implementation of CapsNet
Reference 1: https://www.cs.toronto.edu/~fritz/absps/transauto6.pdf
Reference 2: https://arxiv.org/pdf/1710.09829.pdf
"""
import os
import re
import gzip
import struct
import numpy as np
import scipy.ndimage as ndi
import mxnet as mx
from capsulelayers import primary_caps, CapsuleLayer

from mxboard import SummaryWriter


def margin_loss(y_true, y_pred):
    loss = y_true * mx.sym.square(mx.sym.maximum(0., 0.9 - y_pred)) +\
        0.5 * (1 - y_true) * mx.sym.square(mx.sym.maximum(0., y_pred - 0.1))
    return mx.sym.mean(data=mx.sym.sum(loss, 1))


def capsnet(batch_size, n_class, num_routing, recon_loss_weight):
    """Create CapsNet"""
    # data.shape = [batch_size, 1, 28, 28]
    data = mx.sym.Variable('data')

    input_shape = (1, 28, 28)
    # Conv2D layer
    # net.shape = [batch_size, 256, 20, 20]
    conv1 = mx.sym.Convolution(data=data,
                               num_filter=256,
                               kernel=(9, 9),
                               layout='NCHW',
                               name='conv1')
    conv1 = mx.sym.Activation(data=conv1, act_type='relu', name='conv1_act')
    # net.shape = [batch_size, 256, 6, 6]

    primarycaps = primary_caps(data=conv1,
                               dim_vector=8,
                               n_channels=32,
                               kernel=(9, 9),
                               strides=[2, 2],
                               name='primarycaps')
    primarycaps.infer_shape(data=(batch_size, 1, 28, 28))
    # CapsuleLayer
    kernel_initializer = mx.init.Xavier(rnd_type='uniform', factor_type='avg', magnitude=3)
    bias_initializer = mx.init.Zero()
    digitcaps = CapsuleLayer(num_capsule=10,
                             dim_vector=16,
                             batch_size=batch_size,
                             kernel_initializer=kernel_initializer,
                             bias_initializer=bias_initializer,
                             num_routing=num_routing)(primarycaps)

    # out_caps : (batch_size, 10)
    out_caps = mx.sym.sqrt(data=mx.sym.sum(mx.sym.square(digitcaps), 2))
    out_caps.infer_shape(data=(batch_size, 1, 28, 28))

    y = mx.sym.Variable('softmax_label', shape=(batch_size,))
    y_onehot = mx.sym.one_hot(y, n_class)
    y_reshaped = mx.sym.Reshape(data=y_onehot, shape=(batch_size, -4, n_class, -1))
    y_reshaped.infer_shape(softmax_label=(batch_size,))

    # inputs_masked : (batch_size, 16)
    inputs_masked = mx.sym.linalg_gemm2(y_reshaped, digitcaps, transpose_a=True)
    inputs_masked = mx.sym.Reshape(data=inputs_masked, shape=(-3, 0))
    x_recon = mx.sym.FullyConnected(data=inputs_masked, num_hidden=512, name='x_recon')
    x_recon = mx.sym.Activation(data=x_recon, act_type='relu', name='x_recon_act')
    x_recon = mx.sym.FullyConnected(data=x_recon, num_hidden=1024, name='x_recon2')
    x_recon = mx.sym.Activation(data=x_recon, act_type='relu', name='x_recon_act2')
    x_recon = mx.sym.FullyConnected(data=x_recon, num_hidden=np.prod(input_shape), name='x_recon3')
    x_recon = mx.sym.Activation(data=x_recon, act_type='sigmoid', name='x_recon_act3')

    data_flatten = mx.sym.flatten(data=data)
    squared_error = mx.sym.square(x_recon-data_flatten)
    recon_error = mx.sym.mean(squared_error)
    recon_error_stopped = recon_error
    recon_error_stopped = mx.sym.BlockGrad(recon_error_stopped)
    loss = mx.symbol.MakeLoss((1-recon_loss_weight)*margin_loss(y_onehot, out_caps)+recon_loss_weight*recon_error)

    out_caps_blocked = out_caps
    out_caps_blocked = mx.sym.BlockGrad(out_caps_blocked)
    return mx.sym.Group([out_caps_blocked, loss, recon_error_stopped])


def download_data(url, force_download=False):
    fname = url.split("/")[-1]
    if force_download or not os.path.exists(fname):
        mx.test_utils.download(url, fname)
    return fname


def read_data(label_url, image_url):
    with gzip.open(download_data(label_url)) as flbl:
        magic, num = struct.unpack(">II", flbl.read(8))
        label = np.fromstring(flbl.read(), dtype=np.int8)
    with gzip.open(download_data(image_url), 'rb') as fimg:
        magic, num, rows, cols = struct.unpack(">IIII", fimg.read(16))
        image = np.fromstring(fimg.read(), dtype=np.uint8)
        np.reshape(image, len(label), (rows, cols))
    return label, image


def to4d(img):
    return img.reshape(img.shape[0], 1, 28, 28).astype(np.float32)/255


class LossMetric(mx.metric.EvalMetric):
    """Evaluate the loss function"""
    def __init__(self, batch_size, num_gpus):
        super(LossMetric, self).__init__('LossMetric')
        self.batch_size = batch_size
        self.num_gpu = num_gpus
        self.sum_metric = 0
        self.num_inst = 0
        self.loss = 0.0
        self.batch_sum_metric = 0
        self.batch_num_inst = 0
        self.batch_loss = 0.0
        self.recon_loss = 0.0
        self.n_batch = 0

    def update(self, labels, preds):
        """Update the hyper-parameters and loss of CapsNet"""
        batch_sum_metric = 0
        batch_num_inst = 0
        for label, pred_outcaps in zip(labels[0], preds[0]):
            label_np = int(label.asnumpy())
            pred_label = int(np.argmax(pred_outcaps.asnumpy()))
            batch_sum_metric += int(label_np == pred_label)
            batch_num_inst += 1
        batch_loss = preds[1].asnumpy()
        recon_loss = preds[2].asnumpy()
        self.sum_metric += batch_sum_metric
        self.num_inst += batch_num_inst
        self.loss += batch_loss
        self.recon_loss += recon_loss
        self.batch_sum_metric = batch_sum_metric
        self.batch_num_inst = batch_num_inst
        self.batch_loss = batch_loss
        self.n_batch += 1