#!/bin/env python3

from os import getenv, system
from sys import argv, exit
from os.path import expanduser

program = 'runvdifsim'
version = 0.2
verdate = '20230619'
author = 'Walter Brisken <wbrisken@nrao.edu>'

force = False

dotFile = expanduser('~/.vdifsim')

def usage():
	print('\n%s ver. %s  %s  %s\n' % (program, version, verdate, author))
	print('Usage: %s <input file 1> [ <input file 2> [ ... ] ]\n' % program)
	print('This program spawns instances of vdifsim across multiple computers')
	print('to parallelize generation of synthetic VDIF baseband data using')
	print('vdifsim.  The input files should have corresponding .calc and .im')
	print('files and should have one filename associated with each antenna')
	print('for each scan.  Only VDIF data format is supported.\n')

def loadDotFile():
    try:
        data = open(dotFile, 'r').readlines()
        print('Found %s' % dotFile)
        machines = []
        outPaths = {}
        for d in data:
            s = d.strip().split('#')[0].split()
            if len(s) == 1:
                machines.append(s[0])
            elif len(s) >= 2:
                p = s[1]
                if p[-1] != '/':
                    p += '/'
                outPaths[s[0].upper()] = p
        return machines, outPaths
    except IOError:
        print('Note: %s not found' % dotFile)
        return None, None

# ----------------------------

if len(argv) < 2 or argv[1] in ['-h', '--help']:
	usage()
	exit(0)

machines, outPaths = loadDotFile()

if machines == None:
	print('Error: cannot load %s .\n' % dotFile)
	exit(0)

nFile = 0

for a in argv[1:]:
	if a[-6:] == '.input':
		nFile += 1
	
hosts = ','.join(machines[:nFile])
args = ' '.join(argv[1:])

label = getenv('DIFX_LABEL')
if label == None:
	cmd = 'mpirun --host %s vdifsim %s' % (hosts, args)
	print('\nDIFX_LABEL env. var. must be defined for automatic running.')
	print('You might be able to start the process manually by running:\n')
	print('  %s\n')
	print('Note: you may need to adjust the above command to match your environment.\n')
	exit(0)

else:
	cmd = 'mpirun --host %s rungeneric.%s vdifsim %s' % (hosts, label, args)
	print('Executing: %s' % cmd)
	system(cmd)