import os
import shutil
import subprocess
import threading
import time

###  EDIT DATASET INFORMATION  ###
apix = 0.729
kv = 300
cs = 2.7
##################################


### EDIT COMPUTING INFORMATION ###
Group_Size = 22
batch_Size = 16
mpi = 16
threads = 4
Data_Path ='/mnt/storage_backup/clients/2023-06-28_Paloma/'
Star = 'movies.star'
##################################

InDir = os.listdir(Data_Path) ; #print(InDir)
Processed = os.listdir(os.getcwd())

InDir = [element for element in InDir if '.eer' in element] ; print(len(InDir), 'files to be processed')
Processed = [element.replace('.tif','.eer') for element in Processed if '.tif' in element] ; print('but', len(Processed), 'files already processed')

InDir = [element for element in InDir if element not in Processed] ; print(len(InDir))

steps = [i for i in range(0,len(InDir),batch_Size)]

groups = [InDir[i:i+batch_Size] for i in steps]


for group in groups:

        print('downloading a group of', len(group), 'movies...')

        for element in group: subprocess.call(['cp', Data_Path+element, '.'])

        subprocess.call(['relion_import', '--do_movies', '--i', '*.eer', '--odir', './', '--ofile', Star, '--continue', '--angpix', str(apix), '--kV', str(kv), '--Cs', str(cs)])

        subprocess.run(['mpirun', '-n', str(mpi), 'relion_convert_to_tiff_mpi', '--i', Star, ' --o', './', '--j', str(threads) ,'--eer_grouping', str(Group_Size), '--compression', 'lzw', '--only_do_unfinished'])

        print('now deleting eer files ...')

        for element in group: subprocess.call(['rm', element])

        subprocess.call(['rm', Star])