#!/usr/bin/env python3

###########################################################################################
#  package:   pNbody
#  file:      sw_MoveICsOnOrbit
#  brief:     
#  copyright: GPLv3
#             Copyright (C) 2019 EPFL (Ecole Polytechnique Federale de Lausanne)
#             LASTRO - Laboratory of Astrophysics of EPFL
#  author:    Yves Revaz <yves.revaz@epfl.ch>
#
# This file is part of Gtools.
###########################################################################################


import argparse
from pNbody import *


####################################################################
# option parser
####################################################################

description="move a swift ICs file to an orbit"
epilog     ="""
Examples:
--------
sw_MoveICsToOrbit snap.hdf5  --boxsize 1000 --position 0 0.1 50  --velocity  10  0  0  -o snap.hdf5
"""

class store_as_array(argparse._StoreAction):
    """Allows to pass arrays as arguments to the parser."""

    def __call__(self, parser, namespace, values, option_string=None):
        values = np.array(values)
        return super().__call__(parser, namespace, values, option_string)


parser = argparse.ArgumentParser(description=description,epilog=epilog,formatter_class=argparse.RawDescriptionHelpFormatter)


parser.add_argument(action="store", 
                    dest="file", 
                    metavar='FILE', 
                    type=str,
                    default=None,
                    nargs='*',
                    help='a file name') 

parser.add_argument("-o",
                    action="store",
                    type=str,
                    dest="outputfilename",
                    default=None,
                    help="Name of the output file")  


parser.add_argument("--position", default=None, action=store_as_array,
                    type=float, nargs=3, help="Position from where the "
                    "objects are launched for forward integration. "
                    "Carthesian coordinates in [kpc].")

parser.add_argument("--velocity", default=None, action=store_as_array,
                    type=float, nargs=3, help="Velocity at which the "
                    "objects are launched for forward integration. "
                    "Carthesian coordinates in [km/s].")


parser.add_argument('--boxsize', default=None, type=float,
                        help="Size of the box.")

####################################################################
# main
####################################################################


if __name__ == '__main__':
  
  opt = parser.parse_args()

  nb = Nbody(opt.file)
  
  
  if opt.position is not None:
    nb.pos = nb.pos + opt.position

  if opt.velocity is not None:
    nb.vel = nb.vel + opt.velocity  
  
  if opt.boxsize is not None:
    nb.boxsize = opt.boxsize   
  

  if opt.outputfilename:
    nb.rename(opt.outputfilename)
    nb.write()
