#!/usr/bin/env python

################################################################################
###                         parse script parameters                          ###
################################################################################
from optparse import OptionParser

usage = "usage: %prog [options] [<INPUT>] [<STEP_SIZE>] [<ITERATIONS>]"
arg_parser = OptionParser(usage = usage)
arg_parser.add_option("-i", "--input", action = "store", type = str,
    dest = "input", default = None,
    help = "input file path (required!)")
arg_parser.add_option("-t", "--step_size", action = "store", type = float,
    dest = "step_size", default = -1.0,
    help = "time step size (Delta t) for integration (required!)")
arg_parser.add_option("-N", "--iterations", action = "store", type = int,
    dest = "iterations", default = -1,
    help = "Number of time steps used for Newton's equation integration (required!)")
arg_parser.add_option("-o", "--output", action = "store", type = str,
    dest = "output", default = "task03.xyz",
    help = "output file path (default: 'task03.xyz')")
arg_parser.add_option("-s", "--subset", action = "store", type = int,
    dest = "subset", default = 100,
    help = "only stores every subset state to output file (default: 100)")
arg_parser.add_option("-v", action = "store_true",
    dest = "verbose", default = False,
    help = "turn verbosity mode on (default: off a.k.a. silent)")
# Parse command line arguments (as def. above) or store defaults to `config`
config, args = arg_parser.parse_args()
# overwrite options with positional arguments if supplied
try:
    if len(args) > 0:
        config.input = args[0]
    if len(args) > 1:
        config.step_size = float(args[1])
    if len(args) > 2:
        config.iterations = int(args[2])
except ValueError as expression:
    arg_parser.print_help()
    print(f"Error: {expression}")
    exit(-1)
else:
    # quick and dirty validation
    if not config.step_size > 0.0 \
    or not config.iterations > 0 \
    or config.input is None:
        arg_parser.print_help()
        print("Error: missing or illegal argument")
        exit(-1)

################################################################################
###                      task 3 / trajectory generation                      ###
################################################################################
# note, load module _after_ processing script parameters (no need to load all
# of the heavy numeric modules if only usage or alike is needed)
from molecular_dynamics import load, dump, force, mass, step, print_prog_bar

# load initial state from file
(position, velocity), box_size = load(config.input)

# compute initial acceleration using Newton’s second law of motion
acceleration = force(position, box_size) / mass

# iterate time steps for system time evolution
with open(config.output, "w") as output:
    # store initial state to file
    dump(output, position, velocity, box_size,
         comment = f"Initial State with step size {config.step_size}")
    # perform iterations many time steps
    for iteration in range(1, config.iterations):
        # perform a single time step
        position, velocity, acceleration = step(position, velocity, acceleration,
                                                box_size, config.step_size)
        # dump current state to file
        if (iteration % config.subset) == 0:
            dump(output, position, velocity, box_size,
                 comment = f"Iteration: {iteration}")
        # if verbosity turned on, print a progress bar
        if config.verbose:
            print_prog_bar(iteration, config.iterations)
    # In case of final iteration NOT writen to file yet, write also final state
    if (iteration % config.subset) != 0:
        dump(output, position, velocity, box_size,
             comment = f"Iteration: {iteration}")


# report output file path
print(f"Done: trajectories saved to '{config.output}'")