#!/usr/bin/env python3
'''
Download and format MNIST, CIFAR-10, and derivative datasets.
'''
from os import makedirs
from os.path import join
from shutil import unpack_archive
from tempfile import TemporaryDirectory
from urllib.request import urlretrieve

import numpy as np
import numpy.random as rand
import scipy.io as io
import scipy.misc as misc

################################################################################
# Download MNIST.
################################################################################

def log_mnist_progress(i, chunk_size, file_size):
    print('\rDownloading MNIST — %.2f%% complete.'
          % (100 * i * chunk_size / file_size), end='', flush=True)

mnist_url = 'http://www.cs.nyu.edu/~roweis/data/mnist_all.mat'
mnist_dir = TemporaryDirectory()
mnist_path = mnist_dir.name

urlretrieve(mnist_url, join(mnist_path, 'mnist.mat'), log_mnist_progress)
print(80 * '\b \b' + 'Downloading MNIST — done!')

################################################################################
# Reformat MNIST.
################################################################################

def resize(x, h, w):
    return np.array([
        misc.imresize(x_i[:, :, 0], (h, w), mode='F')[:, :, None]
        for x_i in x])

print('Reformatting MNIST...', end='', flush=True)
mnist_a = io.loadmat(join(mnist_path, 'mnist.mat'))
mnist_b = {
    'm_sym': np.zeros(10),
    'x0_tr': np.vstack([
        resize(np.reshape(mnist_a['train%i' % i], (-1, 28, 28, 1)) / 255,
               32, 32)
        for i in range(10)]),
    'x0_ts': np.vstack([
        resize(np.reshape(mnist_a['test%i' % i], (-1, 28, 28, 1)) / 255,
               32, 32)
        for i in range(10)]),
    'y_tr': np.vstack([
        np.ones((len(mnist_a['train%i' % i]), 1), 'f') * np.identity(10, 'f')[i]
        for i in range(10)]),
    'y_ts': np.vstack([
        np.ones((len(mnist_a['test%i' % i]), 1), 'f') * np.identity(10, 'f')[i]
        for i in range(10)])}

makedirs('data/', exist_ok=True)
np.savez_compressed('data/mnist.npz', mnist_b)
print(80 * '\b \b' + 'Reformatting MNIST — done!')

################################################################################
# Download CIFAR-10.
################################################################################

def log_cifar10_progress(i, chunk_size, file_size):
    print('\rDownloading CIFAR-10 — %.2f%% complete.'
          % (100 * i * chunk_size / file_size), end='', flush=True)

cifar10_url = 'https://www.cs.toronto.edu/~kriz/cifar-10-matlab.tar.gz'
cifar10_dir = TemporaryDirectory()
cifar10_path = cifar10_dir.name

urlretrieve(cifar10_url, join(cifar10_path, '_cifar10.tar.gz'),
            log_cifar10_progress)
print(80 * '\b \b' + 'Downloading CIFAR-10 — done!')

################################################################################
# Reformat and Gamma-Decode CIFAR-10.
################################################################################

print('Reformatting CIFAR-10...', end='', flush=True)
unpack_archive(join(cifar10_path, '_cifar10.tar.gz'), cifar10_path)

cifar10_a_tr = [
    io.loadmat(join(cifar10_path, 'cifar-10-batches-mat/data_batch_%i.mat' % i))
    for i in range(1, 6)]
cifar10_a_ts = [
    io.loadmat(join(cifar10_path, 'cifar-10-batches-mat/test_batch.mat'))]
cifar10_b = {
    'm_sym': np.ones(10),
    'x0_tr': np.vstack([
        np.float32(np.transpose(
            np.reshape(batch['data']**2.2, (-1, 3, 32, 32)),
            (0, 2, 3, 1)) / 255**2.2)
        for batch in cifar10_a_tr]),
    'x0_ts': np.vstack([
        np.float32(np.transpose(
            np.reshape(batch['data']**2.2, (-1, 3, 32, 32)),
            (0, 2, 3, 1)) / 255**2.2)
        for batch in cifar10_a_ts]),
    'y_tr': np.vstack([
        np.float32(batch['labels'] == np.arange(10))
        for batch in cifar10_a_tr]),
    'y_ts': np.vstack([
        np.float32(batch['labels'] == np.arange(10))
        for batch in cifar10_a_ts])}

cifar2_classes = np.array([
    [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
    [1, 1, 1, 1, 1, 1, 1, 0, 1, 1]])
cifar5_classes = np.array([
    [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
    [1, 1, 1, 0, 0, 0, 1, 0, 1, 1]])

cifar2 = {
    'm_sym': np.ones(2),
    'x0_tr': cifar10_b['x0_tr'],
    'x0_ts': cifar10_b['x0_ts'],
    'y_tr': np.dot(cifar10_b['y_tr'], cifar2_classes.T),
    'y_ts': np.dot(cifar10_b['y_ts'], cifar2_classes.T)}
cifar5 = {
    'm_sym': np.ones(5),
    'x0_tr': cifar10_b['x0_tr'],
    'x0_ts': cifar10_b['x0_ts'],
    'y_tr': np.dot(cifar10_b['y_tr'], cifar5_classes.T),
    'y_ts': np.dot(cifar10_b['y_ts'], cifar5_classes.T)}

makedirs('data/', exist_ok=True)
np.savez_compressed('data/cifar-2.npz', cifar2)
np.savez_compressed('data/cifar-5.npz', cifar5)
np.savez_compressed('data/cifar-10.npz', cifar10_b)
print(80 * '\b \b' + 'Reformatting CIFAR-10 — done!')

################################################################################
# Combine MNIST and CIFAR-10.
################################################################################

def recolor(x, d_min=0.3):
    x_tf = []
    for x_i in x:
        c0 = np.float32(rand.rand(3))
        c1 = np.float32(rand.rand(3))
        while np.sqrt(np.sum(np.square(c1 - c0))) < d_min:
            c1 = np.float32(rand.rand(3))
        x_tf.append(c0 + (c1 - c0) * x_i)
    return np.array(x_tf)

rand.seed(0)
print('Combining MNIST and CIFAR-10...', end='', flush=True)

mnist = np.load('data/mnist.npz')['arr_0'][()]
cifar10 = np.load('data/cifar-10.npz')['arr_0'][()]

hybrid = {
    'm_sym': [0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
    'x0_tr': np.vstack([
        recolor(mnist['x0_tr']),
        cifar10['x0_tr']]),
    'x0_ts': np.vstack([
        recolor(mnist['x0_ts']),
        cifar10['x0_ts']]),
    'y_tr': np.vstack([
        np.hstack([mnist['y_tr'], np.zeros_like(mnist['y_tr'])]),
        np.hstack([np.zeros_like(cifar10['y_tr']), cifar10['y_tr']])]),
    'y_ts': np.vstack([
        np.hstack([mnist['y_ts'], np.zeros_like(mnist['y_ts'])]),
        np.hstack([np.zeros_like(cifar10['y_ts']), cifar10['y_ts']])])}

classes = [0, 1, 2, 3, 4, 10, 11, 14, 16, 17]
m_tr = [np.argmax(hybrid['y_tr'], 1) == c for c in classes]
m_ts = [np.argmax(hybrid['y_ts'], 1) == c for c in classes]
hybrid['x0_tr'] = np.vstack([
    hybrid['x0_tr'][m]
    for m in m_tr])
hybrid['x0_ts'] = np.vstack([
    hybrid['x0_ts'][m]
    for m in m_ts])
hybrid['y_tr'] = np.vstack([
    np.float32(np.ones((np.sum(m), 1)) * (np.arange(len(classes)) == c))
    for c, m in enumerate(m_tr)])
hybrid['y_ts'] = np.vstack([
    np.float32(np.ones((np.sum(m), 1)) * (np.arange(len(classes)) == c))
    for c, m in enumerate(m_ts)])

makedirs('data/', exist_ok=True)
np.savez_compressed('data/hybrid.npz', hybrid)
print(80 * '\b \b' + 'Combining MNIST and CIFAR-10 — done!')
