#!/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", "--outputA", action = "store", type = str,
    dest = "outputA", default = "task04A.txt",
    help = "first output file path (default: 'task04A.txt')")
arg_parser.add_option("-B", "--outputB", action = "store", type = str,
    dest = "outputB", default = "task04B.txt",
    help = "second output file path (default: 'task04B.txt')")
arg_parser.add_option("-s", "--cdf_sample_count", action = "store", type = int,
    dest = "cdf_sample_count", default = 100,
    help = "Radius interval sample count for CDF estimate (default: 100)")
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 4 / post processing                         ###
################################################################################
# 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
from jax import jit, numpy as jnp
from molecular_dynamics import iter_load, energy, kinetic


# counte the number of states (system snapshots) in the input file
nr_states = 0
# first output file and collect some info used in the next analysis!
with open(config.outputA, "w") as output:
    # iterate over all state snapshots stored in the `config.input` file
    for (position, velocity), box_size in iter_load(config.input):
        nr_states += 1
        # compute potneital and kinetic energy for current snapshot
        e_pot = energy(position, box_size)
        e_kin = kinetic(velocity)
        # energy results
        print(e_pot, e_kin, file = output)
# report file path
print(f"Saved energy change to '{config.outputA}'")

# CDF radius sample points (`box_size` known from last loop)
radii = np.linspace(0, box_size / 2.0, config.cdf_sample_count)[np.newaxis, :]
# radius Commulative Distribution Function estimate
CDF = np.zeros((config.cdf_sample_count, ))

# extract only one of two distance combinations of non-equal particles
lower_tri = np.tril_indices(position.shape[0], k = -1)
# maximum distance two particles (in folded space) can be apart
max_dist = box_size / 2.0

@jit
def loop_body(position, box_size):
    # magic for all particle pair distances
    diff = position[:, jnp.newaxis, :] - position
    diff = diff[lower_tri[0], lower_tri[1], :]
    # fold space in all directions. The `max_dist` is the maximum distance
    # in the folded space untill the distance through the folded space
    # border is smaller than inside the box. 3-axis parallel fold
    diff = jnp.mod(diff + max_dist, box_size) - max_dist
    # count nr particles pairs with distance smaller than r
    dist = jnp.linalg.norm(diff, axis = 1)
    # accumulate per timepoint partial mean estimates of the CDF
    return jnp.mean(dist[:, jnp.newaxis] < radii, axis = 0)

# second iteration to estimate pairwise distance CDF
nr_time_points = 0
for i, ((position, velocity), box_size) in enumerate(iter_load(config.input)):
    # ignore the first 25% (just keep it simple)
    if 4 * i < nr_states:
        continue
    # Count number of actually considured states / time points
    nr_time_points += 1
    # aggregate CDF sum
    CDF += loop_body(position, box_size)
# final division by timesteps as mean over pairs and time.
CDF /= nr_time_points

# store the CDF estimate to the second file
with open(config.outputB, "w") as output:
    for r, cdf in zip(radii.ravel(), CDF.ravel()):
        print(r, cdf, file = output)
# report file path
print(f"Saved radius CDF estimate to '{config.outputB}'")


# import matplotlib.pyplot as plt
# from mpl_toolkits import mplot3d

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

# for (position, velocity), box_size in iter_load(config.input):
#     e_pot = energy(position, box_size)
#     e_kin = kinetic(velocity)
#     print(f"{e_pot:10.5e} {e_kin:10.5e} {(e_pot + e_kin):10.5e}")

#     if config.animate:
#         count += 1
#         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)



# import numpy as np
# import matplotlib.pyplot as plt
# from mpl_toolkits import mplot3d
# from molecular_dynamics import load, energy, kinetic

# (position, velocity), box_size = load("task02.xyz"):

# fig = plt.figure()
# ax = fig.add_subplot(111, projection = "3d")
# ax.cla()
# 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])