#!/usr/bin/env python
# Usage: ./stack_binning.py <stack.mrcs> <binning factor>
import os, sys, numpy as np
from pyami import mrc
from appionlib.apImage import imagefilter
stack = mrc.read(sys.argv[1])
sampling = int(sys.argv[2])
dimx = len(stack[0])
dimy = len(stack[0][0])
shape = np.asarray(stack[0].shape)
bin2 = sampling * 2
remain = shape % bin2
if remain.any():
	new_dimx = round(int(shape[0]/bin2)*bin2/sampling)
	new_dimy = round(int(shape[1]/bin2)*bin2/sampling)
else:
	new_dimx = int(round(dimx/sampling))
	new_dimy = int(round(dimy/sampling))
binned_stack = np.zeros((len(stack),new_dimx,new_dimy))
for i in range(len(stack)):
	binned_stack[i,:,:] = imagefilter.binImg(stack[i,:,:], bin=sampling)
stack_file = '%s_bin%s.mrcs' % (os.path.splitext(sys.argv[1])[0], sampling)
mrc.write(binned_stack, stack_file)
