#!/usr/bin/env python
from __future__ import print_function

import argparse
import os
import shutil
import subprocess

# This script simply runs TraceWin CLI version.
# Provides a better CLI than going through the manual
# Spelling corrections of commands..

parser = argparse.ArgumentParser(description="Creates a TraceWin batch command")

parser.add_argument('-p', '--project', dest='project',
        help="Project file", required=True)
parser.add_argument('-r', dest='run', action='store_const',
        help="Run TraceWin (default just print command)", const=True, default=False)
# path_cal
parser.add_argument('-o', '--output', dest='outpath',
        help="Path to calculation directory")
# dat_file
parser.add_argument('-l', '--lattice', dest='lattice',
        help="Lattice/structure file (.dat) to use")
# dst_file
parser.add_argument('--dst_file', dest='dst_file1',
        help="Input distribution file of main beam")
parser.add_argument('--dst_file2', dest='dst_file2',
        help="Input distribution file of secondary beam")
# current
parser.add_argument('-c', '--current', dest='current1',
        help="Beam Current of beam 1 [mA]")
parser.add_argument('--current2', dest='current2',
        help="Beam Current of beam 2 [mA]")
# nbr_part
parser.add_argument('-np', dest='nbr_part1',
        help="Number of particles in beam 1")
parser.add_argument('-np2', dest='nbr_part2',
        help="Number of particles in beam 2")
# energy
parser.add_argument('-e', '--energy', dest='energy1',
        help="Input beam energy for beam 1 [MeV]")
parser.add_argument('-e2', dest='energy2',
        help="Input beam energy for beam 2 [MeV]")
# freq 
parser.add_argument('-f', dest='freq1',
        help="Bunch Frequency of beam 1 [MHz]")
parser.add_argument('-f2', dest='freq2',
        help="Bunch Frequency of beam 2 [MHz]")

parser.add_argument('-s', '--seed', dest='seed',
        help="Initial random seed", type=int)

args = parser.parse_args()

# A list to be filled depending on arguments given...
binary = None
for c in ["TraceWin64_noX11", "TraceWin_cli"]:
    if shutil.which(c):
        binary = c
        break
if binary is None:
    raise ValueError("Could not find the TraceWin binary")
cmd = [binary, f"'{args.project}'", "hide_esc"]

# Check that project path exist:
if not os.path.isfile(args.project):
    raise ValueError('Project file does not exist')

if args.lattice:
    cmd.append("dat_file="+args.lattice)

for arg in ['dst_file1', 'dst_file2',
            'current1', 'current2',
            'nbr_part1', 'nbr_part2',
            'energy1', 'energy2',
            'freq1', 'freq2']:
    attr=getattr(args,arg)
    if attr:
        cmd.append(arg+"="+attr)

if args.outpath:
    if not os.path.isdir(args.outpath):
        print("WARNING: output directory does not exist, creating..")
        os.makedirs(args.outpath)
    cmd.append("path_cal='"+os.path.abspath(args.outpath)+"'")
if args.seed:
    cmd.append("random_seed="+str(args.seed))

print(' '.join(cmd))
if args.run:
    # TraceWin gives error when using subprocess.call (both for shell equals True and False)
    #subprocess.call(cmd)
    os.system(" ".join(cmd))