import unittest, os
import numpy
import six.moves.cPickle as pickle
from theano.compat import DictMixin, OrderedDict
import theano
import theano.tensor as T

floatX = 'float32'

def test_graph_opt_caching():
    opt_db_file = theano.config.compiledir+'/optimized_graphs.pkl'
    os.system('rm %s'%opt_db_file)
    
    mode = theano.config.mode
    if mode in ["DEBUG_MODE", "DebugMode"]:
        mode = "FAST_RUN"
    default = theano.config.cache_optimizations
    try:
        theano.config.cache_optimizations = True
        a = T.fmatrix('a')
        b = T.fmatrix('b')
        c = theano.shared(numpy.ones((10, 10), dtype=floatX))
        d = theano.shared(numpy.ones((10, 10), dtype=floatX))
        e = T.sum(T.sum(T.sum(a ** 2 + b) + c) + d)
        f1 = theano.function([a, b], e, mode=mode)

        m = T.fmatrix('x1')
        n = T.fmatrix('x2')
        p = theano.shared(numpy.ones((10, 10), dtype=floatX))
        q = theano.shared(numpy.ones((10, 10), dtype=floatX))
        j = T.sum(T.sum(T.sum(m ** 2 + n) + p) + q)
        f2 = theano.function([m, n], j, mode=mode)
        
        in1 = numpy.ones((10, 10), dtype=floatX)
        in2 = numpy.ones((10, 10), dtype=floatX)
        assert f1(in1, in2) == f2(in1, in2)
    finally:
        theano.config.cache_optimizations = default
        
if __name__ == '__main__':
    test_graph_opt_caching()
