#!/usr/bin/env python

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

usage = "usage: %prog [options] [<INPUT>]"
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("-a", "--animate", action = "store_true",
    dest = "animate", default = False,
    help = "creates an trajectory animation of the states in the input file")
# 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]
except ValueError as expression:
    arg_parser.print_help()
    print(f"Error: {expression}")
    exit(-1)
else:
    # quick and dirty validation
    if 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)
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from molecular_dynamics import iter_load, energy, kinetic

if config.animate:
    # new plot with 3D axis
    fig = plt.figure()
    ax = fig.add_subplot(111, projection = "3d")

count = 0
for (position, velocity), box_size in iter_load(config.input):
    count += 1
    e_pot = energy(position, box_size)
    e_kin = kinetic(velocity)
    print(f"{count: >3} - {e_pot:10.5e} {e_kin:10.5e} {(e_pot + e_kin):10.5e}")

    if config.animate:
        ax.cla()
        plt.title(f"time step {count}")
        plt.xlim([0, box_size])
        plt.ylim([0, box_size])
        ax.set_zlim([0, box_size])
        # create 3D position scatter plot
        position = np.mod(position, box_size)
        ax.scatter(position[:, 0], position[:, 1], position[:, 2])
        # # and save to file
        # plt.savefig(config.plot)
        plt.pause(0.01)