#!/usr/bin/env python
import subprocess
import sys
import re
import os
import gzip
import glob

bin_path = os.path.dirname(__file__)
split_exe = os.path.join(bin_path, 'fast-split')
bc_exe = os.path.join(bin_path, 'segment')

def main():
    ''' Recursively splits all clusters in voxel-file into 
    individual particles.'''
    if len(sys.argv) != 2:
        print 'segmentation [voxel-file]'
        sys.exit(1)
    else:
        input_file = sys.argv[1] 
        output_file = 'particles-fs.gz'

    subprocess.call([split_exe, input_file], stdout=subprocess.PIPE)
    files = glob.glob('p-*.gz')
    while files:
        for f in sorted(files):
            split_cluster(f)
        files = glob.glob('p-*.gz')
    merge_finished()

def merge_finished():
    ''' Merge all finished files into a single one. '''
    finished = gzip.open('segmented-particles.gz', 'wb')
    for i,f in enumerate(sorted(glob.glob('final-*-*.gz'))):
        for line in gzip.open(f, 'rb'):
            m = re.match('begin cluster; id \d+; size (\d+)\n', line)
            if m:
                size = m.group(1)
                finished.write('begin cluster; id {}; size {}\n'.format(i, size))
            else:
                finished.write(line)

def split_cluster(input_file):
    '''  '''
    next_id = highest_file('p-*.gz') + 1
    cmd = [bc_exe, input_file, str(next_id), '2.0']
    pid = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    stdout = pid.communicate()[0]
    print stdout
    files = re.findall('p-\d+.gz', stdout)
    # If no voxels were removed, then move output file to finished.
    if 'Removed 0 bridging voxels' in stdout:
        dest = 'f-{:03d}.gz'.format(highest_file('f-*.gz')+1)
        os.rename(files[0], dest)
    # Delete the file that was just processed.
    os.remove(input_file)
        
def highest_file(pattern):
    files = glob.glob(pattern)
    if not files: 
        return 0
    else:
        return max(int(re.search('-(\d+)', f).group(1)) for f in files)

if __name__ == '__main__': main()
