#!/usr/bin/env python
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.
"""
Script to produce MCPL files for simulating self destructing dark matter. The
arguments to the script are the dark matter mediator mass and energy and the
particle IDs of the two decay products. For example, to simulate a dark matter
mediator with a mass of 100 MeV and total energy 1 GeV decaying to an electron
and a positron:

    $ ./gen-dark-matter -M 100.0 -E 1000.0 -p1 20 -p2 21 -o output

which will create the directory "output_[UUID]" and produce three files:

    1. A MCPL file output.mcpl
    2. A MCPL header file called mcpl_header.dat
    3. A SNOMAN command file which can be used to simulate the dark matter

To simulate it with SNOMAN you can run:

    $ cd output_[UUID]
    $ snoman.exe -c output.cmd
"""

from __future__ import print_function, division
import numpy as np
import string
import uuid
from itertools import islice

SNOMAN_MASS = {
    20: 0.511,
    21: 0.511,
    22: 105.658,
    23: 105.658
}

PSUP_RADIUS = 840.0 # cm

# generate a UUID to append to all the filenames so that if we run the same job
# twice we don't overwrite the first job
ID = uuid.uuid1()

SNOMAN_TEMPLATE = \
"""
$processor_list 'MCO UCL PRU PCK OUT END'

* set up output file
$output_format $full_ds
$zdab_option $zdab_max_mc
file out 1 @output
* Don't know if this next line is necessary
file MCO 1 ./MC_Atm_Nu_No_Osc_Snoman_Genie10000_X_0_rseed.dat checkpoint=10
file mco 2 @mcpl
$mc_event_rate -0.003189 $per_sec


$prune_mc               $keep
$prune_mcpm             $keep
$prune_mcvx_source      $keep
$prune_mcvx_boundary    $drop
$prune_mcvx_interaction $drop
$prune_mcvx_sink        $drop
$prune_mcvx_pre_source  $drop

$mcrun 10000

* simulate run conditions for run 10000
$mc_gen_run_cond $on

* Don't think this is necessary since it's reset in run_mc_atmospherics
$num_events @num_events

* titles files for run 10000
titles /sno/mcprod/dqxx/DQXX_0000010000.dat
titles /sno/mcprod/anxx/titles/pca/anxx_pca_0000010623_p2.dat
titles /sno/output/anxx/titles/neutrino/10000-10999/anxx_nu_0000010000_p12.dat

* set the average number of noise hits per event
* this comes from the autosno generated MC_Atm_Nu_No_Osc_Snoman_Genie10000_X_1_run_mc_atm_nu_genie.cmd file
set bank TRSP 1 word 2 to 2.051309

* From activate_atmospherics.cmd

$killvx                 7
$killvx_neutron		5
$egs4_ds		$off
$store_full_limit	10
$max_cer_ge_errors      2000

* Enable hadron propagation

titles sno_hadron_list.dat
titles chetc_sno.dat
titles flukaaf_sno.dat
$enable_hadrons         $on

* Load information for muons

titles music_sno_info.dat
titles music_double_diff_rock.dat
titles muon.dat
titles muon_param.dat
titles photo_dis.dat
$enable_music_calc      $off

@load_d2o_settings.cmd
@run_mc_atmospherics
""".strip()

MCPL_HEADER="""
*DO  MCPL 1 -i(30I 2I / 2I 17F) -n2000000
#.
#.  This bank is used to run particles through SNOMAN (with file MCO 2 ... command)
#.  This file was derived from a .dat file of Christian Nally.
#.
#.    Contact:  D. Waller (Carleton)
#.
#.       Standard Database Header
#.
19750101        0 20380517 03331900  #.  1..4   Intrinsic validity
       0        0        0           #.  5..7   Data type, Task type, Format no.
       0        0        0           #.  8..10  Creation Date, Time, Source Id.
19750101        0 20380517 03331900  #. 11..14  Effective validity
       0        0                    #. 15..16  Entry Date Time
4*0                                  #. 17..20  Spare
1000000*0                            #. 21..30  Temporary data (not in database)
#.
#.    End of Standard Database Header
#.
#.    User Data.
"""

class MyTemplate(string.Template):
    delimiter = '@'

def rand_ball(R):
    """
    Generates a random point inside a sphere of radius R.
    """
    while True:
        pos = (np.random.rand(3)*2-1)*R

        if np.linalg.norm(pos) < R:
            break

    return pos

def rand_sphere():
    """
    Generates a random point on the unit sphere.
    """
    u = np.random.rand()
    v = np.random.rand()

    phi = 2*np.pi*u
    theta = np.arccos(2*v-1)

    dir = np.empty(3,dtype=float)

    dir[0] = np.sin(theta)*np.cos(phi)
    dir[1] = np.sin(theta)*np.sin(phi)
    dir[2] = np.cos(theta)

    return dir

def lorentz_boost(x,n,beta):
    """
    Performs a Lorentz boost on the 4-vector `x`. Returns the 4-vector as would
    be seen by a frame moving along the direction `n` at a velocity `beta` (in
    units of the speed of light).

    Formula comes from 46.1 in the PDG "Kinematics" Review. See
    http://pdg.lbl.gov/2014/reviews/rpp2014-rev-kinematics.pdf.
    """
    n = np.asarray(n,dtype=float)
    x = np.asarray(x,dtype=float)

    gamma = 1/np.sqrt(1-beta**2)

    # normalize direction
    n /= np.linalg.norm(n)

    # get magnitude of `x` parallel with the boost direction
    x_parallel = np.dot(x[1:],n)

    # compute the components of `x` perpendicular
    x_perpendicular = x[1:] - n*x_parallel

    E = x[0]
    E_new = gamma*E - gamma*beta*x_parallel
    x_new = -gamma*beta*E + gamma*x_parallel

    x_new = [E_new] + (x_new*n + x_perpendicular).tolist()

    return np.array(x_new)

def gen_decay(M, E, m1, m2):
    """
    Generator yielding a tuple (v1,v2) of 4-vectors for the daughter particles
    of a 2-body decay from a massive particle M with total energy E in the lab
    frame.

    Arguments:

            M  - mass of parent particle (MeV)
            E  - Total energy of the parent particle (MeV)
            m1 - mass of first decay product
            m2 - mass of second decay product
    """
    while True:
        # direction of 1st particle in mediator decay frame
        v = rand_sphere()
        # calculate momentum of each daughter particle
        # FIXME: need to double check my math
        p1 = np.sqrt(((M**2-m1**2-m2**2)**2-4*m1**2*m2**2)/(4*M**2))
        E1 = np.sqrt(m1**2 + p1**2)
        E2 = np.sqrt(m2**2 + p1**2)
        v1 = [E1] + (p1*v).tolist()
        v2 = [E2] + (-p1*v).tolist()
        # random direction of mediator in lab frame
        n = rand_sphere()
        beta = np.sqrt(E**2-M**2)/E
        v1_new = lorentz_boost(v1,n,beta)
        v2_new = lorentz_boost(v2,n,beta)

        yield v1_new, v2_new

def test_gen_decay1():
    """
    A super simple test that if we generate a decay with E = mass, the two
    particles come out back to back.
    """
    for v1, v2 in islice(gen_decay(100,100,0.5,0.5),100):
        dir1 = v1[1:]/np.linalg.norm(v1[1:])
        dir2 = v2[1:]/np.linalg.norm(v2[1:])
        assert np.isclose(np.dot(dir1,dir2),-1.0)

def test_gen_decay2():
    """
    A super simple test that if we generate a decay with E >> mass, the two
    particles come out in the same direction.
    """
    for v1, v2 in islice(gen_decay(100,100e3,0.5,0.5),100):
        dir1 = v1[1:]/np.linalg.norm(v1[1:])
        dir2 = v2[1:]/np.linalg.norm(v2[1:])
        assert np.isclose(np.dot(dir1,dir2),1.0,atol=1e-3)

def test_lorentz_boost2():
    """
    Simple test of lorentz_boost() using a problem from
    http://electron6.phys.utk.edu/PhysicsProblems/Mechanics/8-Relativity/decay%20massless.html.
    """
    M = 140.0
    P = 2e3
    E = np.sqrt(P**2 + M**2)
    m1 = 105.0
    m2 = 0.0
    # muon in lab frame is going in +z direction which means that the direction
    # of the muon in the pion decay frame is also in the +z direction
    v = np.array([0,0,1])
    # calculate momentum of each daughter particle
    # FIXME: need to double check my math
    p1 = np.sqrt(((M**2-m1**2-m2**2)**2-4*m1**2*m2**2)/(4*M**2))
    E1 = np.sqrt(m1**2 + p1**2)

    if not np.isclose(E1,109.375):
        print("Energy of muon in pion rest frame = %.2e but expected %.2e" % (E1,109.375))

    assert np.isclose(E1,109.375)

    E2 = np.sqrt(m2**2 + p1**2)
    v1 = [E1] + (p1*v).tolist()
    v2 = [E2] + (-p1*v).tolist()
    # pion in lab frame is going in the +z direction, therefore the lab frame
    # is going in the -z direction relative to the pion frame
    n = np.array([0,0,-1])
    beta = np.sqrt(E**2-M**2)/E

    if not np.isclose(beta,0.99756,rtol=1e-3):
        print("Speed of pion frame relative to lab frame = %.5f but expected %.5f" % (beta,0.99756))

    assert np.isclose(beta,0.99756,rtol=1e-3)

    v1_new = lorentz_boost(v1,n,beta)
    v2_new = lorentz_boost(v2,n,beta)

    if not np.isclose(v1_new[0],2003.8,atol=1e-2):
        print("Energy of muon in lab frame = %.1f but expected %.1f" % (v1_new[0],2003.8))

    assert np.isclose(v1_new[0],2003.8,atol=1e-2)

def test_lorentz_boost():
    """
    Super simple test of the lorentz_boost() function to make sure that if we
    boost forward by a speed equal to the particle's original speed, it's
    4-vector looks like [mass,0,0,0].
    """
    m1 = 100.0
    beta = 0.5
    p1 = m1*beta/np.sqrt(1-beta**2)
    p1 = [np.sqrt(m1**2 + p1**2)] + [p1,0,0]

    p1_new = lorentz_boost(p1,[1,0,0],0.5)

    assert np.allclose(p1_new,[m1,0,0,0])

if __name__ == '__main__':
    import argparse
    import sys
    import os

    parser = argparse.ArgumentParser("generate MCPL files for self destructing dark matter")
    parser.add_argument("-M", type=float, default=100.0,
                        help="mass of mediator")
    parser.add_argument("-E", type=float, default=None,
                        help="total energy of mediator")
    parser.add_argument("-T", type=float, default=0.0,
                        help="kinetic energy of mediator")
    parser.add_argument("-p1", type=int, default=20,
                        help="SNOMAN particle ID for 1st decay product")
    parser.add_argument("-p2", type=int, default=21,
                        help="SNOMAN particle ID for 2nd decay product")
    parser.add_argument("-n", type=int, default=100,
                        help="number of events to generate")
    parser.add_argument("-o", "--output", type=str, required=True,
                        help="output prefix")
    args = parser.parse_args()

    if args.p1 not in SNOMAN_MASS:
        print("%i is not a valid particle ID" % args.p1,file=sys.stderr)
        sys.exit(1)

    m1 = SNOMAN_MASS[args.p1]

    if args.p2 not in SNOMAN_MASS:
        print("%i is not a valid particle ID" % args.p2,file=sys.stderr)
        sys.exit(1)

    m2 = SNOMAN_MASS[args.p2]

    if args.M < m1 + m2:
        print("mediator mass must be greater than sum of decay product masses",file=sys.stderr)
        sys.exit(1)

    if args.E is not None:
        E = args.E
    else:
        E = args.T + args.M

    if E < args.M:
        print("mediator energy must be greater than or equal to the mass",file=sys.stderr)
        sys.exit(1)

    if args.n < 0:
        print("number of events must be positive",file=sys.stderr)
        sys.exit(1)

    # The format for the MCPL files as read in by SNOMAN is:
    #
    #     Word    Type     Description
    #     ----    ----     --------------------------------
    #       1      I       Number of events in the list.
    #       2      I       Number of words per track. (This may vary because
    #                      polarisation is included only for Cerenkov photons).
    #      j+1     I       Number of particles in this event.     
    #      j+2     I       Particle type of this particle.
    #      j+3   F,F,F     X,Y,Z of particle.
    #      j+6     F       Time of particle.
    #      j+7     F       Energy of particle.
    #      j+8   F,F,F     U,V,W of particle.
    #      j+11  F,F,F     x,y,z components of the particle's polarisation.

    
    new_dir = "%s_%s" % (args.output,ID.hex)

    os.mkdir(new_dir)
    os.chdir(new_dir)

    mcpl_filename = args.output + ".mcpl"

    with open(mcpl_filename, "w") as f:
        f.write("%i %i\n" % (args.n, 10))

        for v1, v2 in islice(gen_decay(args.M,E,m1,m2),args.n):
            pos = rand_ball(PSUP_RADIUS)
            p1 = np.linalg.norm(v1[1:])
            p2 = np.linalg.norm(v2[1:])
            f.write("    %i %i %f %f %f %f %f %f %f %f\n" % (2,args.p1,pos[0], pos[1], pos[2],0.0, v1[0], v1[1]/p1, v1[2]/p1, v1[3]/p1))
            f.write("    %i %i %f %f %f %f %f %f %f %f\n" % (2,args.p2,pos[0], pos[1], pos[2],0.0, v2[0], v2[1]/p2, v2[2]/p2, v2[3]/p2))

    # Write out mcpl_header.dat which is needed by the cmd file
    with open("mcpl_header.dat", "w") as f:
        f.write(MCPL_HEADER)

    template = MyTemplate(SNOMAN_TEMPLATE)

    mcds_filename = args.output + ".mcds"
    cmd_filename = args.output + ".cmd"

    with open(cmd_filename, "w") as f:
        f.write(template.safe_substitute(output=mcds_filename, mcpl=mcpl_filename, num_events=str(args.n)))