# Path: distributions/io/stream.py
# class json_stream_load(object):
#     '''
#     Read json data that was created by json_stream_dump or json_costream_dump.
# 
#     Note that this exploits newline formatting in the above dumpers.
#     In particular:
#     - the first line is '['
#     - intermediate lines are of the form '{},'.format(json_parsable_content)
#     - the penultimate line is of the form '{}'.format(json_parsable_content)
#     - the last line is ']'
#     - there is no trailing whitespace
# 
#     An alternative would be to use ijson to streamingly load arbitrary json
#     files, however in practice this is ~40x slower.
#     '''
#     def __init__(self, filename):
#         self.fd = open_compressed(filename, 'rb')
#         line = self.fd.readline(2)
#         if line != '[\n':
#             raise IOError(
#                 'Unhandled format for json_stream_load. '
#                 'Try recreating json file with the compatible '
#                 'json_stream_dump or json_costream_dump.')
# 
#     def __iter__(self):
#         return self
# 
#     def next(self):
#         line = self.fd.readline().rstrip(',\n')
#         if line == ']':
#             self.close()
#             raise StopIteration
#         else:
#             return simplejson.loads(line)
# 
#     def close(self):
#         self.fd.close()
# 
# def json_stream_dump(stream, filename, **kwargs):
#     kwargs['separators'] = (',', ':')
#     stream = iter(stream)
#     with open_compressed(filename, 'w') as f:
#         f.write('[')
#         try:
#             item = next(stream)
#             f.write('\n')
#             simplejson.dump(item, f, **kwargs)
#             for item in stream:
#                 f.write(',\n')
#                 simplejson.dump(item, f, **kwargs)
#         except StopIteration:
#             pass
#         f.write('\n]')


import os
import numpy
import math
import matplotlib
import parsable
from collections import defaultdict
from numpy import log, exp
from matplotlib import pyplot, font_manager
from distributions.lp.special import fast_log
from distributions.io.stream import json_stream_load, json_stream_dump

def get_subprobs(size, max_size):
    '''
    Compute probabilities of shapes of partial assignment vectors.

    Inputs:
        size = sample_size
        max_size = dataset_size
    Returns:
        dict : shape -> prob
    '''
    assert 0 <= size
    assert size <= max_size
    cache_file = '{}/subprobs.{}.{}.json.bz2'.format(TEMP, size, max_size)
    if cache_file not in CACHE:
        if os.path.exists(cache_file):
            flat = json_stream_load(cache_file)
            small_probs = {tuple(key): val for key, val in flat}
        else:
            if size == max_size:
                small_probs = get_probs(size)
            else:
                small_counts = get_counts(size)
                large_counts = get_counts(size + 1)
                large_probs = get_subprobs(size + 1, max_size)
                small_probs = get_smaller_probs(
                    small_counts,
                    large_counts,
                    large_probs)
            print 'caching', cache_file

# Path: distributions/io/stream.py
# class json_stream_load(object):
#     '''
#     Read json data that was created by json_stream_dump or json_costream_dump.
# 
#     Note that this exploits newline formatting in the above dumpers.
#     In particular:
#     - the first line is '['
#     - intermediate lines are of the form '{},'.format(json_parsable_content)
#     - the penultimate line is of the form '{}'.format(json_parsable_content)
#     - the last line is ']'
#     - there is no trailing whitespace
# 
#     An alternative would be to use ijson to streamingly load arbitrary json
#     files, however in practice this is ~40x slower.
#     '''
#     def __init__(self, filename):
#         self.fd = open_compressed(filename, 'rb')
#         line = self.fd.readline(2)
#         if line != '[\n':
#             raise IOError(
#                 'Unhandled format for json_stream_load. '
#                 'Try recreating json file with the compatible '
#                 'json_stream_dump or json_costream_dump.')
# 
#     def __iter__(self):
#         return self
# 
#     def next(self):
#         line = self.fd.readline().rstrip(',\n')
#         if line == ']':
#             self.close()
#             raise StopIteration
#         else:
#             return simplejson.loads(line)
# 
#     def close(self):
#         self.fd.close()
# 
# def json_stream_dump(stream, filename, **kwargs):
#     kwargs['separators'] = (',', ':')
#     stream = iter(stream)
#     with open_compressed(filename, 'w') as f:
#         f.write('[')
#         try:
#             item = next(stream)
#             f.write('\n')
#             simplejson.dump(item, f, **kwargs)
#             for item in stream:
#                 f.write(',\n')
#                 simplejson.dump(item, f, **kwargs)
#         except StopIteration:
#             pass
#         f.write('\n]')


# Path: derivations/clustering.py
import os
import numpy
import math
import matplotlib
import parsable
from collections import defaultdict
from numpy import log, exp
from matplotlib import pyplot, font_manager
from distributions.lp.special import fast_log
from distributions.io.stream import json_stream_load, json_stream_dump


def get_subprobs(size, max_size):
    '''
    Compute probabilities of shapes of partial assignment vectors.

    Inputs:
        size = sample_size
        max_size = dataset_size
    Returns:
        dict : shape -> prob
    '''
    assert 0 <= size
    assert size <= max_size
    cache_file = '{}/subprobs.{}.{}.json.bz2'.format(TEMP, size, max_size)
    if cache_file not in CACHE:
        if os.path.exists(cache_file):
            flat = json_stream_load(cache_file)
            small_probs = {tuple(key): val for key, val in flat}
        else:
            if size == max_size:
                small_probs = get_probs(size)
            else:
                small_counts = get_counts(size)
                large_counts = get_counts(size + 1)
                large_probs = get_subprobs(size + 1, max_size)
                small_probs = get_smaller_probs(
                    small_counts,
                    large_counts,
                    large_probs)
            print 'caching', cache_file

            json_stream_dump(small_probs.iteritems(), cache_file)