"""
   Initialize two stars to a certain age and merge them using MMAMS
"""
import sys
import numpy
from amuse.lab import *
from amuse.plot import plot, xlabel, ylabel
from matplotlib import pyplot 

###BOOKLISTSTART###
def merge_two_stars(Mprim, Msec, t_coll):
    bodies = Particles(mass=[Mprim, Msec] |units.MSun)
        
    stellar = MESA()
    primary = stellar.particles.add_particles(bodies[0].as_set())
    secondary = stellar.particles.add_particles(bodies[1].as_set())

    stellar.evolve_model(t_coll)

    stellar.merge_colliding(primary.copy(), secondary.copy(),
                            MMAMS, return_merge_products=["se"])
    radius = stellar.particles[0].get_radius_profile()
    rho    = stellar.particles[0].get_density_profile()
    stellar.stop()
    return radius, rho
###BOOKLISTSTOP###

def plot_density_profile(radius, rho):
    plot(radius.in_(units.RSun), rho)
    pyplot.xlabel("$R$ [$R_\odot$]")
    pyplot.ylabel("density [$g/cm^3$]")
    pyplot.semilogy()
    pyplot.show()

def new_option_parser():
    from amuse.units.optparse import OptionParser
    result = OptionParser()
    result.add_option("-M", dest="Mprim", type="float",default = 5,
                      help="Mass of the primary star [%default] MSun")
    result.add_option("-m", dest="Msec", type="float",default = 3,
                      help="Mass of the secondary star [%default] MSun")
    result.add_option("-t", unit=units.Myr, dest="t_coll",
                      type="float", default = 1.0|units.Myr,
                      help="end time of the simulation [%default] Myr")
    return result

if __name__ in ('__main__', '__plot__'):
    o, arguments  = new_option_parser().parse_args()
    radius, rho = merge_two_stars(**o.__dict__)
    plot_density_profile(radius, rho)
