#########################################################################
#
# Script: source_utils.py
#
# Author: Igor Volobouev, Feb 2005
#
# Purpose: contains utility code for generation of cosmic ray samples
#
#########################################################################

"""Various utilities for generating collections of cosmic rays."""

import re
import math
import operator
import random as rand
import coords
import particle_info
import EventGenUtils_config as config

from basic_kinematics import *
from timerep import *

class CosmicRay:
    "Represents a cosmic ray in a manner suitable for reading by Gleam."
    def __init__(self, sourceId, particleName, p3, launchPoint):
        if (not particle_info.info.has_key(particleName)):
            raise ValueError, str(particleName) + \
                  " is not a valid particle name"
        self.sourceId    = sourceId
        self.name        = particleName
        self.p           = p3
        self.launchPoint = launchPoint
    def __repr__(self):
        return str(self.sourceId) + \
               ' ' + str(self.name) + \
               ' ' + str(self.p.x) + \
               ' ' + str(self.p.y) + \
               ' ' + str(self.p.z) + \
               ' ' + str(self.launchPoint.x) + \
               ' ' + str(self.launchPoint.y) + \
               ' ' + str(self.launchPoint.z)
    def rigidity(self):
        return abs(self.p)/abs(particle_info.info[self.name].charge)

class CosmicRayLabel:
    "Information pertaining to the cosmic ray origin."
    # Label is placed on the cosmic ray by its source.
    # Each source should have a unique label.
    def __init__(self, sourceTag, filterSet):
        self.tag = sourceTag
        self.filters = filterSet

class LabeledCosmicRay(CosmicRay):
    "Cosmic ray which knows how it was produced."
    def __init__(self, particleName, p3, launchPoint,
                 label, coordsys, filterState = None):
        sourceId = config.source_id.get(label.tag,-1)
        CosmicRay.__init__(self, sourceId, particleName, p3, launchPoint)
        self.label = label
        self.coords = coordsys
        if (filterState is None):
            self.filterState = dict()
        else:
            self.filterState = filterState
        # Filter states:
        # < 0: unknown
        #   0: passed
        # > 0: filtered away
    def applyFilters(self, t, filterList):
        "Applies the given set of filters to this cosmic ray."
        filtered = 0
        for filt in filterList:
            state = self.filterState.get(filt.tag,-1)
            if (state < 0):
                state = filt(t, self)
                assert state >= 0
                self.filterState[filt.tag] = state
            if (state > 0):
                filtered = 1
        return filtered
    def isFiltered(self):
        "Checks whether this cosmic ray has been filtered out."
        return reduce(operator.or_, self.filterState.itervalues(), 0)
    def rotate(self, rot, newcoords):
        # Note that original and rotated ray will share
        # particle name, label, and the filter state
        "Returns the rotated cosmic ray."
        r = LabeledCosmicRay(self.name, rot(self.p),
                             rot(self.launchPoint), self.label,
                             newcoords, self.filterState)
        # Preserve the source id even if it is unknown
        r.sourceId = self.sourceId
        return r

class Event:
    "Represents a cosmic ray event for subsequent simulation by GLEAM."
    def __init__(self, t, raylist):
        self.t = t
        self.rays = raylist
    def __repr__(self):
        return self.filterRep(1)
    def __lt__(self, other):
        return self.t < other.t
    def __cmp__(self, other):
        return cmp(self.t, other.t)
    def filterRep(self, writeFiltered):
        "String representation in which filtered cosmic rays can be removed."
        s = str(self.t)
        for r in self.rays:
            if (writeFiltered or not r.isFiltered()):
                s += ' '
                s += str(r)
        return s
    def applyFilters(self, filterList = None):
        "Applies the given set of filters to this event."
        if (filterList is None):
            filtered = reduce(operator.and_, [LabeledCosmicRay.applyFilters(\
                ray, self.t, ray.label.filters) for ray in self.rays], 1)
        else:
            filtered = reduce(operator.and_, [LabeledCosmicRay.applyFilters(\
                ray, self.t, filterList) for ray in self.rays], 1)
        if filtered:
            return 1
        else:
            return 0
    def isFiltered(self):
        "Checks whether this event has been filtered out."
        return reduce(operator.and_, [r.isFiltered() for r in self.rays], 1)
    def rotate(self, rot, newcoords):
        "Returns the rotated event."
        return Event(self.t, [r.rotate(rot, newcoords) for r in self.rays])
    def local(self, orbit):
        "Returns event representation in the spacecraft local coordinates."
        # Check whether all rays have the same coordinate system
        coord0 = self.rays[0].coords
        allsame = reduce(operator.and_, [r.coords == coord0 for r in self.rays], 1)
        if (allsame):
            if (coord0 == coords.LOCAL):
                return self
            else:
                rot = coords.rotation(coord0, coords.LOCAL, orbit, self.t)
                return self.rotate(rot, coords.LOCAL)
        else:
            return Event(self.t, [r.rotate(\
                coords.rotation(r.coords, coords.LOCAL, orbit, self.t),\
                coords.LOCAL) for r in self.rays])

class EventSequence(list):
    "Represents a sequence of cosmic ray events."
    def __init__(self, info=""):
        list.__init__(self)
        self.header = info
    def write(self, f, orbit=None, writeFiltered=0):
        """Sorts an event sequence and writes it to a file. Provide the orbit
        argument in order to write events in the spacecraft local coordinates."""
        self.sort()
        if (len(self.header) > 0):
            print >> f, '#', "\n# ".join(self.header.splitlines())
        n = 0
        for el in self:
            if (writeFiltered or not el.isFiltered()):
                if (orbit is None):
                    rep = el.filterRep(writeFiltered)
                else:
                    rep = el.local(orbit).filterRep(writeFiltered)
                print >> f, n, rep, ';'
                n += 1
        return n
    def read(self, f, filters=(), coords=coords.LOCAL):
        """Reads an event sequence from the given file object.

        This function may break time ordering. Use sort() as needed."""
        r1 = re.compile(r'^\s*($|#)')
        linenum = 0
        errorline = 0
        inHeader = 1
        event = ""
        self.header = ""
        n0 = len(self)
        for line in f:
            linenum += 1
            # Skip empty lines and lines which start with "#"
            if r1.match(line):
                if (inHeader):
                    self.header += line.lstrip("# \t")
                continue
            inHeader = 0
            event += line
            if (line[-2] == ';'):
                words = event[:-2].split()
                event = ""
                if (len(words) == 11):
                    try:
                        sec  = int(words[1])
                        nsec = int(words[2])
                        sid  = int(words[3])
                        particle = words[4]
                        px = float(words[5])
                        py = float(words[6])
                        pz = float(words[7])
                        x  = float(words[8])
                        y  = float(words[9])
                        z  = float(words[10])
                    except ValueError:
                        errorline = linenum
                    else:
                        label = CosmicRayLabel(\
                            config.source_tag.get(sid,"unknown"), filters)
                        r = LabeledCosmicRay(particle, V3(px, py, pz),
                                             V3(x, y, z), label, coords)
                        # Preserve the source id even if it is unknown
                        r.sourceId = sid
                        self.append(Event(Time(sec, nsec), [r,]))
                else:
                    errorline = linenum
            if (errorline): break
        if errorline:
            raise SyntaxError, "failed to parse line " + \
                  str(errorline) + " in file " + f.name()
        return len(self) - n0
    def grow(self, eventGenerator, N = None):
        """Appends events produced by the given generator.

        This function breaks time ordering. Call sort() as necessary. Returns
        the number of events appended which, depending on the generator,
        may turn out to be smaller than the number requested."""
        count = 0
        if (N is None):
            # Add all events in the given generator
            while 1:
                event = eventGenerator.generate()
                if (event is None):
                    break
                self.append(event)
                count += 1
        else:
            # Add up to the given number of events
            for i in xrange(N):
                event = eventGenerator.generate()
                if (event is None):
                    break
                self.append(event)
                count += 1
        return count
    def applyFilters(self, filterList = None):
        return reduce(operator.add, [Event.applyFilters(event, filterList) \
                                   for event in self], 0)

class BasicSpectrum:
    "Abstract class for cosmic ray energy spectra."
    def __init__(self):
        pass
    def generate(self, lolim):
        # Particle mass will normally play the role of
        # the lower limit for various energy spectra
        raise NotImplementedError, "BasicSpectrum.generate"

class BasicSource:
    "Abstract class for cosmic ray sources."
    def __init__(self, label, coords):
        self.label = label
        self.coords = coords
    def area(self):
        raise NotImplementedError, "BasicSource.area"
    def solidAngle(self):
        raise NotImplementedError, "BasicSource.solidAngle"
    def generate(self):
        # Should return LabeledCosmicRay objects
        raise NotImplementedError, "BasicSource.generate"

class BasicEventGenerator:
    "Abstract class for event generators."
    def __init__(self):
        pass
    def generate(self):
        # Should generate Event objects or return None
        raise NotImplementedError, "BasicGenerator.generate"

class SimpleEventGenerator(BasicEventGenerator):
    "Generates single-ray events by combining a source and a light curve."
    def __init__(self, cosmicRaySource, timeGenerator):
        BasicEventGenerator.__init__(self)
        self.source = cosmicRaySource
        self.timer = timeGenerator
    def generate(self):
        cosmic_ray = self.source.generate()
        t = self.timer.generate()
        if (t is None):
            return None
        else:
            return Event(t, [cosmic_ray,])

class SequenceEventGenerator(EventSequence,BasicEventGenerator):
    "Picks events from a sequence. Can be used as a circular buffer."
    def __init__(self, isCircular, info=""):
        EventSequence.__init__(self, info)
        BasicEventGenerator.__init__(self)
        self.generated = 0
        self.circular = isCircular
    def rewind():
        self.generated = 0
    def generate(self):
        if (self.generated >= len(self)):
            if (self.circular):
                self.rewind()
            else:
                return None
        i = self.generated
        self.generated += 1
        return self[i]

class BasicLightCurve:
    "Abstract class for cosmic event light curves."
    def __init__(self, tmin, tmax):
        self.tmin = tmin.dtime()
        self.tmax = tmax.dtime()
        if (self.tmax < self.tmin):
            self.tmin, self.tmax = self.tmax, self.tmin
    def generate(self):
        # Should generate Time objects or return None.
        # Note that the time generated is normally not
        # going to be sequential.
        raise NotImplementedError, "BasicLightCurve.generate"

class SpectrumLikeLightCurve(BasicLightCurve):
    "Light curve which looks like an existing energy spectrum."
    def __init__(self, tmin, tmax, spectrum):
        BasicLightCurve.__init__(self, tmin, tmax)
        self.spectrum = spectrum
    def generate(self):
        t = self.tmax + 1.0
        while (t >= self.tmax):
            t = self.spectrum.generate(self.tmin)
        return Time(t)

class UniformLightCurve(BasicLightCurve):
    "Light curve uniformly distributed withing the given interval."
    def __init__(self, tmin, tmax):
        BasicLightCurve.__init__(self, tmin, tmax)
    def generate(self):
        return Time(rand.uniform(self.tmin, self.tmax))

class PeriodicLightCurve(BasicLightCurve):
    "Repeats input light curve periodically."
    def __init__(self, tmin, tmax, onePeriodCurve, phase_at_tmin):
        BasicLightCurve.__init__(self, tmin, tmax)
        self.periodGen = onePeriodCurve
        self.period = onePeriodCurve.tmax - onePeriodCurve.tmin
        self.phase0 = phase_at_tmin
        while (self.phase0 < 0.0):
            self.phase0 += 2*math.pi
        while (self.phase0 >= 2*math.pi):
            self.phase0 -= 2*math.pi
        if (self.phase0 > 0.0):
            self.period_start = (1.0 - self.phase0/(2*math.pi))*self.period
            self.add_period = 1
        else:
            self.period_start = 0.0
            self.add_period = 0
        self.nperiods = int((self.tmax - self.period_start)/self.period) + 1
        self.nperiods += self.add_period
    def generate(self):
        t = self.tmax + 1.0
        while (t <= self.tmin or t >= self.tmax):
            trel = self.periodGen.generate().dtime() - self.periodGen.tmin
            shift = self.period*(int(self.nperiods*rand.random()) - \
                                 self.add_period)
            t = self.period_start + trel + shift
        return Time(t)

class ClockTickLightCurve(BasicLightCurve):
    "Generates ticks sequentially at regular time intervals."
    def __init__(self, tmin, tmax, period):
        BasicLightCurve.__init__(self, tmin, tmax)
        self.dt = period.dtime()
        self.counter = 0
    def generate(self):
        t = self.tmin + self.dt*self.counter
        if (t >= self.tmax):
            return None
        else:
            self.counter += 1
            return Time(t)

def shifted_polynomial(x, x0, *coeffs):
    "Calculates a*(x-x0)^N + b*(x-x0)^(N-1) + ..."
    sum = coeffs[0]
    diff = x - x0
    for c in coeffs[1:]:
        sum *= diff
        sum += c
    return sum

def logrange(emin, emax, nsamples):
    "A sequence of numbers equidistant in the log space."
    if (emin <= 0.0):
        raise ValueError, "low range limit must be positive"
    if (emax <= emin):
        raise ValueError, "upper range limit must be higher than the lower one"
    if (nsamples < 2):
        raise ValueError, "must have at least two samples"
    logmin = math.log(emin)
    step = (math.log(emax) - logmin)/(nsamples - 1)
    elist = [math.exp(logmin + step*i) for i in xrange(1, nsamples-1)]
    elist.insert(0,emin)
    elist.append(emax)
    return elist

def eval_regularly_in_log_space(emin, emax, nsamples, callable, *args):
    "Evaluates a function on a regular grid in the log space."
    values = list()
    callargs = list(args)
    for e in logrange(emin, emax, nsamples):
        callargs.insert(0,e)
        values.append(apply(callable, callargs))
        callargs.pop(0)
    return values

def sample_regularly_in_log_space(emin, emax, nsamples, callable, *args):
    "Evaluates a probability density on a regular grid in the log space."
    # This function is intended for sampling statistical densities.
    # It differs from "eval_regularly_in_log_space" by the Jacobian factor.
    values = list()
    callargs = list(args)
    for e in logrange(emin, emax, nsamples):
        callargs.insert(0,e)
        values.append(e*apply(callable, callargs))
        callargs.pop(0)
    return values

def random_direction(cosThetaMin = -1.0, cosThetaMax = 1.0):
    "Generates random directions in the given cos(theta) interval."
    cosTheta = rand.uniform(cosThetaMin, cosThetaMax)
    sinTheta = math.sqrt(1.0 - cosTheta*cosTheta)
    phi      = 2.0*math.pi*rand.random()
    return V3(sinTheta*math.cos(phi), sinTheta*math.sin(phi), cosTheta)

def linear_random(x0, y0, x1, y1):
    "Generates random numbers on the (x0, x1) interval with linear pdf"
    if (x1 < x0):
        raise ValueError, "empty interval"
    if (y0 < 0.0 or y1 < 0.0):
        raise ValueError, "negative probability density"
    if (y0 == 0.0 and y1 == 0.0):
        raise ValueError, "zero probability density"
    if (x1 == x0):
        return x1
    if (math.fabs(y0 - y1)*1.0/(y0 + y1) < 1.0e-8):
        return rand.uniform(x0, x1)
    r = rand.random()
    tmp = (math.sqrt(y0*y0 + r*(y1*y1 - y0*y0)) - y0)/(y1 - y0)
    return x0 + tmp*(x1 - x0)

def random_rotation():
    "Generates a random rotation."
    z = random_direction()
    x = random_direction()
    while (math.fabs(sprod3(z, x)) > 0.99995):
        # Too close for subsequent ortogonalization
        x = random_direction()
    # Assume that the Rotation constructor
    # will ortogonalize the x axis
    return Rotation(x.ra(), x.dec(), z.ra(), z.dec())

def random_point_inside_xy_circle(r):
    "Generates uniformly random points inside a circle with given radius."
    x = 1.0
    y = 1.0
    while (x*x + y*y >= 1.0):
        x = rand.random()*2.0 - 1.0
        y = rand.random()*2.0 - 1.0
    return V3(x*r, y*r, 0.0)

def random_point_on_disk(r, diskNormal):
    "Generates uniformly random points on disk orthogonal to the given direction."
    p = random_point_inside_xy_circle(r)
    normDir = diskNormal.direction()
    if (normDir.x == 0.0 and normDir.y == 0.0):
        # Normal is parallel to the z axis
        return p
    else:
        rot = Rotation(V3(normDir.y, -normDir.x, 0.0), math.acos(normDir.z))
        return rot(p)

def random_launch_point(max_impact_parameter, direction, how_far_away):
    "Generates random cosmic ray launch points."
    return random_point_on_disk(max_impact_parameter, direction) - \
           direction.direction()*how_far_away

def lognormal_parameters(mean, stdev, skewness):
    "Calculates natural parameters of the lognormal distribution from moments."
    # The lognormal distribution is defined by
    # exp(-(log(x-x0) - mu)^2/s^2/2) / (sqrt(2*Pi)*s*(x-x0))
    # The logarithm of (x-x0) has normal distribution.
    # We want x0, mu, and s given mean, standard deviation, and skewness.
    if (stdev <= 0.0):
        raise ValueError, "standard deviation must be positive"
    if (skewness == 0.0):
        raise ValueError, "skewness can not be zero"
    b1 = skewness*skewness
    tmp = pow((2.0+b1+math.sqrt(b1*(4.0+b1)))/2.0, 1.0/3.0)
    w = tmp+1.0/tmp-1.0
    logw = math.log(w)
    if (logw <= 0.0):
        raise ValueError, "the absolute value of skewness is too small,\
                           the distribution is Gaussian"
    s = math.sqrt(logw)
    emgamovd = stdev/math.sqrt(w*(w-1.0))
    xi = mean - emgamovd*math.sqrt(w)
    if (skewness > 0.0):
        x0 = xi
    else:
        x0 = xi - 2.0*mean
    mu = math.log(emgamovd)
    return (x0, mu, s)