#########################################################################
#
# Script: cr_sources.py
#
# Author: Igor Volobouev, Feb 2005
#
# Purpose: contains Monte Carlo generators of various cosmic ray spectra
#          and light curves
#
#########################################################################

"""Contains various cosmic ray sources and spectra as well as an improved
GRB pulse decomposition model."""

import math
import operator
import bisect
import particle_info
import random as rand

from source_utils import *
from basic_kinematics import *
from incomplete_gamma import *
from timerep import *

class FixedEnergySpectrum(BasicSpectrum):
    "Generates a fixed number."
    def __init__(self, energy):
        BasicSpectrum.__init__(self)
        self.energy = energy
    def generate(self, mass):
        if (self.energy <= mass):
            raise ValueError, "particle mass is too high"
        return self.energy


class UniformSpectrum(BasicSpectrum):
    "Generates random numbers uniformly between given limits."
    def __init__(self, lolim, hilim):
        BasicSpectrum.__init__(self)
        if (lolim < 0.0):
            raise ValueError, "can not have negative energy or time limit"
        if (hilim <= lolim):
            raise ValueError, "wrong order of interval boundaries"
        self.lolim  = lolim
        self.hilim  = hilim
    def generate(self, thresh):
        if (self.hilim <= thresh):
            raise ValueError, "lower threshold is too high"
        if (thresh > self.lolim):
            lolim = thresh
        else:
            lolim = self.lolim
        return rand.random(lolim, self.hilim)


class LognormalSpectrum(BasicSpectrum):
    "Generates random numbers according to the lognormal distribution."
    def __init__(self, mean, stdev, skewness):
        BasicSpectrum.__init__(self)
        self.x0, self.mu, self.s = \
                 lognormal_parameters(mean, stdev, skewness)
        self.sign = skewness*1.0/math.fabs(skewness)
    def generate(self, thresh):
        if (self.sign < 0.0 and thresh >= self.x0):
            raise ValueError, "lower threshold is too high"
        e = thresh - 1.0
        while (e <= thresh):
            e = self.sign*(rand.lognormvariate(self.mu, self.s) + self.x0)
        return e


class EGRETBetaSpectrum(LognormalSpectrum):
    "Random high energy power-law spectral indices compatible with EGRET data"
    #
    # This model comes from the 15 EGRET GRBs presented at the 4th Huntsville
    # GRB Symposium (J.R. Catelli et al., AIP 428, 311, 1998). I have fitted
    # the spectral indices (unbinned maximum likelihood, errors not used)
    # using the lognormal density which seems to provide a reasonable
    # description of the distribution.
    #
    def __init__(self):
        LognormalSpectrum.__init__(self, 2.25083, 0.4871, 3.07227)


class PowerLawSpectrum(BasicSpectrum):
    """Generates random numbers distributed according to a power law with
given index and limits."""
    def __init__(self, power, lolim, hilim):
        BasicSpectrum.__init__(self)
        if (lolim < 0.0):
            raise ValueError, "can not have negative energy or time limit"
        if (hilim <= lolim):
            raise ValueError, "wrong order of interval boundaries"
        if (power <= -1 and lolim == 0.0):
            raise ValueError, "integral of the spectral density diverges"
        self.power = power
        self.lolim  = lolim
        self.hilim  = hilim
        self._maxarea = self._integral(hilim)
    def _integral(self, uplim):
        if (self.power == -1):
            return math.log(uplim*1.0/self.lolim)
        else:
            pp1 = self.power + 1
            return (pow(uplim,pp1) - pow(self.lolim,pp1))/pp1        
    def integral(self, uplim = None):
        if (uplim is None):
            return self._maxarea
        elif (uplim <= self.lolim):
            return 0.0
        elif (uplim >= self.hilim):
            return self._maxarea
        else:
            return self._integral(uplim)
    def cdf(self, e):
        "Cumulative distrubution function."
        return self.integral(e)/self._maxarea
    def generate(self, thresh):
        if (self.hilim <= thresh):
            raise ValueError, "lower threshold is too high"
        if (thresh > self.lolim):
            lolim = thresh
        else:
            lolim = self.lolim
        r = rand.random()
        if (self.power == -1):
            # logarithmic cdf
            lgmin = math.log(lolim)
            return math.exp(r*(math.log(self.hilim) - lgmin) + lgmin)
        else:
            # power cdf
            pp1 = self.power + 1
            ppmin = pow(lolim, pp1)
            return pow(r*(pow(self.hilim, pp1) - ppmin) + ppmin, 1.0/pp1)


class PiecewisePowerSpectrum(BasicSpectrum):
    """Generates random numbers distributed according to power laws with
given indices and limits."""
    def __init__(self, ranges, powers):
        BasicSpectrum.__init__(self)
        nbins = len(powers)
        if (len(ranges) != nbins + 1):
            raise ValueError, "incompatible # of ranges and powers"
        # Start with scaling coefficient 1.0 for the first bin
        # and calculate scales for all other bins so that the
        # spectrum is continuous.
        self.ranges = list()
        self.ranges.append(ranges[0])
        self.bins = list()
        self.cdflist = list()
        self.emin = ranges[0]
        self.emax = ranges[0]
        self._0scale = 1.0
        self._lastscale = 1.0
        self._0power = powers[0]
        self._lastpower = powers[0]
        self._area = 0.0
        for i in xrange(nbins):
            self.appendPiece(powers[i], ranges[i+1])
    def appendPiece(self, power, hilim):
        if (hilim <= self.emax):
            raise ValueError, "bad piece boundary"
        s = PowerLawSpectrum(power, self.emax, hilim)
        self.bins.append(s)
        self.ranges.append(hilim)
        self._lastscale *= pow(self.emax, self._lastpower-power)
        self._area += self._lastscale*s.integral()
        self.cdflist.append(self._area)
        self._lastpower = power
        self.emax = hilim
    def prependPiece(self, power, lolim):
        if (lolim >= self.emin):
            raise ValueError, "bad piece boundary"
        s = PowerLawSpectrum(power, lolim, self.emin)
        self.bins.insert(0, s)
        self.ranges.insert(0, lolim)
        self._0scale /= pow(self.emin, power-self._0power)
        newarea = self._0scale*s.integral()
        self._area += newarea
        self.cdflist = [c + newarea for c in self.cdflist]
        self.cdflist.insert(0, newarea)
        self._0power = power
        self.emin = lolim
    def _integral(self, uplim = None):
        if (uplim is None):
            return self._area
        elif (uplim <= self.emin):
            return 0.0
        elif (uplim >= self.emax):
            return self._area
        binnum = bisect.bisect(self.ranges, uplim) - 1
        if (binnum > 0):
            mininteg = self.cdflist[binnum-1]
        else:
            mininteg = 0.0
        binarea = self.cdflist[binnum] - mininteg
        return mininteg + binarea*self.bins[binnum].cdf(uplim)
    def cdf(self, e):
        "Cumulative distrubution function."
        return self._integral(e)/self._area
    def generate(self, thresh):
        if (thresh >= self.emax):
            raise ValueError, "threshold is too high"
        elif (thresh < self.emin):
            r = self._area*rand.random()
            i = bisect.bisect(self.cdflist, r)
            return self.bins[i].generate(thresh)
        else:
            cdfmin = self.cdf(thresh)
            e = thresh - 1.0
            while (e <= thresh):
                r = self._area*rand.uniform(cdfmin, 1.0)
                i = bisect.bisect(self.cdflist, r)
                e = self.bins[i].generate(thresh)
            return e


class HistogramSpectrum(BasicSpectrum):
    "Generates bin numbers randomly distributed according to the bin height."
    def __init__(self, values):
        BasicSpectrum.__init__(self)
        self.nbins = len(values)
        if (self.nbins <= 0):
            raise ValueError, "number of bins must be positive"
        cdf = list()
        sum = 0.0
        for v in values:
            cdf.append(sum)
            if (v < 0.0):
                raise ValueError, "all bin weights must be non-negative"
            sum += v
        self.cdf = cdf
        self.sum = sum
        if (sum == 0.0):
            raise ValueError, "all bin weights are zero"
    def generate(self, threshold):
        minbin = threshold + 1
        if (minbin >= self.nbins):
            raise ValueError, "lower threshold is too high"
        if (self.nbins == 1):
            return 0
        if (minbin < 0):
            minbin = 0
        i = -1
        while (i < minbin or i >= self.nbins):
            r = rand.uniform(self.cdf[minbin], self.sum)
            i = bisect.bisect(self.cdf, r) - 1
        return i


class UniformlySampledSpectrum(BasicSpectrum):
    "Linearly interpolated spectrum from equidistant samples."
    def __init__(self, lolim, hilim, values):
        BasicSpectrum.__init__(self)
        nbins = len(values)
        if (nbins < 2):
            raise ValueError, "not enough data points"
        step = (hilim - lolim)*1.0/(nbins - 1)
        cdf = list()
        elist = list()
        sum = 0.0
        for i in xrange(nbins-1):
            e = lolim + i*step
            elist.append(e)
            cdf.append(sum)
            area = (values[i] + values[i+1])/2.0
            if (area < 0.0):
                raise ValueError, "sample sizes must be non-negative"
            sum += area
        elist.append(hilim)
        cdf.append(sum)
        self.values = values
        self.elist = elist
        self.cdf = cdf
        self.sum = sum
        if (sum == 0.0):
            raise ValueError, "sample weights sum is zero"
    def generate(self, threshold):
        if (self.elist[-1] <= threshold):
            raise ValueError, "lower threshold is too high"
        e = threshold - 1.0
        while (e <= threshold):
            r = self.sum*rand.random()
            i = bisect.bisect(self.cdf, r)
            e = linear_random(self.elist[i-1], self.values[i-1],\
                              self.elist[i], self.values[i])
        return e


class LogSampledSpectrum(UniformlySampledSpectrum):
    """Energy spectrum from samples equidistant in log(energy).

    It is assumed that the spectrum values in each bin are given in counts,
    not counts/MeV (that is, counts are not divided by the bin width).
    """
    def __init__(self, emin, emax, values):
        if (emin <= 0.0 or emax <= 0.0):
            raise ValueError, "energy must be positive"
        UniformlySampledSpectrum.__init__(self, math.log(emin), \
                                          math.log(emax), values)
    def generate(self, mass):
        if (mass > 0.0):
            loge = UniformlySampledSpectrum.generate(self, math.log(mass))
        else:
            loge = UniformlySampledSpectrum.generate(self, -1.0e100)
        return math.exp(loge)


class LogPlottedSpectrum(PiecewisePowerSpectrum):
    """Energy spectrum from points equidistant in log(energy).

    It is assumed that the spectrum values for each point are given in
    counts/MeV, not just counts (that is, already divided by the bin width).
    Spectrum is linearly interpolated between points in the log-log plot.
    """
    def __init__(self, emin, emax, values):
        npoints = len(values)
        points = logrange(emin, emax, npoints)
        step = math.log(emax*1.0/emin)/(npoints - 1)
        powers = [math.log(values[i+1]*1.0/values[i])/step \
                  for i in xrange(npoints-1)]
        PiecewisePowerSpectrum.__init__(self, points, powers)


class AMSHeliumSpectrum(LogPlottedSpectrum):
    """AMS helium spectrum for zenith pointing (predominantly galactic).

    See J. Alcaraz et al., Physics Letters B 494, 193 (2000).
    This implementation covers rigidity range from 0.76 GV to 500 GV.
    """
    def __init__(self, minTotalEnergy=0.0):
        # Spectrum in the paper is flux as a function of rigidity
        # in GV. The log-log spectrum plot is fitted to a cubic
        # polynomial below 22.1 GV and to a straight line above
        # (Table 1 in the paper). The fit CL is 96.4%.
        # Here, we will ignore He(3) contributions and will
        # assume 100% of He(4).
        self.pmass = particle_info.info["He"].mass/1000.0
        x0 = 1.34390145192
        a = 0.662291799321
        b = 5.48756096914e-05
        c = -2.70168646949
        d = -0.30934009085
        rigimin = max(.76, self._rigidity(minTotalEnergy))
        # Assume that linear portion of the spectrum extends
        # at least until 1 TeV
        rigimax = self._rigidity(1.0e6)
        #
        rigibreak = pow(10.0, x0)
        logmin = math.log10(rigimin)
        if (rigimin < rigibreak):
            npoints = 100
            logstep = (x0 - logmin)/(npoints - 1)
            values = [pow(10.0, shifted_polynomial( \
                logmin + i*logstep, x0, a, b, c, d)) for i in xrange(npoints)]
            LogPlottedSpectrum.__init__(self, rigimin, rigibreak, values)
            self.appendPiece(c, rigimax)
        else:
            vmin = pow(10.0, shifted_polynomial(logmin, x0, c, d))
            vmax = pow(10.0, shifted_polynomial(math.log10(rigimax),x0,c,d))
            LogPlottedSpectrum.__init__(self, rigimin, rigimax, (vmin, vmax))
    def _rigidity(self, emev):
        e = emev/1000.0
        if (e <= self.pmass):
            return 0.0
        else:
            return math.sqrt(e*e - self.pmass*self.pmass)/2.0
    def cdf(self, e):
        "Cumulative distrubution function."
        return LogPlottedSpectrum.cdf(self, self._rigidity(e))
    def generate(self, thresh):
        rigithresh = self._rigidity(thresh)
        e = thresh - 1.0
        while (e <= thresh):
            r = LogPlottedSpectrum.generate(self, rigithresh)
            e = 1000.0*math.sqrt(4.0*r*r + self.pmass*self.pmass)
        return e


class AMSGalacticProtonSpectrum(LogPlottedSpectrum):
    """Galactic proton spectrum observed by the AMS experiment.
    
    See J. Alcaraz et al., Physics Letters B 490, 27 (2000).
    This implementation covers Ek range from 220 MeV to 1 TeV.
    """
    def __init__(self, minTotalEnergy=0.0):
        # The proton spectrum in the AMS paper is fitted to a cubic
        # polynomial below 24.3 GeV and to a straight line above
        # (Fig 2 in the paper). The fit is very good (CL = 99.989%).
        self.pmass = particle_info.info["proton"].mass
        x0 = 4.38576633497
        a  = 0.113508547339
        b  = -0.428649737546
        c  = -2.72277196397
        d  = -2.65757902666
        emin = max(220.0, minTotalEnergy - self.pmass)
        # Assume that linear portion of the spectrum extends
        # at least until 1 TeV
        emax = 1.0e6 - self.pmass
        #
        ebreak = pow(10.0, x0)
        logmin = math.log10(emin)
        if (emin < ebreak):
            npoints = 100
            logstep = (x0 - logmin)/(npoints - 1)
            values = [pow(10.0, shifted_polynomial( \
                            logmin + i*logstep, x0, a, b, c, d)) \
                      for i in xrange(npoints)]
            LogPlottedSpectrum.__init__(self, emin, ebreak, values)
            self.appendPiece(c, emax)
        else:
            vmin = pow(10.0, shifted_polynomial(logmin, x0, c, d))
            vmax = pow(10.0, shifted_polynomial(math.log10(emax),x0,c,d))
            LogPlottedSpectrum.__init__(self, emin, emax, (vmin, vmax))
    def cdf(self, e):
        "Cumulative distrubution function."
        return LogPlottedSpectrum.cdf(self, e - self.pmass)
    def generate(self, thresh):
        ethresh = thresh - self.pmass
        return LogPlottedSpectrum.generate(self, ethresh) + self.pmass


def band_function(E, A, alpha, beta, E0):
    """GRB energy spectrum function.

    See D. Band et al., ApJ, 413, 281-292, 1993.
    """
    Ebreak = (alpha - beta)*E0
    if (Ebreak >= E):
        return A*pow(E/0.1, alpha)*math.exp(-E*1.0/E0)
    else:
        return A*pow(Ebreak/0.1,alpha-beta)* \
               math.exp(beta-alpha)*pow(E/0.1,beta)


def sbpl_function(E, A, Epiv, lambda1, lambda2, Eb, delta):
    """Smoothly-broken power law spectrum.

    See R.D. Preece et al., astro-ph/9908119.
    """
    m = (lambda2 - lambda1)/2.0
    b = (lambda2 + lambda1)/2.0
    alpha = math.log10(E*1.0/Eb)/delta
    alpiv = math.log10(Epiv*1.0/Eb)/delta
    beta = m*delta*math.log(math.cosh(alpha))
    bpiv = m*delta*math.log(math.cosh(alpiv))
    return A*pow(E*1.0/Epiv, b)*pow(10.0,beta-bpiv)


class GRBModelSpectrum(BasicSpectrum):
    """GRB energy spectrum model by D. Band et al., ApJ, 413, 281-292, 1993.

    The parameterization using Epeak is from R.D. Preece et al.,
    ApJ Supplement, 126, 19-36 (2000).
    """
    def __init__(self, emin, emax, alpha, beta, Epeak, nsamples=200):
        BasicSpectrum.__init__(self)
        if (alpha == -2.0):
            raise ValueError, "bad value of parameter alpha"
        if (beta >= 0.0):
            # In principle, beta should be less than -1, otherwise
            # the spectrum becomes divergent. But we have introduced
            # an energy cutoff...
            raise ValueError, "upper power law index must be negative"
        E0 = Epeak/(2.0+alpha)
        if (emin >= (alpha - beta)*E0):
            # The spectrum becomes a simple power law
            self.helperSpectrum = PowerLawSpectrum(beta, emin, emax)
        else:
            # Exact modeling of this spectrum is somewhat complicated,
            # and requires at least the ability to calculate exponential
            # integrals (see, e.g., "Numerical Recipes"). I think that
            # for the purpose of flight software testing an approximation
            # by a sampled generator should suffice. We will sample
            # the spectral density in the log(E) space.
            values = sample_regularly_in_log_space(emin, emax, nsamples,
                                                   band_function, 1.0,
                                                   alpha, beta, E0)
            self.helperSpectrum = LogSampledSpectrum(emin, emax, values)
    def generate(self, thresh):
        return self.helperSpectrum.generate(thresh)


class SBPLSpectrum(LogSampledSpectrum):
    """Smoothly-broken power law empirical spectrum model.

    See R.D. Preece et al., ApJ Supplement, 126, 19-36 (2000).
    """
    #
    # dflux/dE in the log-log plot should be fitted
    # to the function which looks like this:
    #
    # log10(flux) = const + (lambdalow + lambdahi)/2*log10(E) +
    #               width*(lambdahi-lambdalow)/2*log(cosh(y))
    #
    # where y = (log10(E) - log10(Ebreak))/breakwidth
    #
    def __init__(self, emin, emax, lambdalow, lambdahi,
                 Ebreak, breakwidth, nsamples=200):
        # Epiv is fixed
        Epiv = 0.1
        if (breakwidth <= 0.0):
            raise ValueError, "break scale must be positive"
        if (lambdahi >= 0.0):
            raise ValueError, "upper power law index must be negative"
        # Again, we are not going to attempt an exact MC generation,
        # and will use a sampled generator. Since the function is
        # pretty smooth in log(E), this approximation should be fine.
        values = sample_regularly_in_log_space(emin, emax, nsamples,
                                               sbpl_function, 1.0, Epiv,
                                               lambdalow, lambdahi,
                                               Ebreak, breakwidth)
        LogSampledSpectrum.__init__(self, emin, emax, values)


class GRBPulse(BasicSpectrum):
    "Simple GRB pulse model, as in ApJ, 459, 393-412 (1996)."
    def __init__(self, tpeak, s_l, s_r, nu, amplitude=1):
        BasicSpectrum.__init__(self)
        if (nu <= 0.0):
            raise ValueError, "bad peakedness parameter"
        if (s_l < 0.0):
            raise ValueError, "bad rise time"
        if (s_r < 0.0):
            raise ValueError, "bad decay time"
        self.tpeak  = tpeak
        self.s_l = s_l
        self.s_r = s_r
        self.alpha = 1.0/nu
        if (s_l == 0.0 and s_r == 0.0):
            self.degenerate = 1
        else:
            self.degenerate = 0
        self.amplitude = amplitude
    def _getWidth(self):
        return self.s_l + self.s_r
    def _setWidth(self, w):
        if (w < 0.0):
            raise ValueError, "width can't be negative"
        elif (w == 0.0):
            self.s_l = 0.0
            self.s_r = 0.0
            self.degenerate = 1
        else:
            if (self.degenerate):
                raise ValueError, "can't determine width ratio"
            scale = w*1.0/(self.s_l + self.s_r)
            self.s_r *= scale
            self.s_l *= scale
    width = property(_getWidth, _setWidth, doc="pulse width")
    def generate(self, thresh):
        if (self.degenerate):
            if (thresh >= self.tpeak):
                raise ValueError, "threshold is too high"
            else:
                return self.tpeak
        t = thresh - 1.0
        while (t <= thresh):
            # Should we generate left or right sample?
            if (rand.random() < self.s_l*1.0/self.width):
                # Left sample
                s = self.s_l
                sign = -1.0
            else:
                # Right sample
                s = self.s_r
                sign = 1.0
            # Generate exp(-x^n) density
            r = rand.gammavariate(self.alpha, 1.0)
            expxn = pow(r, self.alpha)
            # Scale and shift to get t
            t = sign*expxn*s + self.tpeak
        return t


class RandomGRBPulse(GRBPulse):
    def __init__(self, tpeak):
        # Generate pulse parameters using an approximate
        # BATSE GRB model. Based on ApJ, 459, 393-412 (1996).
        #
        # Peakedness is generated like in Gleam.
        peakedness = pow(10.0, rand.normalvariate(0.16, 0.3))
        # The width is not like in Gleam. I use shorter pulses
        # which look more like BATSE channel 4. Also, unlike Gleam,
        # there is no correlation between width and peakedness here.
        # Note that the width variation is not completely accounted
        # for in the formula below because subsequently the width
        # may be fudged in order to introduce correlation
        # with amplitude.
        fwhm  = pow(10.0, rand.normalvariate(-0.6, 0.5))
        width = 0.8493218*fwhm
        # The ratio between decay and rise time in Gleam is fixed to 3.
        # This assumption doesn't look very good to me.
        tratio = pow(10.0, rand.normalvariate(0.4, 0.1))
        s_l = width/(1.0 + tratio)
        s_r = width - s_l
        # Unfortunately, the discussion of amplitude distributions
        # of individual pulses in the ApJ paper is incomplete.
        # Table 2 does give some idea about the dependence of
        # the average pulse width on the normalized amplitude.
        # However, there is no info about the scatter, and
        # the average over the log of the pulse width would be 
        # much more useful than the average over the width.
        # Besides, this kind of correlation can be realized
        # only using higher level code which is responsible
        # for the whole light curve.
        #
        # From figure 2 in AIP 428, 261-265 (1997), (proceedings
        # of the 4th Huntsville GRB Symposium) it looks like
        # there is an inverse correlation between pulse width
        # and amplitude. In that plot the typical pulse width
        # drops by about two decades when the amplitude rises
        # by three decades. However, it is not at all clear
        # how much scatter in the amplitudes should be attributed
        # to within-GRB variation and how much to between-GRB
        # variation. I could not find more information in the
        # literature about the correlations between pulse
        # widths and amplitudes in the context of GRB pulse
        # decomposition analysis. So, short of redoing the whole
        # BATSE decomposition analysis, at this point I don't
        # have any other choice but to assume something for the
        # amplitude distribution. It appears that this path was
        # taken by the Gleam developers as well.
        amplitude = rand.uniform(0.01, 1.0)
        GRBPulse.__init__(self, tpeak, s_l, s_r, peakedness, amplitude)


# The following generator makes random number of pulses for a GRB
# pulse decomposition (from 1 to 40). It generates this number
# according to N_tot pulses from BATSE channel 4, table 1 in
# ApJ, 459, 393-412 (1996). The distribution of the number of
# pulses in that table was fitted to a decaying exponential, and
# the width turned out to be 6.467. Note that this generator is
# only intended for "long" GRBs, with T_90 above 1.5 sec.
__n_pulses_in_long_grb = HistogramSpectrum((0.0, 0.132475411329, 0.113496107063,
     0.0972359035483, 0.0833052444133, 0.0713703837113, 0.0611453901489,
     0.0523852968423, 0.0448802324847, 0.0384503933221, 0.0329417354763,
     0.0282222844146, 0.0241789731495, 0.0207149334184, 0.0177471749472,
     0.0152045971979, 0.0130262859659, 0.011160054019, 0.00956119081323,
     0.00819139133296, 0.00701783839277, 0.0060124164144, 0.00515103784342,
     0.00441306606787, 0.00378082101344, 0.00323915557026, 0.00277509270368,
     0.00237751455495, 0.00203689608333, 0.00174507686846, 0.00149506560583,
     0.00128087261148, 0.00109736632322, 0.000940150360422, 0.000805458197047,
     0.000690062924508, 0.0005911999425, 0.000506500725656, 0.000433936079232,
     0.000371767524351, 0.000318505648128))


def random_number_of_grb_pulses():
    # Are we going to generate a short GRB or a long GRB?
    # It is a common belief that long and short GRBs are generated
    # by different production mechanisms. The long/short fraction
    # is chosen to qualitatively reproduce the bimodal distribution
    # of BATSE T_90 GRB times, as in
    # http://f64.nsstc.nasa.gov/batse/grb/duration/
    if (rand.random() < 0.3):
        # Short GRB
        return 1
    else:
        # Long GRB
        return __n_pulses_in_long_grb.generate(0)


def random_grb_pulse_interval():
    """Generates random time interval between pulses inside a GRB.

    Distribution is loosely based on ApJ, 459, 393-412 (1996).
    Compared to the paper, the intervals are increased in order to
    better approximate the T_90 distribution.
    """
    return pow(10.0, 3.5*rand.betavariate(3.5,5.0) - 1.0)


class GRBLightCurve(BasicLightCurve):
    """GRB light curve decomposed into a collection of pulses.

    Based on ApJ, 459, 393-412 (1996). Approximate distributions
    are used for number of pulses, pulse separation, pulse peakedness,
    amplitude, width, and rise/decay time ratio. Argument t0 is the peak
    time of the first pulse. Note that the simplified model used here
    has no correlation between the light curve and the spectrum.
    """
    def __init__(self, tmin, tmax, t0, npulses = None):
        BasicLightCurve.__init__(self, tmin, tmax)
        if (npulses is None):
            npulses = random_number_of_grb_pulses()
        if (npulses <= 0):
            raise ValueError, "number of pulses must be positive"
        self.npulses = npulses
        self.pulses = list()
        tcurrent = t0.dtime()
        for i in xrange(npulses):
            pulse = RandomGRBPulse(tcurrent)
            self.pulses.append(pulse)
            tcurrent += random_grb_pulse_interval()
        if (npulses > 1):
            # Fudge pulse widths so that there is a trend similar
            # to that exhibited by Table 2 (channel 4) in the ApJ
            # paper. Keep average log of the pulse widths unmodified.
            amplitudes = [p.amplitude for p in self.pulses]
            norm = max(amplitudes)
            logfactors = [-4.5*a/norm for a in amplitudes]
            logaverage = reduce(operator.add, logfactors)*1.0/npulses
            for p, logf in zip(self.pulses, logfactors):
                p.width = p.width*math.exp(logf - logaverage)
            fluxes = [p.width * p.amplitude for p in self.pulses]
            self.pulse_chooser = HistogramSpectrum(fluxes)
        else:
            self.pulse_chooser = HistogramSpectrum((1,))
    def generate(self):
        t = self.tmin - 1.0
        while (t <= self.tmin or t >= self.tmax):
            i = self.pulse_chooser.generate(-1)
            t = self.pulses[i].generate(-1.0e300)
        return Time(t)


class PencilSource(BasicSource):
    "Point-like cosmic rays with fixed direction."
    def __init__(self, label, coords, particleName, eSpectrum, \
                 pointOnTheRay, rayDirection, launchBackoff):
        BasicSource.__init__(self, label, coords)
        self.pname     = particleName
        self.mass      = particle_info.info[particleName].mass
        self.eSpectrum = eSpectrum
        self.direction = rayDirection.direction()
        self.launch    = pointOnTheRay - self.direction*launchBackoff
    def area(self):
        return 0.0
    def solidAngle(self):
        return 0.0
    def generate(self):
        # Generate the particle energy
        energy = self.eSpectrum.generate(self.mass)
        # Calculate the particle momentum
        pmag = math.sqrt(energy*energy - self.mass*self.mass)
        return LabeledCosmicRay(self.pname, pmag*self.direction,
                                self.launch, self.label, self.coords)


class TubeSource(BasicSource):
    "Base class for cosmic rays generated uniformly over a circular patch."
    def __init__(self, label, coords, particleName, eSpectrum, max_impact_par,
                 backoff_distance, constant_shift = V3(0, 0, 0)):
        BasicSource.__init__(self, label, coords)
        self.pname      = particleName
        self.max_impact = max_impact_par
        self.backoff    = backoff_distance
        self.shift      = constant_shift
        self.mass       = particle_info.info[particleName].mass
        self.eSpectrum  = eSpectrum
    def area(self):
        return self.max_impact*self.max_impact*math.pi
    def solidAngle(self):
        raise NotImplementedError, "TubeSource.solidAngle"
    def direction(self):
        raise NotImplementedError, "TubeSource.direction"
    def generate(self):
        # Generate the particle energy
        energy = self.eSpectrum.generate(self.mass)
        # Calculate the particle momentum
        pmag = math.sqrt(energy*energy - self.mass*self.mass)
        p = self.direction()*pmag
        launch_point = self.shift + \
                       random_launch_point(self.max_impact, p, self.backoff)
        return LabeledCosmicRay(self.pname, p, launch_point,
                                self.label, self.coords)


class IsotropicTubeSource(TubeSource):
    "Uniform and isotropic cosmic rays."
    def solidAngle(self):
        return 4.0*math.pi
    def direction(self):
        return random_direction()


class ThetaRangeTubeSource(TubeSource):
    "Uniform and isotropic cosmic rays within a range of cos(theta)."
    def __init__(self, label, coords, particleName, eSpectrum,
                 max_impact_par, backoff_distance, thetaMin, thetaMax,
                 constant_shift = V3(0, 0, 0)):
        TubeSource.__init__(self, label, coords, particleName, eSpectrum,
                            max_impact_par, backoff_distance, constant_shift)
        self.__cosThetaMin = math.cos(thetaMin)
        self.__cosThetaMax = math.cos(thetaMax)
        if (self.__cosThetaMin > self.__cosThetaMax):
            self.__cosThetaMin,self.__cosThetaMax = \
            self.__cosThetaMax,self.__cosThetaMin
    def solidAngle(self):
        return 2.0*math.pi*(self.__cosThetaMax - self.__cosThetaMin)
    def direction(self):
        return random_direction(self.__cosThetaMin, self.__cosThetaMax)


class DirectedTubeSource(TubeSource):
    "Uniform circular patch cosmic rays with fixed direction."
    def __init__(self, label, coords, particleName, eSpectrum, max_impact_par,
                 backoff_distance, rayDirection, constant_shift = V3(0, 0, 0)):
        TubeSource.__init__(self, label, coords, particleName, eSpectrum, \
                            max_impact_par, backoff_distance, constant_shift)
        self.__dir = rayDirection.direction()
    def solidAngle(self):
        return 0.0
    def direction(self):
        return self.__dir