import os, sys, time

import numpy
import theano
from theano import tensor, config
from theano.sandbox import rng_mrg
from theano.sandbox.rng_mrg import MRG_RandomStreams
from theano.sandbox.cuda import cuda_available, cuda_enabled
if cuda_available:
    from theano.sandbox.cuda import float32_shared_constructor

import unittest
from theano.tests import unittest_tools as utt

#TODO: test gpu
# Done in test_consistency_GPU_{serial,parallel}

#TODO: test MRG_RandomStreams
# Partly done in test_consistency_randomstreams

#TODO: test optimizer mrg_random_make_inplace
#TODO: make tests work when no flags gived. Now need: THEANO_FLAGS=device=gpu0,floatX=float32
# Partly done, in test_consistency_GPU_{serial,parallel}

#TODO: bug fix test_normal0, in normal() fct, n_samples currently need to be numpy.prod(size) not self.n_streams(size)

mode = theano.config.mode
utt.seed_rng()

## Results generated by Java code using L'Ecuyer et al.'s code, with:
# main seed: [12345]*6 (default)
# 12 streams
# 7 substreams for each stream
# 5 samples drawn from each substream
java_samples = numpy.loadtxt(os.path.join(os.path.split(theano.__file__)[0],
                                          'sandbox','samples_MRG31k3p_12_7_5.txt'))


def test_deterministic():
    seed = utt.fetch_seed()
    sample_size = (10, 20)

    test_use_cuda = [False]
    if cuda_enabled:
        test_use_cuda.append(True)

    for use_cuda in test_use_cuda:
        print 'use_cuda =', use_cuda
        R = MRG_RandomStreams(seed=seed, use_cuda=use_cuda)
        u = R.uniform(size=sample_size)
        f = theano.function([], u)

        fsample1 = f()
        fsample2 = f()
        assert not numpy.allclose(fsample1, fsample2)

        R2 = MRG_RandomStreams(seed=seed, use_cuda=use_cuda)
        u2 = R2.uniform(size=sample_size)
        g = theano.function([], u2)
        gsample1 = g()
        gsample2 = g()
        assert numpy.allclose(fsample1, gsample1)
        assert numpy.allclose(fsample2, gsample2)


def test_consistency_randomstreams():
    '''Verify that the random numbers generated by MRG_RandomStreams
    are the same as the reference (Java) implementation by L'Ecuyer et al.
    '''

    seed = 12345
    n_samples = 5
    n_streams = 12
    n_substreams = 7

    test_use_cuda = [False]
    if cuda_enabled:
        test_use_cuda.append(True)

    for use_cuda in test_use_cuda:
        print 'use_cuda =', use_cuda
        samples = []
        rng = MRG_RandomStreams(seed = seed, use_cuda=False)
        for i in range(n_streams):
            stream_samples = []
            u = rng.uniform(size=(n_substreams,), nstreams=n_substreams)
            f = theano.function([], u)
            for j in range(n_samples):
                s = f()
                stream_samples.append(s)
            stream_samples = numpy.array(stream_samples)
            stream_samples = stream_samples.T.flatten()
            samples.append(stream_samples)

        samples = numpy.array(samples).flatten()
        assert(numpy.allclose(samples, java_samples))

def test_consistency_cpu_serial():
    '''Verify that the random numbers generated by mrg_uniform, serially,
    are the same as the reference (Java) implementation by L'Ecuyer et al.
    '''
    seed = 12345
    n_samples = 5
    n_streams = 12
    n_substreams = 7

    samples = []
    curr_rstate = numpy.array([seed] * 6, dtype='int32')

    for i in range(n_streams):
        stream_rstate = curr_rstate.copy()
        for j in range(n_substreams):
            rstate = tensor.shared(numpy.array([stream_rstate.copy()], dtype='int32'))
            new_rstate, sample = rng_mrg.mrg_uniform.new(rstate, ndim=None, dtype=config.floatX, size=(1,))
            # Not really necessary, just mimicking rng_mrg.MRG_RandomStreams' behavior
            sample.rstate = rstate
            sample.update = (rstate, new_rstate)

            rstate.default_update = new_rstate
            f = theano.function([], sample)
            for k in range(n_samples):
                s = f()
                samples.append(s)

            # next substream
            stream_rstate = rng_mrg.ff_2p72(stream_rstate)

        # next stream
        curr_rstate = rng_mrg.ff_2p134(curr_rstate)

    samples = numpy.array(samples).flatten()
    assert(numpy.allclose(samples, java_samples))

def test_consistency_cpu_parallel():
    '''Verify that the random numbers generated by mrg_uniform, in parallel,
    are the same as the reference (Java) implementation by L'Ecuyer et al.
    '''
    seed = 12345
    n_samples = 5
    n_streams = 12
    n_substreams = 7 # 7 samples will be drawn in parallel

    samples = []
    curr_rstate = numpy.array([seed]*6, dtype='int32')

    for i in range(n_streams):
        stream_samples = []
        rstate = [curr_rstate.copy()]
        for j in range(1, n_substreams):
            rstate.append(rng_mrg.ff_2p72(rstate[-1]))
        rstate = numpy.asarray(rstate)
        rstate = tensor.shared(rstate)

        new_rstate, sample = rng_mrg.mrg_uniform.new(rstate, ndim=None,
                dtype=config.floatX, size=(n_substreams,))
        # Not really necessary, just mimicking rng_mrg.MRG_RandomStreams' behavior
        sample.rstate = rstate
        sample.update = (rstate, new_rstate)

        rstate.default_update = new_rstate
        f = theano.function([], sample)

        for k in range(n_samples):
            s = f()
            stream_samples.append(s)

        samples.append(numpy.array(stream_samples).T.flatten())

        # next stream
        curr_rstate = rng_mrg.ff_2p134(curr_rstate)

    samples = numpy.array(samples).flatten()
    assert(numpy.allclose(samples, java_samples))

def test_consistency_GPU_serial():
    '''Verify that the random numbers generated by GPU_mrg_uniform, serially,
    are the same as the reference (Java) implementation by L'Ecuyer et al.
    '''
    if not cuda_available:
        raise SkipTest('Optional package cuda not available')
    if config.mode == 'FAST_COMPILE':
        mode = 'FAST_RUN'
    else:
        mode = config.mode

    seed = 12345
    n_samples = 5
    n_streams = 12
    n_substreams = 7

    samples = []
    curr_rstate = numpy.array([seed] * 6, dtype='int32')

    for i in range(n_streams):
        stream_rstate = curr_rstate.copy()
        for j in range(n_substreams):
            substream_rstate = numpy.array(stream_rstate.copy(), dtype='int32')
            # HACK - we transfer these int32 to the GPU memory as float32
            # (reinterpret_cast)
            tmp_float_buf = numpy.frombuffer(substream_rstate.data, dtype='float32')
            rstate = float32_shared_constructor(tmp_float_buf) # Transfer to device

            new_rstate, sample = rng_mrg.GPU_mrg_uniform.new(rstate, ndim=None,
                    dtype='float32', size=(1,))
            rstate.default_update = new_rstate

            # Not really necessary, just mimicking rng_mrg.MRG_RandomStreams' behavior
            sample.rstate = rstate
            sample.update = (rstate, new_rstate)

            # We need the sample back in the main memory
            cpu_sample = tensor.as_tensor_variable(sample)
            f = theano.function([], cpu_sample, mode=mode)
            for k in range(n_samples):
                s = f()
                samples.append(s)

            # next substream
            stream_rstate = rng_mrg.ff_2p72(stream_rstate)

        # next stream
        curr_rstate = rng_mrg.ff_2p134(curr_rstate)

    samples = numpy.array(samples).flatten()
    assert(numpy.allclose(samples, java_samples))

def test_consistency_GPU_parallel():
    '''Verify that the random numbers generated by GPU_mrg_uniform, in parallel,
    are the same as the reference (Java) implementation by L'Ecuyer et al.
    '''
    if not cuda_available:
        raise SkipTest('Optional package cuda not available')
    if config.mode == 'FAST_COMPILE':
        mode = 'FAST_RUN'
    else:
        mode = config.mode

    seed = 12345
    n_samples = 5
    n_streams = 12
    n_substreams = 7 # 7 samples will be drawn in parallel

    samples = []
    curr_rstate = numpy.array([seed]*6, dtype='int32')

    for i in range(n_streams):
        stream_samples = []
        rstate = [curr_rstate.copy()]
        for j in range(1, n_substreams):
            rstate.append(rng_mrg.ff_2p72(rstate[-1]))
        rstate = numpy.asarray(rstate).flatten()
        # HACK - transfer these int32 to the GPU memory as float32
        # (reinterpret_cast)
        tmp_float_buf = numpy.frombuffer(rstate.data, dtype='float32')
        rstate = float32_shared_constructor(tmp_float_buf) # Transfer to device

        new_rstate, sample = rng_mrg.GPU_mrg_uniform.new(rstate, ndim=None,
                dtype='float32', size=(n_substreams,))
        rstate.default_update = new_rstate

        # Not really necessary, just mimicking rng_mrg.MRG_RandomStreams' behavior
        sample.rstate = rstate
        sample.update = (rstate, new_rstate)

        # We need the sample back in the main memory
        cpu_sample = tensor.as_tensor_variable(sample)
        f = theano.function([], cpu_sample, mode=mode)

        for k in range(n_samples):
            s = f()
            stream_samples.append(s)

        samples.append(numpy.array(stream_samples).T.flatten())

        # next stream
        curr_rstate = rng_mrg.ff_2p134(curr_rstate)

    samples = numpy.array(samples).flatten()
    assert(numpy.allclose(samples, java_samples))

def test_rng0():

    def basictest(f, steps, prefix="", allow_01=False):
        dt = 0.0
        for i in xrange(steps):
            t0 = time.time()
            ival = f()
            dt += time.time() - t0
            ival = numpy.asarray(ival)
            if i == 0:
                mean = numpy.array(ival, copy=True)
                min_ = ival.min()
                max_ = ival.max()
            else:
                alpha = 1.0 / (1+i)
                mean = alpha * ival + (1-alpha)*mean
                min_ = min(min_,ival.min())
                max_ = max(max_,ival.max())
            if not allow_01:
                assert min_ > 0
                assert max_ < 1

        print prefix, 'mean', numpy.mean(mean)
        assert abs(numpy.mean(mean) - 0.5) < .01, 'bad mean?'
        print prefix, 'time', dt
        print prefix, 'elements', steps*sample_size[0]*sample_size[1]
        print prefix, 'samples/sec', steps*sample_size[0]*sample_size[1] / dt
        print prefix, 'min',min_,'max',max_

    if mode in ['DEBUG_MODE','FAST_COMPILE']:
        sample_size = (10,100)
        steps = int(1e2)
    else:
        sample_size = (1000,100)
        steps = int(1e3)

    print ''
    print 'ON CPU:'

    R = MRG_RandomStreams(234, use_cuda=False)
    u = R.uniform(size=sample_size)
    f = theano.function([], u, mode=mode)
    theano.printing.debugprint(f)
    print 'random?[:10]\n', f()[0,0:10]
    basictest(f, steps, prefix='mrg  cpu')

    if mode!='FAST_COMPILE':
        print ''
        print 'ON GPU:'
        R = MRG_RandomStreams(234, use_cuda=True)
        u = R.uniform(size=sample_size, dtype='float32')
        assert u.dtype == 'float32' #well, it's really that this test w GPU doesn't make sense otw
        f = theano.function([], theano.Out(
                theano.sandbox.cuda.basic_ops.gpu_from_host(u),
                borrow=True), mode=mode)
        theano.printing.debugprint(f)
        print 'random?[:10]\n', numpy.asarray(f())[0,0:10]
        basictest(f, steps, prefix='mrg  gpu')

    print ''
    print 'ON CPU w NUMPY:'
    RR = theano.tensor.shared_randomstreams.RandomStreams(234)

    uu = RR.uniform(size=sample_size)
    ff = theano.function([], uu, mode=mode)
    # It's not our problem if numpy generates 0 or 1
    basictest(ff, steps, prefix='numpy', allow_01=True)




def test_normal0():

    if config.mode == 'FAST_COMPILE':
        mode = 'FAST_RUN'
    else:
        mode = config.mode
    def basictest(f, steps, target_avg, target_std, prefix=""):
        dt = 0.0
        avg_std = 0.0
        for i in xrange(steps):
            t0 = time.time()
            ival = f()
            dt += time.time() - t0
            ival = numpy.asarray(ival)
            if i == 0:
                mean = numpy.array(ival, copy=True)
                avg_std = numpy.std(ival)
            else:
                alpha = 1.0 / (1+i)
                mean = alpha * ival + (1-alpha)*mean
                avg_std = alpha * numpy.std(ival) + (1-alpha)*avg_std

        print prefix, 'mean', numpy.mean(mean)
        assert abs(numpy.mean(mean) - target_avg) < .01, 'bad mean?'
        print prefix, 'std', avg_std
        assert abs(avg_std - target_std) < .01, 'bad std?'
        print prefix, 'time', dt
        print prefix, 'elements', steps*sample_size[0]*sample_size[1]
        print prefix, 'samples/sec', steps*sample_size[0]*sample_size[1] / dt

    sample_size = (999,100)
    print ''
    print 'ON CPU:'

    R = MRG_RandomStreams(234, use_cuda=False)
    n = R.normal(size=sample_size, avg=-5.0, std=2.0)
    f = theano.function([], n, mode=mode)
    theano.printing.debugprint(f)
    print 'random?[:10]\n', f()[0,0:10]
    basictest(f, 50, -5.0, 2.0, prefix='mrg ')

    sys.stdout.flush()

    # now with odd number of samples
    sample_size = (999,99)


    print ''
    print 'ON GPU:'
    R = MRG_RandomStreams(234, use_cuda=True)
    n = R.normal(size=sample_size, avg=-5.0, std=2.0, dtype='float32')
    assert n.dtype == 'float32' #well, it's really that this test w GPU doesn't make sense otw
    f = theano.function([], theano.Out(
        theano.sandbox.cuda.basic_ops.gpu_from_host(n),
        borrow=True), mode=mode)
    
    theano.printing.debugprint(f)
    sys.stdout.flush()
    print 'random?[:10]\n', numpy.asarray(f())[0,0:10]
    print '----'
    sys.stdout.flush()
    basictest(f, 50, -5.0, 2.0, prefix='gpu mrg ')


    print ''
    print 'ON CPU w NUMPY:'
    RR = theano.tensor.shared_randomstreams.RandomStreams(234)

    nn = RR.normal(size=sample_size, avg=-5.0, std=2.0)
    ff = theano.function([], nn)

    basictest(ff, 50, -5.0, 2.0, prefix='numpy ')


   
