#!/usr/bin/env python
# Creates a data-based fourier mask for a given tomogram by averaging 9 fourier sub-tomogram transforms.
# Particle size must be even and smaller than the tomogram.
# Requires Appion, EMAN2, and IMOD
# Usage: ./auto_fourier_masking_tomogram.py <tomogram> <sigma threshold> <particle size> <distance from edges>
import os, sys, glob, subprocess, numpy as np
from pyami import mrc
from appionlib.apImage import imagefilter
proc=subprocess.Popen('header %s |grep grid>tempppp;awk \'{print $11}\' tempppp;rm tempppp' % sys.argv[1], stdout=subprocess.PIPE, shell=True)
(x, err) = proc.communicate()
x=int(x)
proc=subprocess.Popen('header %s |grep grid>tempppp;awk \'{print $12}\' tempppp;rm tempppp' % sys.argv[1], stdout=subprocess.PIPE, shell=True)
(y, err) = proc.communicate()
y=int(y)
proc=subprocess.Popen('header %s |grep grid>tempppp;awk \'{print $13}\' tempppp;rm tempppp' % sys.argv[1], stdout=subprocess.PIPE, shell=True)
(z, err) = proc.communicate()
z=int(z)
size = int(sys.argv[3])
dist = int(sys.argv[4])
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 1.mrc' % (dist, dist+size-1, dist, dist+size-1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 2.mrc' % (x/2 - size/2, x/2 + size/2 - 1, dist, dist+size-1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 3.mrc' % (x/2 - size/2, x/2 + size/2 - 1,y/2 - size/2, y/2 + size/2 - 1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 4.mrc' % (dist, dist+size-1, y/2 - size/2, y/2 + size/2 - 1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 5.mrc' % (x-dist-size, x-dist-1, dist, dist+size-1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 6.mrc' % (x-dist-size, x-dist-1, y/2 - size/2, y/2 + size/2 - 1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 7.mrc' % (x-dist-size, x-dist-1, y-dist-size, y-dist-1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 8.mrc' % (x/2 - size/2, x/2 + size/2 - 1, y-dist-size, y-dist-1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
os.system('trimvol -x %s,%s -y %s,%s -z %s,%s %s 9.mrc' % (dist, dist+size-1, y-dist-size, y-dist-1, z/2 - size/2, z/2 + size/2 - 1, sys.argv[1]))
fftsum = np.zeros((size,size,size))
for i in range(1,10):
	v=mrc.read('%s.mrc' % i)
	ft_big = imagefilter.scaleImage(np.square(np.real(np.fft.fftn(v))),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)
	fftsum = fftsum + ft
	os.system('rm %s.mrc' % i)
fftsum = fftsum/9
fftsum2 = fftsum
low = fftsum2 < fftsum.mean() + float(sys.argv[2])*fftsum.std()
high = fftsum2 >= fftsum.mean() + float(sys.argv[2])*fftsum.std()
fftsum2[low] = 0
fftsum2[high] = 1
mrc.write(fftsum2, '%s.pfmask%s.mrc' % (os.path.splitext(sys.argv[1])[0],size))
os.system('e2proc3d.py %s.pfmask%s.mrc %s.pfmask%s.em; rm %s.pfmask%s.mrc' % (os.path.splitext(sys.argv[1])[0],size, os.path.splitext(sys.argv[1])[0],size, os.path.splitext(sys.argv[1])[0],size))
print 'Link this fourier mask to your data folder for each particle, for example in bash:'
print 'for i in {#####..#####};do ln %s data/pfmask_${i}.em;done'
print 'where #####..##### is the particle number range for the given tomogram.'