#!/usr/bin/env python
# Creates a data-based fourier mask for each particle named pfmask_#####.em
# Sigma = 2.3 works for me. Run in Dynamo data folder.
# Files must be named particle_#####.em from 00001 and without gaps (ie. Dynamo format)
# Requires Appion and EMAN2
# Usage: ./auto_fourier_masking_particles.py <sigma threshold> <# procs>
import os, sys, glob, numpy as np, multiprocessing as mp
from pyami import mrc
from appionlib.apImage import imagefilter
filelist=glob.glob('particle_*.em')
def createFourierMask(i):
	if os.path.isfile(os.path.join(os.getcwd(),'particle_%05d.em') % i) is True:
		fname = 'particle_%05d' % i
		mname = 'pfmask_%05d' % i
	else:
		fname = 'particle_%06d' % i
		mname = 'pfmask_%06d' % i
	if os.path.isfile(os.path.join(os.getcwd(),mname+'.em')) is False:
		os.system('e2proc3d.py %s.em %s.mrc' % (fname, fname))
		f = mrc.read('%s.mrc' % fname)
		ft_big = imagefilter.scaleImage(np.sqrt(np.square(np.real(np.fft.fftn(f)))),2)
		roll_length = len(ft_big)/2 - 1
		ft_big = np.roll(np.roll(np.roll(ft_big,roll_length-1),roll_length,0),roll_length-1,1)
		ft = imagefilter.scaleImage(ft_big,0.5)
		ft2 = ft
		low = ft2 < ft.mean() + float(sys.argv[1])*ft.std()
		high = ft2 >= ft.mean() + float(sys.argv[1])*ft.std()
		ft2[low] = 0
		ft2[high] = 1
		mrc.write(ft2,'%s.mrc' % mname)
		os.system('e2proc3d.py %s.mrc %s.em;rm %s.mrc %s.mrc' % (mname, mname, mname, fname))
last_particle=glob.glob('particle_*.em')
last_particle.sort()
last_particle=last_particle[-1]
last_particle=int(last_particle.split('_')[1].split('.')[0])
for i in range(1,last_particle+1):
	p = mp.Process(target=createFourierMask, args=(i,))
	p.start()
	if (i % int(sys.argv[2]) == 0) and (i != 0):
		[p.join() for p in mp.active_children()]