# 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 logging
import math
import mxnet as mx
import numpy as np
from .base import STNAccuracy
from .ctc.ctc_loss import CTCLoss as C_CTCLoss
class CTCSTNAccuracy(STNAccuracy):

    def __init__(self, label_width, num_timesteps, blank_label=0):
        self.label_width = label_width
        self.num_timesteps = num_timesteps
        self.blank_label = blank_label
        super(CTCSTNAccuracy, self).__init__()

    def accuracy(self, labels, preds):
        labels = labels.reshape(self.label_width, -1)
        labels = labels.transpose(1, 0)
        _, num_classes = preds.shape
        preds = preds.reshape(self.num_timesteps, -1, num_classes)
        preds = np.transpose(preds, (1, 0, 2))

        data_iter = zip(
            np.split(labels, len(labels), axis=0),
            np.split(preds, len(preds), axis=0),
        )

        hits = 0
        for label, prediction in data_iter:
            label = np.squeeze(label, axis=0)
            prediction = np.squeeze(prediction, axis=0)

            label = remove_blank(label, blank_label=self.blank_label)
            pred = np.argmax(prediction, axis=1)
            pred = strip_prediction(pred, blank_label=self.blank_label)

            if len(label) == len(pred):

# 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/ctc_metrics.py
import logging
import math
import mxnet as mx
import numpy as np
from .base import STNAccuracy
from .ctc.ctc_loss import CTCLoss as C_CTCLoss

class CTCSTNAccuracy(STNAccuracy):

    def __init__(self, label_width, num_timesteps, blank_label=0):
        self.label_width = label_width
        self.num_timesteps = num_timesteps
        self.blank_label = blank_label
        super(CTCSTNAccuracy, self).__init__()

    def accuracy(self, labels, preds):
        labels = labels.reshape(self.label_width, -1)
        labels = labels.transpose(1, 0)
        _, num_classes = preds.shape
        preds = preds.reshape(self.num_timesteps, -1, num_classes)
        preds = np.transpose(preds, (1, 0, 2))

        data_iter = zip(
            np.split(labels, len(labels), axis=0),
            np.split(preds, len(preds), axis=0),
        )

        hits = 0
        for label, prediction in data_iter:
            label = np.squeeze(label, axis=0)
            prediction = np.squeeze(prediction, axis=0)

            label = remove_blank(label, blank_label=self.blank_label)
            pred = np.argmax(prediction, axis=1)
            pred = strip_prediction(pred, blank_label=self.blank_label)

            if len(label) == len(pred):

                match = True