#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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.
#
# Part of the code originally from Tensorflow


import gzip

import numpy

from bigdl.dataset import base

SOURCE_URL = 'http://yann.lecun.com/exdb/mnist/'

TRAIN_MEAN = 0.13066047740239506 * 255
TRAIN_STD = 0.3081078 * 255
TEST_MEAN = 0.13251460696903547 * 255
TEST_STD = 0.31048024 * 255


def _read32(bytestream):
    dt = numpy.dtype(numpy.uint32).newbyteorder('>')
    return numpy.frombuffer(bytestream.read(4), dtype=dt)[0]


def extract_images(f):
    """Extract the images into a 4D uint8 numpy array [index, y, x, depth].


    Args:
      f: A file object that can be passed into a gzip reader.


    Returns:
      data: A 4D unit8 numpy array [index, y, x, depth].


    Raises:
      ValueError: If the bytestream does not start with 2051.


    """
    print('Extracting', f.name)
    with gzip.GzipFile(fileobj=f) as bytestream:
        magic = _read32(bytestream)
        if magic != 2051:
            raise ValueError(
                'Invalid magic number %d in MNIST image file: %s' %
                (magic, f.name))
        num_images = _read32(bytestream)
        rows = _read32(bytestream)
        cols = _read32(bytestream)
        buf = bytestream.read(rows * cols * num_images)
        data = numpy.frombuffer(buf, dtype=numpy.uint8)
        data = data.reshape(num_images, rows, cols, 1)
        return data


def extract_labels(f):
    print('Extracting', f.name)
    with gzip.GzipFile(fileobj=f) as bytestream:
        magic = _read32(bytestream)
        if magic != 2049:
            raise ValueError(
                'Invalid magic number %d in MNIST label file: %s' %
                (magic, f.name))
        num_items = _read32(bytestream)
        buf = bytestream.read(num_items)
        labels = numpy.frombuffer(buf, dtype=numpy.uint8)
        return labels


def read_data_sets(train_dir, data_type="train"):
    """
    Parse or download mnist data if train_dir is empty.

    :param train_dir: The directory storing the mnist data
    :param data_type: Reading training set or testing set.It can be either "train" or "test"
    :return: (ndarray, ndarray) representing (features, labels)
            features is a 4D unit8 numpy array [index, y, x, depth]
            representing each pixel valued from 0 to 255. labels
            is 1D unit8 nunpy array representing the label valued
             from 0 to 9.
    """
    TRAIN_IMAGES = 'train-images-idx3-ubyte.gz'
    TRAIN_LABELS = 'train-labels-idx1-ubyte.gz'
    TEST_IMAGES = 't10k-images-idx3-ubyte.gz'
    TEST_LABELS = 't10k-labels-idx1-ubyte.gz'

    if data_type == "train":
        local_file = base.maybe_download(TRAIN_IMAGES, train_dir,
                                         SOURCE_URL + TRAIN_IMAGES)
        with open(local_file, 'rb') as f:
            train_images = extract_images(f)

        local_file = base.maybe_download(TRAIN_LABELS, train_dir,
                                         SOURCE_URL + TRAIN_LABELS)
        with open(local_file, 'rb') as f:
            train_labels = extract_labels(f)
        return train_images, train_labels

    else:
        local_file = base.maybe_download(TEST_IMAGES, train_dir,
                                         SOURCE_URL + TEST_IMAGES)
        with open(local_file, 'rb') as f:
            test_images = extract_images(f)

        local_file = base.maybe_download(TEST_LABELS, train_dir,
                                         SOURCE_URL + TEST_LABELS)
        with open(local_file, 'rb') as f:
            test_labels = extract_labels(f)
        return test_images, test_labels


if __name__ == "__main__":
    train, _ = read_data_sets("/tmp/mnist/", "train")
    test, _ = read_data_sets("/tmp/mnist", "test")
    assert numpy.abs(numpy.mean(train) - TRAIN_MEAN) / TRAIN_MEAN < 1e-7
    assert numpy.abs(numpy.std(train) - TRAIN_STD) / TRAIN_STD < 1e-7
    assert numpy.abs(numpy.mean(test) - TEST_MEAN) / TEST_MEAN < 1e-7
    assert numpy.abs(numpy.std(test) - TEST_STD) / TEST_STD < 1e-7

@@ -1116,24 +1116,19 @@
 %5D.%0A%0A
-%0A
     
-Args:%0A     
+:param:
  f: 
@@ -1184,28 +1184,20 @@
 er.%0A
-%0A%0A
     
-R
+:r
 eturn
-s:%0A     
+:
  dat
@@ -1248,27 +1248,19 @@
 h%5D.%0A
-%0A%0A
     
-R
+:r
 aise
-s:%0A     
+:
  Val
@@ -1313,17 +1313,16 @@
  2051.%0A%0A
-%0A
     %22%22%22%0A
@@ -2480,16 +2480,17 @@
   :param
+:
  train_d
@@ -2530,16 +2530,17 @@
 st data%0A
+%0A
     :par
@@ -2541,16 +2541,17 @@
   :param
+:
  data_ty
@@ -2625,16 +2625,17 @@
  %22test%22%0A
+%0A
     :ret
@@ -2638,16 +2638,29 @@
 :return:
+%0A%0A    %60%60%60%0A   
  (ndarra
@@ -2703,24 +2703,16 @@
 labels)%0A
-        
     feat
@@ -2762,28 +2762,16 @@
 , depth%5D
-%0A           
  represe
@@ -2812,27 +2812,19 @@
 255.
+%0A   
  labels
-%0A           
  is 
@@ -2873,29 +2873,16 @@
 l valued
-%0A            
  from 0 
@@ -2887,16 +2887,25 @@
 0 to 9.%0A
+    %60%60%60%0A%0A
     %22%22%22%0A

fix doc-gen (#1151)