#!/usr/bin/env python

# ---- Import standard modules to the python path.
import argparse

import os

import numpy as np

import astropy.units as u
import astropy.constants as C

from scipy.integrate import ode
from scipy.stats import maxwell
from scipy.stats import rv_continuous
from scipy.integrate import quad

from matplotlib import use
use('agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

import pandas as pd

import h5py

import pdb

from astro_traj import __version__
from astro_traj.system import System
from astro_traj.galaxy import Hernquist_NFW
from astro_traj import constr_dict
from astro_traj.sample import Sample
from astro_traj.plot import Plot


def parse_commandline():
    """Parse the arguments given on the command-line.
    """
    parser = argparse.ArgumentParser(description="""
    Main function. Samples parameters, implements kicks, and sees whether the systems is a 'successful' merger
    by merging at the proper projected distance.

    Writes to csv file in following format:
    [M2, Mns, Mhe, Apre, Apost, epre, epost, R, galcosth, galphi, Vkick, Tmerge, Rmerge, Rmerge_proj, Vfinal, flag]

    Flag description:
    0: Binary doesn't merge at the proper projected distance
    1: Binary mergers at the correct projected distance
    2: Binary's inspiral time is greater than 10 Gyr, does not evolve
    3: Binary is disrupted by SN kick
    4: Energy conservation requirement not met!!!
    """)
    parser.add_argument('-V', '--version', action='version', version=__version__)
    parser.add_argument('--galaxy', help="name of galaxy", required=True)
    parser.add_argument('--offset', type=float, help="offset from center of galaxy of  event", required=True)
    parser.add_argument('--effective-radius', type=float, help="effective radius", required=True)
    parser.add_argument('-s', '--filename-samples', help="Path to posterior samples", required=True)
    parser.add_argument('-t', '--telescope', help="Name of telescope", required=True)
    parser.add_argument('-T', '--trials', type=int, default=1, help="Specify the number of trials. Default=1")
    parser.add_argument('-f', '--outfile', type=str, default='output', help="Name of the output file. Default=output")
    args = parser.parse_args()

    return args

def main(samples, galaxy_name, offset, r_eff, telescope):
    """
    Main function. Samples parameters, implements kicks, and sees whether the systems is a 'successful' merger
    by merging at the proper projected distance.

    Writes to csv file in following format:
    [M2, Mns, Mhe, Apre, Apost, epre, epost, R, galcosth, galphi, Vkick, Tmerge, Rmerge, Rmerge_proj, Vfinal, flag]

    Flag description:
    0: Binary doesn't merge at the proper projected distance
    1: Binary mergers at the correct projected distance
    2: Binary's inspiral time is greater than 10 Gyr, does not evolve
    3: Binary is disrupted by SN kick
    4: Energy conservation requirement not met!!!
    """

    # set cosmology to Reiss [h = H0/(100 km/s/Mpc)]
    h = 0.73
    # Info from "GW"
    GW = constr_dict.GW(samples)
    # Info about the galaxy
    Galaxy = constr_dict.galaxy(galaxy_name, samples, r_eff, offset, h)
    # Infer about the telescope that made the measurements (for angular resolution)
    tele = constr_dict.telescope(telescope)

    # Calculate angular resolution of the telescope, convert to physical size at the distance of the kilonova
    theta = 1.22 * 1e-9 * tele['lambda'] / tele['D']
    D_theta = Galaxy['d']*np.tan(theta)                # units of Mpc
    Galaxy['offset_uncer'] = D_theta*1000.0            # offset uncertainty due to angular resolution of telescope (kpc)


    # Initialize potential with galactic parameters, choose from one of the definied potentials in galaxy class
    gal=Hernquist_NFW(Galaxy['Mspiral'], Galaxy['Mbulge'], Galaxy['Mhalo'], Galaxy['R_eff'], h, rcut=100)
    samp=Sample(gal)
    # Initialize R
    PDF = samp.initialize_R()


    # Create empty output file
    df = pd.DataFrame([], columns=['M2', 'Mns', 'Mhe', 'Apre', 'Apost', 'epre', 'epost', 'R', 'galcosth', 'galphi', 'Vkick', 'Tmerge', 'Rmerge', 'Rmerge_proj', 'Vfinal', 'flag'])
    df.to_csv(args.outfile+'.dat', sep=',', index=False)



    #########################
    ####### MAIN LOOP #######
    #########################
    Nsys=args.trials
    dEfrac = 0.0

    # Initialize random draws for some parameters based on number of trials

    # Msun
    Mcomp_dist, Mns_dist = samp.sample_masses(GW['m1'], GW['m2'], GW['m1_sigma'], GW['m2_sigma'], size=Nsys)

    # Rsun
    Apre_dist = samp.sample_Apre(Amin=0.1, Amax=10.0, size=Nsys)

    # Msun
    Mhe_dist = samp.sample_Mhe(Mmin=Mns_dist)

    # km/s
    Vkick_dist = samp.sample_Vkick_maxwellian(size=Nsys)

    # kpc
    R_dist = samp.sample_R(PDF, Nsys)

    for R, Mcomp, Mns, Apre, Mhe, Vkick in zip(R_dist, Mcomp_dist, Mns_dist, Apre_dist, Mhe_dist, Vkick_dist):

        # Sample necessary parameters
        epre = 0.0                                  # assume binary is circularized before SN2

        # initialize System class with pertinent parameters
        T=System(gal, R, Mns, Mcomp, Mhe, Apre, epre, Vkick)
        # implement a SN kick to the system, and check to see if binary is disrupted
        T.SN()
        # if binary is disrupted, move on to next system
        if T.flag == 3:
            continue

        # choose random location on sphere of radius R
        T.setXYZ_0()
        # choose random direction for circular velocity, and add SN-imparted velocity to get V0
        T.setVxyz_0()

        # set merger time as tmax for trajectory integration
        T.setTmerge()
        # if binary does not meet inspiral time requirements, move on to next system
        if T.flag == 2:
            outdata = T.write_data()
            f = open(args.outfile+'.dat', 'a')
            f.write(str(outdata)[1:-1]+'\n')
            f.close()
            continue

        # integrate trajectory until Tmerge
        T.doMotion()

        # now, see if a successful binary for replicating GW has been created!
        T.check_success(Galaxy['offset'], uncer=Galaxy['offset_uncer'])

        # check for energy conservation, and hold onto highest offset
        T.energy_check()
        if T.dEfrac > dEfrac:
            dEfrac = T.dEfrac


        # hold onto output data if successful binary is produced
        outdata = T.write_data()
        f = open(args.outfile+'.dat', 'a')
        f.write(str(outdata)[1:-1]+'\n')
        f.close()


    ########################
    ####### END LOOP #######
    ########################


    # print maximum energy difference to check for conservation of energy
    print 'Maximum energy difference between initial post-SN energy and final energy: %.20f' % dEfrac

if __name__ == '__main__':
    args = parse_commandline()
    main(args.filename_samples, args.galaxy, args.offset, args.effective_radius, args.telescope)
