#!/bin/env python

import collections
import getopt
import numpy
import os
import random
import sys
import unittest 
from time import process_time as cputime
from time import time as wallclocktime

from amuse.community.ph4.interface import ph4 as grav
from amuse.units import nbody_system
from amuse.units import units
from amuse.units.quantities import zero

from amuse import datamodel
from amuse.datamodel import particle_attributes as pa
from amuse.rfi.core import is_mpd_running
from amuse.ic.plummer import new_plummer_model
from amuse.ic.salpeter import new_salpeter_mass_distribution_nbody

def print_log(pre, time, gravity, E0 = 0.0 | nbody_system.energy,
              cpu0 = 0.0, wall0 = 0.0):
    cpu = cputime()
    wall = wallclocktime()
    N = len(gravity.particles)
    M = gravity.total_mass

    print('')
    print(pre+"time=", time.number)
    print(pre+"cpu=", cpu-cpu0)
    print(pre+"wall=", wall-wall0)
    print(pre+"Ntot=", N)
    print(pre+"mass=", M.number)
    E = 0 | nbody_system.energy

    if N > 0:
        U = gravity.potential_energy
        T = gravity.kinetic_energy
        Etop = T + U
        E = Etop
        if E0 == 0 | nbody_system.energy: E0 = E
        Rvir = -0.5*M*M/U
        Q = -T/U
        com = pa.center_of_mass(gravity.particles)
        comv = pa.center_of_mass_velocity(gravity.particles)
        if N > 15:
            dcen,rcore,rhocore \
              = pa.densitycentre_coreradius_coredens(gravity.particles)
        else:
            dcen = com
            rcore = zero
            rhocore = zero
            cmx,cmy,cmz = dcen
        lagr,mf = pa.LagrangianRadii(gravity.particles, cm=dcen)  # no units!

        print(pre+"Etot=", E.number)
        print(pre+"dE/E=", E/E0 - 1)
        print(pre+"Rvir=", Rvir.number)
        print(pre+"Qvir=", Q)
        cmx,cmy,cmz = com
        print(pre+"cmpos[3]= %.8f %.8f %.8f" % (cmx.number, cmy.number, cmz.number))
        cmx,cmy,cmz = comv
        print(pre+"cmvel[3]= %.8f %.8f %.8f" % (cmx.number, cmy.number, cmz.number))
        cmx,cmy,cmz = dcen
        print(pre+"dcpos[3]= %.8f %.8f %.8f" % (cmx.number, cmy.number, cmz.number))
        print(pre+"Rcore=", rcore.number)
        print(pre+"Mlagr[9]=", end=' ')
        for m in mf: print("%.4f" % (m), end=' ')
        print('')
        print(pre+"Rlagr[9]=", end=' ')
        for r in lagr.number: print("%.8f" % (r), end=' ')
        print('')

    sys.stdout.flush()
    return E,cpu,wall

def run_ph4(infile = None, number_of_stars = 40,
             end_time = 10 | nbody_system.time,
             delta_t = 1 | nbody_system.time,
             n_workers = 1, use_gpu = 1, gpu_worker = 1, gpu_id = -1,
             accuracy_parameter = 0.1,
             softening_length = -1 | nbody_system.length,
             manage_encounters = 1):

    if infile != None: print("input file =", infile)
    print("end_time =", end_time.number)
    print("delta_t =", delta_t.number)
    print("n_workers =", n_workers)
    print("use_gpu =", use_gpu)
    print("manage_encounters =", manage_encounters)
    print("\ninitializing the gravity module")
    sys.stdout.flush()

    # Note that there are actually three GPU options to test:
    #
    #	1. use the GPU code and allow GPU use (default)
    #	2. use the GPU code but disable GPU use (-g)
    #	3. use the non-GPU code (-G)

    #print "1"; sys.stdout.flush()
    
    gpu = 0
    if gpu_worker == 1:
        try:
            gravity = grav(number_of_workers = n_workers,
                           redirection = "none", mode = "gpu")
            #              debugger='valgrind')
            gpu = 1
        except Exception as ex:
            print('*** GPU worker code not found. Reverting to non-GPU code. ***')
            gpu = 0

    if gpu == 0:
        gravity = grav(number_of_workers = n_workers,
                       redirection = "none")
        #              debugger='valgrind')

    #print "2"; sys.stdout.flush()
    gravity.initialize_code()

    #print "3"; sys.stdout.flush()
    gravity.parameters.set_defaults()

    gravity.parameters.gpu_id = gpu_id

    #-----------------------------------------------------------------

    #print "4"; sys.stdout.flush()
    if infile == None:

        print("making a Plummer model")
        stars = new_plummer_model(number_of_stars)

        id = numpy.arange(number_of_stars)
        stars.id = id+1

        print("setting particle masses and radii")
        if number_of_stars > 0:
            stars.mass = (1.0 / number_of_stars) | nbody_system.mass
        if 0:
            scaled_mass = new_salpeter_mass_distribution_nbody(number_of_stars) 
            stars.mass = scaled_mass
        stars.radius = 0.0 | nbody_system.length

        print("centering stars")
        stars.move_to_center()
        if 0:
            print("scaling stars to virial equilibrium")
            stars.scale_to_standard(smoothing_length_squared
                                    = gravity.parameters.epsilon_squared)

        time = 0.0 | nbody_system.time
        sys.stdout.flush()

    else:

        # Read the input data.  Units are dynamical.

        print("reading file", infile)

        id = []
        mass = []
        pos = []
        vel = []

        f = open(infile, 'r')
        count = 0
        for line in f:
            if len(line) > 0:
                count += 1
                cols = line.split()
                if count == 1: snap = int(cols[0])
                elif count == 2: number_of_stars = int(cols[0])
                elif count == 3: time = float(cols[0]) | nbody_system.time
                else:
                    if len(cols) >= 8:
                        id.append(int(cols[0]))
                        mass.append(float(cols[1]))
                        pos.append((float(cols[2]),
                                    float(cols[3]), float(cols[4])))
                        vel.append((float(cols[5]),
                                    float(cols[6]), float(cols[7])))
        f.close()

        stars = datamodel.Particles(number_of_stars)
        stars.id = id
        stars.mass = mass | nbody_system.mass
        stars.position = pos | nbody_system.length
        stars.velocity = vel | nbody_system.speed
        stars.radius = 0. | nbody_system.length

    # print "IDs:", stars.id.number
    sys.stdout.flush()

    #-----------------------------------------------------------------

    #print "5"; sys.stdout.flush()
    if softening_length == -1 | nbody_system.length and number_of_stars > 0:
        eps2 = 0.25*(float(number_of_stars))**(-0.666667) \
                | nbody_system.length**2
    else:
        eps2 = softening_length*softening_length

    #print "6"; sys.stdout.flush()
    gravity.parameters.timestep_parameter = accuracy_parameter
    gravity.parameters.epsilon_squared = eps2
    gravity.parameters.use_gpu = use_gpu
    gravity.parameters.manage_encounters = manage_encounters

    print("adding particles")
    # print stars
    sys.stdout.flush()
    gravity.particles.add_particles(stars)
    gravity.commit_particles()

    print('Taking step')
    sys.stdout.flush()
    gravity.evolve_model(time)

    print('')
    print("number_of_stars =", number_of_stars)
    print("evolving to time =", end_time.number, \
          "in steps of", delta_t.number)
    sys.stdout.flush()

    E0,cpu0,wall0 = print_log('', time, gravity)
    
    # Channel to copy values from the code to the set in memory.
    channel = gravity.particles.new_channel_to(stars)

    stopping_condition = gravity.stopping_conditions.collision_detection
    stopping_condition.enable()
    
    while time < end_time:
        time += delta_t
        gravity.evolve_model(time)

        # Ensure that the stars list is consistent with the internal
        # data in the module.

        ls = len(stars)

        # Update the bookkeeping: synchronize stars with the module data.

        try:
            gravity.update_particle_set()
            gravity.particles.synchronize_to(stars)
        except:
            pass
    
        # Copy values from the module to the set in memory.

        channel.copy()
    
        # Copy the index (ID) as used in the module to the id field in
        # memory.  The index is not copied by default, as different
        # codes may have different indices for the same particle and
        # we don't want to overwrite silently.

        channel.copy_attribute("index_in_code", "id")

        if stopping_condition.is_set():
            star1 = stopping_condition.particles(0)[0]
            star2 = stopping_condition.particles(1)[0]
            print('\nstopping condition set at time', \
                gravity.get_time().number,'for:\n')
            print(star1)
            print('')
            print(star2)
            print('')
            raise Exception("no encounter handling")

        if len(stars) != ls:
            if 0:
                print("stars:")
                for s in stars:
                    print(" ", s.id.number, s.mass.number, \
                    s.x.number, s.y.number, s.z.number)
            else:
                print("number of stars =", len(stars))
            sys.stdout.flush()

        print_log('', time, gravity, E0, cpu0, wall0)
        sys.stdout.flush()

    print('')
    gravity.stop()

if __name__ == '__main__':

    infile = None
    N = 250
    t_end = 5.0 | nbody_system.time
    delta_t = 1.0 | nbody_system.time
    n_workers = 1
    use_gpu = 1
    gpu_worker = 1
    gpu_id = -1
    accuracy_parameter = 0.1
    softening_length = -1  | nbody_system.length
    random_seed = -1
    manage_encounters = 1

    try:
        opts, args = getopt.getopt(sys.argv[1:], "a:c:d:e:f:gGi:n:s:t:w:")
    except getopt.GetoptError as err:
        print(str(err))
        sys.exit(1)

    for o, a in opts:
        if o == "-a":
            accuracy_parameter = float(a)
        elif o == "-c":
            manage_encounters = int(a)
        elif o == "-d":
            delta_t = float(a) | nbody_system.time 
        elif o == "-e":
            softening_length = float(a) | nbody_system.length
        elif o == "-f":
            infile = a
        elif o == "-g":
            use_gpu = 0
        elif o == "-G":
            use_gpu = 0
            gpu_worker = 0
        elif o == "-i":
            gpu_id = int(a)
        elif o == "-n":
            N = int(a)
        elif o == "-s":
            random_seed = int(a)
        elif o == "-t":
            t_end = float(a) | nbody_system.time
        elif o == "-w":
            n_workers = int(a)
        else:
            print("unexpected argument", o)

    if random_seed <= 0:
        numpy.random.seed()
        random_seed = numpy.random.randint(1, pow(2,31)-1)
    numpy.random.seed(random_seed)
    print("random seed =", random_seed)


    #os.system('env')


    assert is_mpd_running()
    run_ph4(infile, N, t_end, delta_t, n_workers,
             use_gpu, gpu_worker, gpu_id,
             accuracy_parameter, softening_length,
             manage_encounters)
