import random, json, os
import numpy as np

def make_n_blocks(p, n):
    """Returns a block, i.e., a random number of samples of a Bernoulli distribution of probability p."""
    """The number of samples is randomly sampled from a geometric distribution with mean n"""
    samples = []

    samp_num = np.random.geometric(1/n)
    #The number of samples for a given block is capped at 18
    if samp_num > 18:
        samp_num = 18

    for _ in range(samp_num):
        samples.append( 1 if random.random() < p else 0 )
    return {'p':p, 'samples':samples}

def make_blockset(N, n, minp=0., maxp=1.):
    """Make a blockset, i.e., a list of N blocks of n samples on average, with randomized probabilities.
    N will be the number of blocks one subject sees. """
    r = []
    for _ in range(N):
        p = random.uniform(minp, maxp)
        r.append(make_n_blocks(p, n))
        #r.append(make_n_blocks(p, np.random.geometric(0.2)))
    return r

def make_M_blocksets(M, N, n, minp=0., maxp=1.):
    '''Make M blocksets, each containing N blocks of on average n samples.
    M should be comparable to the number of subjects.'''
    r = []
    for _ in range(M):
        #np.random.seed(1)
        r.append( make_blockset(N,n, minp, maxp) )
    return r

def make_and_save_blocksets(M,N,n,filename, minp=0., maxp=1.):
    blocksets = make_M_blocksets(M, N, n, minp, maxp)
    f = open(filename, 'w')
    json.dump(blocksets, f)
    f.close()

def load_blocksets(filename):
    f = open(filename, 'r')
    blocksets = json.load(f)
    f.close()
    return blocksets

def get_or_make_blocksets(Msets, Nblocks, ntrials):
    fpath = 'ringvar/blocksets-%d-%d-%d.json' % (Msets, Nblocks, ntrials)
    if not os.path.isfile(fpath):
        make_and_save_blocksets(Msets, Nblocks, ntrials, fpath)
    return load_blocksets(fpath)