#########################################################################
#
# Script: generate_ams_spectra.py
#
# Author: Igor Volobouev, March 2005
#
# Purpose: Generates galactic protons or He ions according to the AMS
#          spectra for a typical GLAST orbit using the 6m^2 round patch.
#          Note that this code assumes that the cutoff rigidity maps
#          have been generated for the altitude simular to the orbit
#          altitude.
#
#########################################################################

"""
Usage: generate_ams_spectra particle tmin_sec tmax_sec orbit_file output_file
"""

import random_seed
import random
import sys
import math

from cr_sources import *
from source_utils import *
from ascii_orbit import *
from timerep import *
from standard_header import *

import coords
import particle_info
import filters


def generate_ams(particle, tmin_sec, tmax_sec, orbit):
    # Some parameters.
    # Patch area in square meters:
    patch_area = 6.0
    max_impact_param = 1000.0*math.sqrt(patch_area/math.pi)
    
    # Source backoff distance in mm:
    backoff_distance = 3000.0
    
    # Flux for the full primary proton spectrum in the rigidity
    # range to which AMS is sensitive. It is determined by simply
    # adding all differential fluxes presented in the AMS papers.
    # It is converted from units in the paper into particles/(mm^2 sec sr)
    primary_flux = {"proton":2.7055e-3, "He":0.2404e-3}
    
    # Create the cosmic ray generator.
    # Unfiltered arrival times will be uniformly distributed.
    timeGen = UniformLightCurve(Time(tmin_sec), Time(tmax_sec))

    # The particles will be affected by the Earth occultation
    # and by the rigidity cutoff. Import of the rigidity module
    # was delayed until this point because it is slow.
    filter = list()
    filter.append(filters.EarthOccultationFilter("occultation", orbit))
    import rigidity
    filter.append(rigidity.RigidityFilter("rigidity", orbit))
    
    # We will use the AMS proton spectrum with the minimum energy
    # determined from the smallest rigidity cutoff.
    min_cutoff = rigidity.min_cutoff()
    pmass = particle_info.info[particle].mass
    pmin = min_cutoff*abs(particle_info.info[particle].charge)
    min_energy = math.sqrt(pmin*pmin + pmass*pmass)

    # Figure out which spectrum generator to use
    spectrum = {"proton":AMSGalacticProtonSpectrum, "He":AMSHeliumSpectrum}
    
    # Make a uniform and isotropic source producing this spectrum.
    # The tag in the label will be mapped to a unique integer source
    # number used by Gleam if the tag is present in the "source_id"
    # dictionary defined in the EventGenUtils_config.py configuration script.
    label = CosmicRayLabel("AMS" + particle, filter)
    source = IsotropicTubeSource(label, coords.CELESTIAL, particle,
                                 spectrum[particle](min_energy),
                                 max_impact_param, backoff_distance)
    
    # Figure out how many events we should generate.
    # We need to take the rigidity cutoff into account.
    fraction = 1.0 - spectrum[particle](0.0).cdf(min_energy)
    mean_events = source.area()*source.solidAngle()*(tmax_sec - tmin_sec)*\
                  primary_flux[particle]*fraction
    # There is no Poisson random number generator in the "random"
    # module. Approximate by the Gaussian distribution.
    nevents = 0
    while (nevents <= 0):
        nevents = int(random.gauss(mean_events, math.sqrt(mean_events)) + 0.5)

    # Generate the event sequence.
    events = EventSequence()
    print "Generating", nevents, particle, "events"
    events.grow(SimpleEventGenerator(source, timeGen), nevents)
    return events


def main(argv=None):
    # parse the command line
    if argv is None:
        argv = sys.argv
    argc = len(argv)-1
    if (argc == 0):
        print __doc__
        return 0
    elif (argc == 5):
        i = 1
        particle = argv[i]; i+=1
        tmin_sec = float(argv[i]); i+=1
        tmax_sec = float(argv[i]); i+=1
        orbit_file = argv[i]; i+=1
        output_file = argv[i]; i+=1
    else:
        print __doc__
        return 1

    if (tmin_sec >= tmax_sec):
        print >> sys.stderr, "Bad time limits"
        return 1

    # Check particle name
    if ("proton" != particle and "He" != particle):
        print >> sys.stderr, 'Bad particle name, must be "proton" or "He"'
        return 1

    # Read in the orbit information
    orbit = AsciiOrbit()
    forb = open(orbit_file, 'r')
    orbit.read(forb)
    forb.close()
    
    # Open the output file now so that we immediately
    # get an error if the file name is bogus.
    outfile = open(output_file, 'w')
    try:
        # Generate the event sequence
        events = generate_ams(particle, tmin_sec, tmax_sec, orbit)

        # Filter the event sequence
        print "Applying filters"
        nfiltered = events.applyFilters()

        # Write events out in local coordinates.
        # By default, filtered events are not written.
        events.header = StandardHeader(output_file,
                                       events[0].rays[0].label.tag).header()
        nout = events.write(outfile, orbit)
        assert(nfiltered + nout == len(events))
    finally:
        outfile.close()

    # We are done
    print "Wrote", nout, "events to file", output_file
    return 0


if __name__=='__main__':
    sys.exit(main())