import random, json, os

def make_n_blocks(p, n):
    """Returns a block, i.e., n samples of a Bernoulli distribution of probability p."""
    samples = []
    for _ in range(n):
        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, with randomized probabilities.
    N will be the number of blocks one subject sees.
    (and N*n the number of rings they see)."""
    r = []
    for _ in range(N):
        p = random.uniform(minp, maxp)
        r.append(make_n_blocks(p, n))
    return r

def make_M_blocksets(M, N, n, minp=0., maxp=1.):
    '''Make M blocksets, each containing N blocks of n samples.
    M should be comparable to the number of subjects.'''
    r = []
    for _ in range(M):
        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)