import os
import numpy as np
from bobo.util import remkdir
import gzip, struct
from array import array
from bobo.image import Image, ImageCategory
from itertools import product, combinations
from bobo.app import setverbosity

TRAIN_IMG_URL = 'http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz'
TRAIN_IMG_SHA1 =  '6c95f4b05d2bf285e1bfb0e7960c31bd3b3f8a7d'
TRAIN_LBL_URL = 'http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz'
TRAIN_LBL_SHA1 = '2a80914081dc54586dbdf242f9805a6b8d2a15fc'
TEST_IMG_URL = 'http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz'
TEST_IMG_SHA1 = 'c3a25af1f52dad7f726cce8cacb138654b760d48'
TEST_LBL_URL = 'http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz'
TEST_LBL_SHA1 = '763e7fa3757d93b0cdec073cef058b2004252c17'

class MNIST(object):
    def __init__(self, outdir):
        """download URLS above to outdir, then run export()"""        
        self.outdir = remkdir(outdir)
        if not self._downloaded():
            print '[viset.mnist]: downloading MNIST to "%s"' % self.outdir
            setverbosity(1)                        
            self._wget()
            print '[viset.mnist]: exporting MNIST to "%s"' % self.outdir            
            self._export()

    def _downloaded(self):
        gzip_downloaded = (os.path.exists(os.path.join(self.outdir, 'train-images-idx3-ubyte.gz')) and
                os.path.exists(os.path.join(self.outdir, 'train-labels-idx1-ubyte.gz')) and
                os.path.exists(os.path.join(self.outdir, 't10k-images-idx3-ubyte.gz')) and
                os.path.exists(os.path.join(self.outdir, 't10k-labels-idx1-ubyte.gz')))             
        unpacked_downloaded = (os.path.exists(os.path.join(self.outdir, 'train-images-idx3-ubyte')) and
                os.path.exists(os.path.join(self.outdir, 'train-labels-idx1-ubyte')) and
                os.path.exists(os.path.join(self.outdir, 't10k-images-idx3-ubyte')) and
                os.path.exists(os.path.join(self.outdir, 't10k-labels-idx1-ubyte')))             
        return (unpacked_downloaded or gzip_downloaded)

            
    def _wget(self):
        os.system('wget --directory-prefix=%s %s' % (self.outdir, TRAIN_IMG_URL))
        os.system('wget --directory-prefix=%s %s' % (self.outdir, TRAIN_LBL_URL))
        os.system('wget --directory-prefix=%s %s' % (self.outdir, TEST_IMG_URL))
        os.system('wget --directory-prefix=%s %s' % (self.outdir, TEST_LBL_URL))
    
    def _labels(self, gzfile):
        with gzip.open(gzfile, 'rb') as file:
            magic, size = struct.unpack(">II", file.read(8))
            if magic != 2049:
                raise ValueError('Magic number mismatch, expected 2049,'
                                'got %d' % magic)
            labels = array("B", file.read())
        return labels

    def _imread(self, dataset, index):
        """Read MNIST encoded images, adapted from: https://github.com/sorki/python-mnist/blob/master/mnist/loader.py"""
        gzfile = None
        
        with gzip.open(gzfile, 'rb') as file:
            magic, size, rows, cols = struct.unpack(">IIII", file.read(16))
            if magic != 2051:
                raise ValueError('Magic number mismatch, expected 2051, got %d' % magic)
            file.seek(index*rows*cols + 16)
            image = numpy.asarray(array("B", file.read(rows*cols)).tolist())
            return numpy.reshape(image, (rows,cols))
		
    def trainset(self):
        y_train = self._labels(os.path.join(self.outdir, 'train-labels-idx1-ubyte.gz')).tolist() 
        return [ImageCategory(filename=os.path.join(self.outdir, 'train', '%06d.png' % k), category=str(y)) for (k, y) in enumerate(y_train) if k <= 0.9*len(y_train)]

    def subset(self, dataset, N=256):
        """Balanced subset"""
        D = {int(c):[im for im in dataset if im.category()==str(c)] for c in self.categories()}
        V = []
        for c in self.categories():
            V = V + (np.random.choice(D[c], np.ceil(N/10.0))).tolist()
        np.random.shuffle(V)
        return V[0:N]
        
    def validationset(self, N=None):
        y_val = self._labels(os.path.join(self.outdir, 'train-labels-idx1-ubyte.gz')).tolist() 
        return [ImageCategory(filename=os.path.join(self.outdir, 'train', '%06d.png' % k), category=str(y)) for (k, y) in enumerate(y_val) if k > 0.9*len(y_val)]
                                            
    def categories(self):
        return range(0,10)
    
    def pairs(self, dataset, N=1E6):
        """Balanced set of randomly selected N positive and negative pairs given dataset"""
        D = {int(c):[im for im in dataset if im.category()==str(c)] for c in self.categories()}
        M = np.maximum(N, 256)
        T = []
        for (p,q) in combinations(self.categories(), 2):  # 45 category combinations 
            T = T + zip(*(np.random.choice(D[q], size=np.ceil(M/(45*4))), np.random.choice(D[p], size=np.ceil(M/(45*4)))))  # different (p,q)
            T = T + zip(*(np.random.choice(D[p], size=np.ceil(M/(45*4))), np.random.choice(D[q], size=np.ceil(M/(45*4)))))  # different (q,p)
            T = T + zip(*(np.random.choice(D[p], size=np.ceil(M/(45*4))), np.random.choice(D[p], size=np.ceil(M/(45*4)))))  # same (p,p)
            T = T + zip(*(np.random.choice(D[q], size=np.ceil(M/(45*4))), np.random.choice(D[q], size=np.ceil(M/(45*4)))))  # same (q,q)
        return T[0:int(N)]

    def triplets(self, dataset, N=1E6):
        """Balanced set of randomly selected N positive and negative pairs given dataset"""
        D = {int(c):[im for im in dataset if im.category()==str(c)] for c in self.categories()}
        M = np.maximum(N, 256)
        T = []
        for (p,q) in combinations(self.categories(), 2):  # 45 category combinations 
            T = T + zip(*(np.random.choice(D[p], size=np.ceil(M/(45*2))), np.random.choice(D[p], size=np.ceil(M/(45*2))), np.random.choice(D[q], size=np.ceil(M/(45*2)))))  # (p,p,q)
            T = T + zip(*(np.random.choice(D[q], size=np.ceil(M/(45*2))), np.random.choice(D[q], size=np.ceil(M/(45*2))), np.random.choice(D[p], size=np.ceil(M/(45*2)))))  # different (q,q,p)
        return T[0:int(N)]
        
    def testset(self):
        y_test = self._labels(os.path.join(self.outdir, 't10k-labels-idx1-ubyte.gz')).tolist() 
        return [ImageCategory(filename=os.path.join(self.outdir, 'test', '%06d.png' % k), category=str(y)) for (k, y) in enumerate(y_test)]

    def dataset(self):
        return (self.trainset(), self.validationset(), self.testset())
    
    def _export(self):
        # Fetch data necessary to initial construction (do not gunzip!)
        test_img_file = os.path.join(self.outdir, 't10k-images-idx3-ubyte.gz')
        train_img_file = os.path.join(self.outdir, 'train-images-idx3-ubyte.gz')
        
        # Write dataset as 28x28 PNG
        remkdir(os.path.join(self.outdir, 'train'))
        remkdir(os.path.join(self.outdir, 'test'))        
        with gzip.open(train_img_file, 'rb') as gzfile:
            magic, size, rows, cols = struct.unpack(">IIII", gzfile.read(16))
            if magic != 2051:
                raise ValueError('Magic number mismatch, expected 2051, got %d' % magic)

            for k in range(60000):
                img = np.asarray(array("B", gzfile.read(rows*cols)).tolist()).reshape( (rows, cols) ).astype(np.uint8)
                im = Image().array(img).resize(rows=32).saveas(os.path.join(self.outdir, 'train', '%06d.png' % k))
                if (k % 10000) == 0:
                    print '[viset.mnist][%d/%d]: exporting training images to "%s"' % (k, 60000, os.path.join(self.outdir, 'train'))
    
        with gzip.open(test_img_file, 'rb') as gzfile:
            magic, size, rows, cols = struct.unpack(">IIII", gzfile.read(16))
            if magic != 2051:
                raise ValueError('Magic number mismatch, expected 2051, got %d' % magic)
                
            for k in range(10000):
                img = np.asarray(array("B", gzfile.read(rows*cols)).tolist()).reshape( (rows, cols) ).astype(np.uint8)
                im = Image().array(img).resize(rows=32).saveas(os.path.join(self.outdir, 'test', '%06d.png' % k))
                if (k % 10000) == 0:            
                    print '[viset.mnist][%d/%d]: exporting testing images to "%s"' % (k, 10000, os.path.join(self.outdir, 'test'))
        
        # Cleanup    
        return self
    
    
