# Path: mxnet/metrics/base.py
# class STNAccuracy(mx.metric.EvalMetric):
# 
#     def __init__(self, make_label_time_major=True):
#         super(STNAccuracy, self).__init__('Accuracy')
#         self.make_label_time_major = make_label_time_major
# 
#     def accuracy(self, labels, preds):
#         if preds.shape != labels.shape:
#             preds = np.argmax(preds, axis=1)
#         preds = preds.astype(np.int32)
#         labels = labels.astype(np.int32)
# 
#         sum_metric = (preds.flat == labels.flat).sum()
#         return sum_metric, len(preds.flat)
# 
#     def update(self, labels, preds):
#         cls_prob = preds[0].asnumpy()
#         labels = labels[0].asnumpy()
#         cls_labels = labels
# 
#         if len(cls_labels.shape) == 2:
#             # make labels time major
#             if self.make_label_time_major:
#                 cls_labels = np.transpose(cls_labels, axes=(1, 0))
#             cls_labels = cls_labels.ravel()
# 
#         accuracy_sum, num_instances = self.accuracy(cls_labels, cls_prob)
#         self.sum_metric += accuracy_sum
#         self.num_inst += num_instances


import mxnet as mx
import numpy as np
from .base import STNAccuracy

def tile_batch_size(array, size=4):
    array = array[:, np.newaxis, ...]
    array = np.tile(array, (1, size, 1))
    array_shape = array.shape
    return array.reshape(-1, *array_shape[2:])


class FSNSPretrainAccuracy(STNAccuracy):

    def update(self, labels, preds):
        cls_prob = preds[0].asnumpy()
        labels = labels[0].asnumpy()

        # copy the same labels 4 times because we have 4 times the same image (at least we assume that)
        labels = tile_batch_size(labels)

        accuracy_sum, num_instances = self.accuracy(labels, cls_prob)
        self.sum_metric += accuracy_sum
        self.num_inst += num_instances


class FSNSPretrainCrossEntropy(mx.metric.EvalMetric):

    def __init__(self, eps=1e-8):
        super(FSNSPretrainCrossEntropy, self).__init__('cross-entropy')
        self.eps = eps

    def update(self, labels, preds):
        for label, pred in zip(labels, preds):

# Path: mxnet/metrics/base.py
# class STNAccuracy(mx.metric.EvalMetric):
# 
#     def __init__(self, make_label_time_major=True):
#         super(STNAccuracy, self).__init__('Accuracy')
#         self.make_label_time_major = make_label_time_major
# 
#     def accuracy(self, labels, preds):
#         if preds.shape != labels.shape:
#             preds = np.argmax(preds, axis=1)
#         preds = preds.astype(np.int32)
#         labels = labels.astype(np.int32)
# 
#         sum_metric = (preds.flat == labels.flat).sum()
#         return sum_metric, len(preds.flat)
# 
#     def update(self, labels, preds):
#         cls_prob = preds[0].asnumpy()
#         labels = labels[0].asnumpy()
#         cls_labels = labels
# 
#         if len(cls_labels.shape) == 2:
#             # make labels time major
#             if self.make_label_time_major:
#                 cls_labels = np.transpose(cls_labels, axes=(1, 0))
#             cls_labels = cls_labels.ravel()
# 
#         accuracy_sum, num_instances = self.accuracy(cls_labels, cls_prob)
#         self.sum_metric += accuracy_sum
#         self.num_inst += num_instances


# Path: mxnet/metrics/fsns_metrics.py
import mxnet as mx
import numpy as np
from .base import STNAccuracy


def tile_batch_size(array, size=4):
    array = array[:, np.newaxis, ...]
    array = np.tile(array, (1, size, 1))
    array_shape = array.shape
    return array.reshape(-1, *array_shape[2:])


class FSNSPretrainAccuracy(STNAccuracy):

    def update(self, labels, preds):
        cls_prob = preds[0].asnumpy()
        labels = labels[0].asnumpy()

        # copy the same labels 4 times because we have 4 times the same image (at least we assume that)
        labels = tile_batch_size(labels)

        accuracy_sum, num_instances = self.accuracy(labels, cls_prob)
        self.sum_metric += accuracy_sum
        self.num_inst += num_instances


class FSNSPretrainCrossEntropy(mx.metric.EvalMetric):

    def __init__(self, eps=1e-8):
        super(FSNSPretrainCrossEntropy, self).__init__('cross-entropy')
        self.eps = eps

    def update(self, labels, preds):
        for label, pred in zip(labels, preds):

            label = label.asnumpy()