Merges cands/noise pkl files from multiple segments to single cands/noise file.

    Expects segment cands pkls with have (1) state dict and (2) cands dict.
    Writes tuple state dict and duple of numpy arrays
    A single pkl written per scan using root name fileroot.
    if cleanup, it will remove segments after merging.
    if sizelimit, it will reduce the output file to be less than this many MB.

def merge_segments(filename, scan, cleanup=True, sizelimit=0):
    """ Merges cands/noise pkl files from multiple segments to single cands/noise file.

    Expects segment cands pkls with have (1) state dict and (2) cands dict.
    Writes tuple state dict and duple of numpy arrays
    A single pkl written per scan using root name fileroot.
    if cleanup, it will remove segments after merging.
    if sizelimit, it will reduce the output file to be less than this many MB.
    """

    workdir = os.path.dirname(filename)
    fileroot = os.path.basename(filename)
    candslist = glob.glob(os.path.join(workdir, 'cands_' + fileroot + '_sc' + str(scan) + 'seg*.pkl'))
    noiselist = glob.glob(os.path.join(workdir, 'noise_' + fileroot + '_sc' + str(scan) + 'seg*.pkl'))
    candssegs = sorted([candsfile.rstrip('.pkl').split('seg')[1] for candsfile in candslist])
    noisesegs = sorted([noisefile.rstrip('.pkl').split('seg')[1] for noisefile in noiselist])

    # test for good list with segments
    if not candslist and not noiselist:
        logger.warn('candslist and noiselist are empty.')
        return

    # aggregate cands over segments
    if not os.path.exists(os.path.join(workdir, 'cands_' + fileroot + '_sc' + str(scan) + '.pkl')):
        logger.info('Aggregating cands over segments %s for fileroot %s, scan %d' % (str(candssegs), fileroot, scan))
        logger.debug('%s' % candslist)

        cands = {}
        for candsfile in candslist:
            with open(candsfile, 'r') as pkl:
                state = pickle.load(pkl)
                result = pickle.load(pkl)
            for kk in result.keys():
                cands[kk] = result[kk]
            segment = state.pop('segment')  # remove this key, as it has no meaning after merging segments

        # optionally limit size
        if sizelimit and len(cands):
            logger.debug('Checking size of cands dictionary...')
            if 'snr2' in state['features']:
                snrcol = state['features'].index('snr2')
            elif 'snr1' in state['features']:
                snrcol = state['features'].index('snr1')

            candsize = sys.getsizeof(cands[cands.keys()[0]])/1e6
            maxlen = int(sizelimit/candsize)
            if len(cands) > maxlen:  # need to reduce length to newlen
                logger.info('cands dictionary of length %.1f would exceed sizelimit of %d MB. Trimming to strongest %d candidates' % (len(cands), sizelimit, maxlen))
                snrs = [abs(cands[k][snrcol]) for k in cands.iterkeys()]  # take top snrs
                snrsort = sorted(snrs, reverse=True)
                snrmax = snrsort[maxlen]  # get min snr for given length limit
                cands = {k: v for k,v in cands.items() if abs(v[snrcol]) > snrmax} # new cands dict

        # write cands to single file
        with open(os.path.join(workdir, 'cands_' + fileroot + '_sc' + str(scan) + '.pkl'), 'w') as pkl:
            pickle.dump(state, pkl, protocol=2)
            pickle.dump( (np.array(cands.keys()), np.array(cands.values())), pkl, protocol=2)
            
        if cleanup:
            if os.path.exists(os.path.join(workdir, 'cands_' + fileroot + '_sc' + str(scan) + '.pkl')):
                for candsfile in candslist:
                    os.remove(candsfile)
    else:
        logger.warn('Merged candsfile already exists for scan %d. Not merged.' % scan)

    # aggregate noise over segments
    if not os.path.exists(os.path.join(workdir, 'noise_' + fileroot + '_sc' + str(scan) + '.pkl')):
        logger.info('Aggregating noise over segments %s for fileroot %s, scan %d' % (str(noisesegs), fileroot, scan))
        logger.debug('%s' % noiselist)

        noise = []
        for noisefile in noiselist:
            with open(noisefile, 'r') as pkl:
                result = pickle.load(pkl)   # gets all noises for segment as list
            noise += result

        # write noise to single file
        if len(noise):
            with open(os.path.join(workdir, 'noise_' + fileroot + '_sc' + str(scan) + '.pkl'), 'w') as pkl:
                pickle.dump(noise, pkl, protocol=2)

        if cleanup:
            if os.path.exists(os.path.join(workdir, 'noise_' + fileroot + '_sc' + str(scan) + '.pkl')):
                for noisefile in noiselist:
                    os.remove(noisefile)

    else:
        logger.warn('Merged noisefile already exists for scan %d. Not merged.' % scan)