import nnets.aiobrain.personality as aiop
import nnets.core.trainingset as nts
import nnets.core.utils as nut
import nnets.gru_layer as ngrul

import nnets.nnetworks.multilayer_perceptron as mlp

import sklearn
from sklearn.datasets import make_moons
import matplotlib.pyplot as plt
from PIL import Image
import unittest,wikipedia,numpy,math,os,pylab
import io,gzip,time
import six.moves.cPickle as pickle



import pdb

#python3 -m unittest discover -v ./ "*tests.py"


def generate_sentence(model,wtoi,itow,min_length = 5):
    # We start the sentence with the start token
    new_sentence = [wtoi[nut.SENTENCE_START_TOKEN]]
    # Repeat until we get an end token
    while not new_sentence[-1] == wtoi[nut.SENTENCE_END_TOKEN]:
        next_word_probs = model.predict(new_sentence)[-1]
        samples = numpy.random.multinomial(1,next_word_probs)
        sampled_word = numpy.argmax(samples)
        new_sentence.append(sampled_word)
        # Seomtimes we get stuck if the sentence becomes too long, e.g. "........" :(
        # And: We don't want sentences with UNKNOWN_TOKEN's
        if len(new_sentence) > 100 or sampled_word == wtoi[nut.UNKNOWN_TOKEN]:
            return None
    if len(new_sentence) < min_length:
        return None
    return new_sentence


class test_brain(unittest.TestCase):


    def problem(self,name,features,categories,targetlevel,
            tpgen,traincnt,validcnt,testcnt,nkws,tkws,dsets = None):
        if dsets is None:
            # create training data using tpgen generator
            train_dset = nts.trainingpair(name+'_train',features,categories,traincnt)
            valid_dset = nts.trainingpair(name+'_valid',features,categories,validcnt)
            test_dset = nts.trainingpair(name+'_test',features,categories,testcnt)
            #train_x,train_y = train_dset.accumulate(tpgen).save().xy()
            #valid_x,valid_y = valid_dset.accumulate(tpgen).save().xy()
            #test_x,test_y   =  test_dset.accumulate(tpgen).save().xy()
            train_x,train_y = train_dset.generate(tpgen)
            valid_x,valid_y = valid_dset.generate(tpgen)
            test_x,test_y   =  test_dset.generate(tpgen)
            dsets = (train_x,train_y,valid_x,valid_y,test_x,test_y)

            # show a histogram of the data sets
            #dss = (train_y,valid_y,test_y)
            #self.tset_distributions(dss)

        # create and train a network using the training set
        nn = aiop.initnetwork(name,features,len(categories),**nkws)
        r = nn.train(dsets,targetlevel,**tkws)
        self.assertGreaterEqual(nn.level,targetlevel)


    def tset_distributions(self,dss,
            dns = ('train','valid','test'),dcs = ('r','b','g')):
        # compute a histogram of outputs for the test data
        def distribution(y):
            def h(xs,ys):
                hy,bs = numpy.histogram(ys,bins = xs,density = True)
                fm = lambda u,v : (u+v)/2.0
                bx = numpy.array([fm(bs[y-1],bs[y]) for y in range(1,bs.size)])
                return bx,hy
            def step(xs,ys,bf):
                xs = [xs[0]-bf,xs[0]+bf,xs[1]-bf,xs[1]+bf]
                ys = [ys[0],ys[0],ys[1],ys[1]]
                return xs,ys
            bf = 0.45
            bx = [x-0.5 for x in range(3)]
            hx,hy = h(bx,y)
            return step(hx,hy,bf)
        ax = plt.gca()
        ax.plot([-2,2],[1,1],color = 'k',linewidth = 5,linestyle = '--')
        ax.plot([-2,2],[0,0],color = 'k',linewidth = 5,linestyle = '--')
        for ds,dn,dc in zip(dss,dns,dcs):
            ax.plot(*distribution(ds),color = dc,linewidth = 2,
                label = ('%s-distribution' % (dn)))
        ax.set_xlim(-0.1,1.1)
        ax.set_ylim(-0.1,1.1)
        ax.legend()
        plt.show()


    def setUp(self):
        self.networks = {}
        print('RESET\n')


    def mnist_data(self):
        datapath = os.path.join(os.getcwd(),'data','mnist.pkl.gz')
        if not os.path.isfile(datapath):
            origin = ('http://www.iro.umontreal.ca/~lisa/deep/data/mnist/mnist.pkl.gz')
            nut.urlfetch(origin,datapath)
        with gzip.open(datapath,'rb') as f:
            try:tr,vd,ts = pickle.load(f,encoding='latin1')
            except:tr,vd,ts = pickle.load(f)
        dsets = tr[0],tr[1],vd[0],vd[1],ts[0],ts[1]
        return dsets


    def atest_new_mlp(self):
        name,features,categories,targetlevel = 'xor_strange',20,(0,1),1.0
        traincnt,validcnt,testcnt = 20000,50000,10000

        def yf(x):return x,1 if (x[len(x)//2] == 1 and sum(x) < 10) else 0
        def pgen(f):
            v = lambda : numpy.random.choice((0,1))
            while True:yield yf([v() for i in range(f)])
        tpgen = pgen(features)
        train_set = nts.trainingpair(name+'_train',features,categories,traincnt)
        valid_set = nts.trainingpair(name+'_valid',features,categories,validcnt)
        test_set  = nts.trainingpair(name+'_test' ,features,categories, testcnt)
        #train_x,train_y = train_set.accumulate(tpgen).save().xy()
        #valid_x,valid_y = valid_set.accumulate(tpgen).save().xy()
        #test_x,test_y   =  test_set.accumulate(tpgen).save().xy()
        train_x,train_y = train_set.generate(tpgen)
        valid_x,valid_y = valid_set.generate(tpgen)
        test_x,test_y   =  test_set.generate(tpgen)

        # show a histogram of the data sets
        print(train_set)
        dss = (train_y,valid_y,test_y)
        self.tset_distributions(dss)

        n_in,n_out,n_hiddens = features,len(categories),[(100,'tanh')]
        tkws = {
            'learning_rate':0.1,'batch_size':1000,
            'max_time':30.0,'target':targetlevel}
        mlpnn = mlp.multilayer_perceptron(name,n_in,n_out,n_hiddens)
        mlpnn.train(train_x,train_y,valid_x,valid_y,test_x,test_y,**tkws)
        self.assertGreaterEqual(mlpnn.level,targetlevel)


    # attempt to break the wikipedia page on python into complete thoughts (sentences)
    def test_token_document(self):
        #####
        wikicache = nut.persistent_container('data','wikicache')
        if not wikicache.load():
            queries = [
                #'Python programming language',
                #'Tensor',
                #'Twitter',
                'Philosophy',
                #'Spoon',
                #'Computer',
                #'Canada',
                #'Canning',
                    ]
            traindoc = io.StringIO()
            for q in queries:traindoc.write(nut.wiki(q).content)
            wikicache.traindoc = traindoc.getvalue()

            train_x,train_y,wtoi,itow = nut.tokenize([wikicache.traindoc])
            def print_sentence(s):
                xs = list(range(len(itow)))
                worded = [itow[x] if x in xs else 'UU' for x in s[1:-1]]
                print('-'*40+'\n'+' '.join(worded)+'\n'+'-'*40)
            for xsi,ysi in zip(train_x,train_y):
                xsw = [itow[j] for j in xsi]
                print('x-sentence %s' % (xsw))
                ysw = [itow[j] for j in ysi]
                print('y-sentence %s' % (ysw))

            wikicache.trainset = train_x,train_y,wtoi,itow
        wikicache.save()
        #####

        #####
        train_x,train_y,wtoi,itow = wikicache.trainset
        def print_sentence(s):
            xs = list(range(len(itow)))
            worded = [itow[x] if x in xs else 'UU' for x in s[1:-1]]
            print('-'*40+'\n'+' '.join(worded)+'\n'+'-'*40)
        print('sentence count: %i' % (len(train_x)))

        #cb = lambda m,n : print('model:\t',m,'\tand # examples seen:\t',n)
        def cb(m,n): 
            print('model:\t',m.name,'\tand # examples seen:\t',n)

        vocab_size = 8000
        learning_rate = 1e-3
        nepoch = 20
        hidden_dim = 64
        layercount = 5

        model = ngrul.gru_layer('grumodel',vocab_size,hidden_dim,layercount)
        if not model.load():
            print('training new grumodel')
            t1 = time.time()
            model.sgd_step(train_x[0],train_y[0],learning_rate)
            t2 = time.time()
            print("SGD Step time: ~%f milliseconds" % ((t2 - t1) * 1000.))

        ngrul.train_with_sgd(model,train_x,train_y,
            learning_rate,nepoch,decay = 0.9,callback_every = 10,callback = cb,
            timeout = 900)

        model.level = -1
        model.save()
        print('... saved network \'%s\'; trained to level %f ...' % (model.name,model.level))
        #####

        for i in range(10):
            sent = None
            while not sent:sent = generate_sentence(model,wtoi,itow,min_length = 5)
            print_sentence(sent)

        pdb.set_trace()





        goallevel = 0.95
        sn,fs,cs = 'doctoken',vocab.window,2
        dsets = tset.datasets()

        numpy_rng = numpy.random.RandomState(123)
        nkws = {
            #'n_hidden':100,
            #'n_classes':3,
            #'n_embeddings':vocab.vindex,
            #'n_embdimensions':100,
            #'n_windowsize':vocab.window,
            #'ntype':'enn',
            'ntype':'dbn',
            'numpy_rng':numpy_rng,
            'hidden':[
                #(100,'tanh'),
                #(100,'tanh')
                (100,'sigm'),
                (100,'sigm'),
                    ],
                }
        tkws = {
            'max_time':10.0,'learning_rate':0.1,
            #'batch_size':100,
                }
        #self.nnet(self.networks,sn,fs,cs,dsets,goallevel,nkws,tkws)
        raise NotImplemented
        #self.problem(name,features,categories,targetlevel,
        #    None,None,None,None,nkws,tkws,dsets = dsets)
        s = self.networks[sn]

        vdoc = vocab.documents[-1]
        tokens = [[]]
        for vwordcontext in vdoc:
            vword = vocab.numword(vwordcontext[vocab.window//2])
            tokens[-1].append(vword)
            isbreak = s([vwordcontext])[0]
            if isbreak == 1:
                #print('isbreak',isbreak)
                tokens.append([])

        for t in tokens:
            tok = ' '.join(t)
            print('WE DID IT!',s.level)
            print('token:',tok)
            print('WE DID IT!',tok.endswith('.'))

            pdb.set_trace()

        #assert (s.level >= gl)
        return




  
        # do i need a function for each possible sentence length

        import theano, numpy
        from theano import tensor as T

        # nv :: size of our vocabulary
        # de :: dimension of the embedding space
        # cs :: context window size
        nv,de,cs = textvocab.vindex,50,textvocab.window
        embarray = 0.2*numpy.random.uniform(-1.0,1.0,(nv+1,de))
        embeddings = theano.shared(embarray.astype(theano.config.floatX))

        # as many columns as words in the context window 
        #   and as many lines as words in the sentence
        idxs = T.imatrix()
        x = embeddings[idxs].reshape((idxs.shape[0], de*cs)) 
        y_sentence = T.ivector('y_sentence')  # labels

        f = theano.function(inputs = [idxs],outputs = x)
        #print('hea it is boss:',f(numwordcontexts))
        #print('hea it is again boss:',x.get_value(borrow = True))
        pdb.set_trace()

        sn,fs,cs,gl = 'readtext',3,2,0.99
        dsets = xor_dgen(1000,fs,cs)
        nkws = {'ntype':'mlp','n_hidden':100}
        tkws = {'max_time':60.0,'learning_rate':0.8}
        #s = ns.initnetwork(networks,sn,fs,cs,**nkws)
        s = aiop.initnetwork(networks,sn,fs,cs,**nkws)
        r = ns.pracnetwork(networks,sn,dsets,gl,**tkws)
        assert (s.level >= gl)

        #nh,nc,ne,de,cs = 100,,
        #elman_network(nh,nc,ne,de,cs)

        pdb.set_trace()


    def atest_dbn(self):
        name,features,categories,targetlevel = 'mnistdbn',28*28,list(range(10)),0.95
        dsets = self.mnist_data()
        numpy_rng = numpy.random.RandomState(123)
        nkws = {
            'hidden':[
                (100,'sigm'),
                (100,'sigm'),
                #(100,'sigm'),
                    ],
            'ntype':'dbn','numpy_rng':numpy_rng,
                }
        tkws = {'max_time':10.0,'learning_rate':0.1,'batch_size':10}
        #self.nnet(self.networks,sn,fs,cs,dsets,goallevel,nkws,tkws)
        self.problem(name,features,categories,targetlevel,
            None,None,None,None,nkws,tkws,dsets = dsets)


    def atest_xor_strange(self):
        # define the problem:
        name,features,categories,targetlevel = 'xor_strange',20,(0,1),0.99
        traincnt,validcnt,testcnt = 1000,1000,10000

        def yf(x):return x,1 if (x[len(x)//2] == 1 and sum(x) < 10) else 0
        def pgen(f):
            v = lambda : numpy.random.choice((0,1,2,3))
            while True:yield yf([v() for i in range(f)])
        tpgen = pgen(features)
        nkws = {'hidden':[(100,'tanh')],'ntype':'mlp'}
        tkws = {'max_time':30.0,'learning_rate':0.1,'batch_size':100}

        self.problem(name,features,categories,targetlevel,
                tpgen,traincnt,validcnt,testcnt,nkws,tkws)


    def atest_sum(self):
        name,features,categories,targetlevel = 'sum',4,(0,1,2,3,4),0.99
        traincnt,validcnt,testcnt = 1000,1000,10000

        def yf(x):return x,sum(x)
        def pgen(f):
            v = lambda : numpy.random.choice((0,0,1))
            while True:yield yf([v() for i in range(f)])
        tpgen = pgen(features)
        nkws = {'hidden':[(100,'tanh')],'ntype':'mlp'}
        tkws = {'max_time':30.0,'learning_rate':0.1,'batch_size':100}

        self.problem(name,features,categories,targetlevel,
                tpgen,traincnt,validcnt,testcnt,nkws,tkws)


    def atest_mean(self):
        name,features,categories,targetlevel = 'mean',4,(0,1,2,3,4),0.99
        traincnt,validcnt,testcnt = 1000,1000,10000

        def yf(x):return x,int(numpy.mean(x))
        def pgen(f):
            v = lambda : float(numpy.random.choice(categories))
            while True:yield yf([v() for i in range(f)])
        tpgen = pgen(features)
        nkws = {'hidden':[(100,'tanh')],'ntype':'mlp'}
        tkws = {'max_time':10.0,'learning_rate':0.1,'batch_size':100}

        self.problem(name,features,categories,targetlevel,
                tpgen,traincnt,validcnt,testcnt,nkws,tkws)


    def atest_moon(self):
        name,features,categories,targetlevel = 'moon',2,(0,1),0.95
        traincnt,validcnt,testcnt = 1000,1000,10000

        numpy.random.seed(0)
        tx,ty = sklearn.datasets.make_moons(
            traincnt+validcnt+testcnt,noise = 0.2)
        ty_onehot = numpy.eye(2)[ty]
        def pgen(f):
            j = 0
            while True:
                yield tx[j],ty[j]
                j += 1
        tpgen = pgen(features)
        nkws = {'hidden':[(100,'tanh')],'ntype':'mlp'}
        tkws = {'max_time':10.0,'learning_rate':0.1,'batch_size':100}

        self.problem(name,features,categories,targetlevel,
                tpgen,traincnt,validcnt,testcnt,nkws,tkws)


    def atest_mnist(self):
        name,features,categories,targetlevel = 'mnist',28*28,list(range(10)),0.95
        dsets = self.mnist_data()
        nkws = {'hidden':[(100,'tanh')],'ntype':'mlp'}
        tkws = {'max_time':10.0,'learning_rate':0.8}
        self.problem(name,features,categories,targetlevel,
            None,None,None,None,nkws,tkws,dsets = dsets)


    def atest_mnist_ncnn(self):
        name,features,categories,targetlevel = 'mnistcnn',28*28,list(range(10)),0.95
        dsets = self.mnist_data()
        nkws = {'ntype':'cnn','n_hidden':100,'batch_size':1}
        tkws = {'max_time':10.0,'learning_rate':0.01,'batch_size':1}
        #self.nnet(self.networks,sn,fs,cs,dsets,goallevel,nkws,tkws)
        self.problem(name,features,categories,targetlevel,
            None,None,None,None,nkws,tkws,dsets = dsets)




    def atest_3wolf(self):
        # open random image of dimensions 639x516
        dpath = os.path.join(os.getcwd(),'data','3wolfmoon.jpg')
        img = Image.open(open(dpath))

        # dimensions are (height, width, channel)
        img = numpy.asarray(img, dtype='float64') / 256.

        # put image in 4D tensor of shape (1, 3, height, width)
        ishape = (1,3,639,516)
        img_ = img.transpose(2, 0, 1).reshape(*ishape)




        #f = self.networks['mnistcnn']

        '''#
        rng = numpy.random.RandomState(1234)
        input = T.tensor4(name = 'input')   # the data is presented as rasterized images
        #nkerns = [20,50]
        #nk1,nk2 = nkerns
        clay = nlc.lenet_convpool_layer(
            rng,input = input,poolsize = (2,2),
            image_shape = ishape,filter_shape = ishape)

        filtered_img = clay.f(img_)
        '''#

        filtered_img = img_



        # plot original image and first and second components of output
        pylab.subplot(1, 3, 1); pylab.axis('off'); pylab.imshow(img)
        pylab.gray();

        # recall that the convOp output (filtered image) is actually a "minibatch",
        # of size 1 here, so we take index 0 in the first dimension:
        pylab.subplot(1, 3, 2); pylab.axis('off'); pylab.imshow(filtered_img[0, 0, :, :])
        pylab.subplot(1, 3, 3); pylab.axis('off'); pylab.imshow(filtered_img[0, 1, :, :])
        pylab.show()


if __name__ == '__main__':
    unittest.main()





