#!/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/>.

from subprocess import check_call
from os.path import split, splitext, join
import os

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser("script to convert full GENIE root files to the reduced gst ROOT format")
    parser.add_argument("filenames", nargs="+", help="GENIE root files")
    parser.add_argument("--dest", required=True, help="destination directory")
    args = parser.parse_args()

    for filename in args.filenames:
        head, tail = split(filename)
        root, ext = splitext(tail)
        output = join(args.dest, root) + ".ntuple.root"
        cmd = ["gntpc","-f","gst","-i",filename,"-o",output]
        print(" ".join(cmd))
        wi
#!/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

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_sphere2(R):
    """
    Generates a random point inside a sphere of radius R.
    """
    while True:
        pos = np.random.rand(3)*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` 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_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
    from itertools import islice
    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=100.0,
                        help="total 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 < 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,args.E,m1,m2),args.n):
            pos = rand_sphere2(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)))