import glob, os
import numpy as np
import cv2
import random


class OCR_data(object):
    def __init__(self, num, data_dir, num_classes, batch_size=50, len_code=5, height=60, width=180, resize_height=24,
                 resize_width=88, num_channels=1):
        self.num = num
        self.data_dir = data_dir
        self.num_classes = num_classes
        self.batch_size = batch_size
        self.len_code = len_code
        self.height = height
        self.width = width
        self.resize_height = resize_height
        self.resize_width = resize_width
        self.num_channels = num_channels
        self.index_in_epoch = 0
        self._imgs = []
        self._labels = []
        for pathAndFilename in glob.iglob(os.path.join(data_dir, '*.png')):
            img, label = self.create_captcha(pathAndFilename)
            self._imgs.append(img)
            self._labels.append(label)
            self._imgs = np.array(self._imgs).reshape((-1, resize_height, resize_width, num_channels)).astype(
                np.float32)
            self._labels = np.array(self._labels)

    def create_captcha(self, pathAndFilename):
        img = cv2.imread(pathAndFilename, cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, (self.resize_width, self.resize_height), interpolation=cv2.INTER_AREA)
        filename, ext = os.path.splitext(os.path.basename(pathAndFilename))
        label = self.create_label(filename)
        return (img, label)

    def create_label(self, filename):
        label = []
        for c in filename:
            ascii_code = ord(c)
            if ascii_code < 58:
                char_value = ascii_code - 48
            else:
                char_value = ascii_code - 87
            label.append(char_value)
        return self.dense_to_one_hot(label, self.num_classes)

    def dense_to_one_hot(self, labels_dense, num_classes):
        num_labels = len(labels_dense)
        index_offest = np.arange(num_labels) * num_classes
        labels_one_hot = np.zeros((num_labels, num_classes))
        labels_one_hot.flat[index_offest + labels_dense] = 1
        labels_one_hot = labels_one_hot.reshape(num_labels * num_classes)
        return labels_one_hot

    def next_batch(self, batch_size):
        start = self.index_in_epoch
        self.index_in_epoch += batch_size
        if self.index_in_epoch > self.num:
            perm = np.arange(self.num)
            np.random.shuffle(perm)
            self._imgs = self._imgs[perm]
            self._labels = self._labels[perm]
            start = 0
            self.index_in_epoch = batch_size
            assert batch_size <= self.num
        end = self.index_in_epoch
        return self._imgs[start:end], self._labels[start:end]
