#########################################################################
#
# Script: rigidity.py
#
# Author: Igor Volobouev, March 2005
#
# Purpose: contains code relevant to calculation of cosmic ray
#          rigidity cutoffs
#
#########################################################################

import re
import math
import numarray
import EventGenUtils_config as config
from basic_kinematics import *
from filters import *

# Optimal covering of the sphere by 75 caps. This is a rotated
# point set (so that there is a point at zenith) from 
# http://www.research.att.com/~njas/coverings/dim3/cover.3.75.txt
# The covering radius is just below 15 degrees. For details see
# http://www.research.att.com/~njas/coverings/index.html
# This is the covering which was used to generate directional
# rigidity cutoff maps.
covering = (\
    V3(0.0, 0.0, 1.0),
    V3(-0.508592462897, 0.857118617202, -0.0817397255271),
    V3(-0.0173704421737, -0.628828507687, -0.777349969871),
    V3(-0.98574286312, 0.167576428111, 0.0151376533825),
    V3(0.590503517675, -0.793307651746, -0.148217965491),
    V3(0.249115551445, 0.704570676733, 0.664470919993),
    V3(0.351242790288, 0.485158303051, -0.800780820981),
    V3(0.828766507356, -0.327578973347, 0.453693831241),
    V3(0.705646765389, 0.706515476478, -0.0538379419473),
    V3(-0.157825337332, -0.987464766109, 0.00212098792495),
    V3(0.852897258653, -0.518859858272, 0.057885349235),
    V3(0.00184398746583, -0.197460456259, -0.980309118556),
    V3(0.849438608846, 0.409541055288, -0.332761436822),
    V3(-0.708634305645, 0.191140772912, -0.6791926279),
    V3(-0.964488514743, -0.263702559769, -0.0149286603867),
    V3(0.923503275327, -0.0337659224413, -0.382101508687),
    V3(0.374902228588, 0.0553809983974, -0.92540870107),
    V3(-0.895774623614, 0.317832443256, -0.310757721873),
    V3(0.18826733282, 0.846548841072, -0.497906086598),
    V3(0.961702807442, 0.272948157225, -0.0250402401265),
    V3(-0.376471000652, 0.075835555124, -0.923319313264),
    V3(0.543487309237, 0.68903232817, -0.479432993686),
    V3(-0.386000349351, 0.488806832359, -0.782350056522),
    V3(0.606093404331, -0.727371667386, 0.321840399439),
    V3(0.370830099824, -0.147707768099, 0.916879191774),
    V3(0.899448077858, 0.0955707617853, 0.426449744671),
    V3(-0.332297774893, -0.380409120007, -0.863056829077),
    V3(-0.0355951210871, 0.655537080638, -0.754323619717),
    V3(0.316332605987, -0.769854835658, -0.554307869692),
    V3(-0.14704543296, -0.874766874781, -0.461693139901),
    V3(-0.630054190183, 0.612088758548, -0.477890227032),
    V3(-0.581055320882, 0.719398946894, 0.380578335277),
    V3(0.302346768265, -0.935979854632, 0.180355602747),
    V3(0.687176002239, -0.159952451579, -0.70866378148),
    V3(0.628061833996, 0.443743916956, 0.639241479288),
    V3(0.342815315823, 0.304366145974, 0.888728816019),
    V3(-0.732196898752, 0.296857562555, 0.612995341754),
    V3(0.693759012111, 0.272437317818, -0.666690588635),
    V3(0.201481736647, -0.956597013183, -0.210540409821),
    V3(0.63294767167, -0.562510205318, -0.531939389254),
    V3(-0.0527094299475, -0.918531942266, 0.391817287816),
    V3(-0.54312143303, -0.539069893497, 0.643756754456),
    V3(-0.275796517398, 0.932927631562, 0.23147854168),
    V3(0.102687168486, -0.455359542416, 0.884365892919),
    V3(-0.662712120049, -0.237410154041, -0.71024577767),
    V3(0.146802586258, 0.950523772882, 0.273776474253),
    V3(-0.0476220877626, 0.987834149423, -0.148039960793),
    V3(-0.470944385892, -0.660286311139, -0.585007156127),
    V3(0.68368653748, 0.00826949897734, 0.729728945469),
    V3(-0.00751439563972, 0.254373338967, -0.967076904016),
    V3(-0.766162647791, -0.524725453165, -0.371022904859),
    V3(-0.780322990697, 0.624690221875, -0.0292943148662),
    V3(-0.200326013435, -0.69993007446, 0.685541668469),
    V3(-0.44066074127, 0.541586290518, 0.715892730112),
    V3(0.346487290169, -0.399658428891, -0.848657585819),
    V3(0.556807392985, -0.454926698882, 0.69498721266),
    V3(-0.796444095884, -0.461141325873, 0.391184713025),
    V3(0.839735132422, 0.470338736078, 0.271341815284),
    V3(-0.474885756049, -0.826600964357, 0.302017159157),
    V3(-0.778479495404, -0.625647390638, 0.0503489605003),
    V3(0.381041711814, 0.920119732839, -0.0904814406292),
    V3(0.990924675272, -0.119507043957, 0.0615333599049),
    V3(0.269246499752, -0.737497323556, 0.61935774809),
    V3(-0.0422476217786, 0.439606850557, 0.897196163276),
    V3(-0.916165515787, -0.0817759997025, 0.392369001776),
    V3(-0.670656570155, -0.120551897197, 0.731906418191),
    V3(-0.136512577199, 0.791484469215, 0.595748815574),
    V3(-0.868062117546, 0.390160896678, 0.306989633026),
    V3(-0.285346553811, 0.844495600826, -0.45321575923),
    V3(-0.492093302821, -0.854667549914, -0.165491874249),
    V3(-0.917225778945, -0.109307833299, -0.383077887665),
    V3(-0.291805749004, -0.294821974274, 0.909906263487),
    V3(0.861476161074, -0.405058745115, -0.306245386752),
    V3(0.516364602699, 0.780992027964, 0.351310474277),
    V3(-0.393246866011, 0.152133511585, 0.906759227704),
)

_index = numarray.zeros((181,360))
_cutoff = numarray.zeros((config.rigidity_latitude_nsteps,
                          config.rigidity_longitude_nsteps,
                          len(covering)), numarray.Float32)

def _fill_index():
    # Really dumb O(N*M) index mapper
    coverset = zip(covering, xrange(len(covering)))
    maxitheta, maxiphi = _index.getshape()
    phirange = xrange(maxiphi)
    cosphi = [math.cos(iphi/180.0*math.pi) for iphi in phirange]
    sinphi = [math.sin(iphi/180.0*math.pi) for iphi in phirange]
    phiset = zip(cosphi, sinphi, phirange)
    for itheta in xrange(maxitheta):
        theta = itheta/180.0*math.pi
        sintheta = math.sin(theta)
        costheta = math.cos(theta)
        for cp, sp, iphi in phiset:
            v = V3(sintheta*cp, sintheta*sp, costheta)
            maxcos = -2.0
            minindex = -1
            for dir, ic in coverset:
                cosa = sprod3(v, dir)
                if (cosa > maxcos):
                    maxcos = cosa
                    minindex = ic
            _index[itheta,iphi] = minindex

_fill_index()

def index(theta, phi):
    "Covering direction number from theta and phi in geographic coordinates."
    itheta = int(theta*180.0/math.pi + 0.5)
    phideg = phi*180.0/math.pi
    if (phideg >= -0.5):
        iphi = int(phideg + 0.5)
    else:
        iphi = int(phideg + 360.5)
    if (iphi >= 360):
        iphi -= 360
    return _index[itheta,iphi]

def _print(f):
    # Just print the map from theta, phi into the covering
    maxitheta, maxiphi = _index.getshape()
    for itheta in xrange(maxitheta):
        for iphi in xrange(maxiphi):
            print >> f, itheta, iphi, _index[itheta,iphi]

def _zenith_and_azimuth(f):
    # Print zenith direction and azimuth of the covering vectors
    for dir, ic in zip(covering, xrange(len(covering))):
        print >> f, ic, dir.theta()*180.0/math.pi, 90.0-dir.phi()*180.0/math.pi

_longitude_step = 360.0/config.rigidity_longitude_nsteps

def _read_rigidity_map(icover, f):
    r1 = re.compile(r'^\s*($|#)')
    linenum = 0
    errorline = 0
    count = 0
    ilat = 0
    ilon = 0
    for line in f:
        linenum += 1
        # Skip empty lines and lines which start with "#"
        if r1.match(line): continue
        # Normal data should be a 5-element list
        items = line.split()
        if (len(items) == 5):
            ilat = count / config.rigidity_longitude_nsteps
            ilon = count % config.rigidity_longitude_nsteps
            expected_lat = config.rigidity_min_latitude + \
                           ilat*config.rigidity_latitude_step
            expected_lon = ilon*_longitude_step
            lat = float(items[0])
            lon = float(items[1])
            Ru  = float(items[2])
            Rc  = float(items[3])
            Rl  = float(items[4])
            if (lat != expected_lat):
                raise ValueError, "unexpected latitude in line " + \
                      str(linenum) + ", file " + f.name()
            if (lon != expected_lon):
                raise ValueError, "unexpected longitude in line " + \
                      str(linenum) + ", file " + f.name()
            # Translate from GV into MV
            _cutoff[ilat,ilon,icover] = 1000.0*Rc
            count += 1
        else:
            errorline = linenum
        if (errorline): break
    if errorline:
        raise SyntaxError, "failed to parse line " + \
              str(errorline) + " in file " + f.name()
    if ilat != config.rigidity_latitude_nsteps - 1 or \
           ilon != config.rigidity_longitude_nsteps - 1:
        raise ValueError, "unexpected data size in file " + f.name()

def _load_rigidity_maps():
    "Loads the rigidity cutoff maps."
    for imap in xrange(len(covering)):
        filename = config.rigidity_maps % imap
        f = open(filename, "r")
        try:
            _read_rigidity_map(imap, f)
        finally:
            f.close()

_load_rigidity_maps()

_min_latitude = config.rigidity_min_latitude - \
                config.rigidity_latitude_step/2.0
_max_latitude = _min_latitude + \
                config.rigidity_latitude_step*config.rigidity_latitude_nsteps

def cutoff(latidude_deg, longitude_deg, theta_rad, phi_rad):
    """Returns rigidity cutoff for the given spacecraft location and
    cosmic ray direction in the geographic coordinate system."""
    if (latidude_deg < _min_latitude or latidude_deg > _max_latitude):
        raise ValueError, "latitude is out of range"
    ilat = int((latidude_deg - _min_latitude)/(_max_latitude - _min_latitude)*\
               config.rigidity_latitude_nsteps)
    if (ilat >= config.rigidity_latitude_nsteps):
        ilat = config.rigidity_latitude_nsteps - 1
    if (longitude_deg < 0.0):
        ilon = int((longitude_deg + 360.0)/_longitude_step + 0.5)
    else:
        ilon = int(longitude_deg/_longitude_step + 0.5)
    if (ilon >= config.rigidity_longitude_nsteps):
        ilon -= config.rigidity_longitude_nsteps
    # Cutoffs are calculated for looking direction,
    # not particle direction. Reverse the direction.
    ind = index(math.pi - theta_rad, math.pi + phi_rad)
    c = _cutoff[ilat,ilon,ind]
    # When the cutoffs were calculated, 100 GV was
    # the topmost starting point, so it essentially
    # plays the role of infinity.
    if (c < 100000.0):
        return c
    else:
        return 1.1e300

def min_cutoff():
    "Returns the smallest cutoff for the loaded set of cutoff rigidity maps."
    return _cutoff.min()

class RigidityFilter(EarthOccultationFilter):
    """Class for cosmic ray filtering based on particle rigidity."""
    def __call__(self, t, cosmicRay):
        rot = coords.rotation(cosmicRay.coords,coords.GEOGRAPHIC,self.orbit,t)
        p = rot(cosmicRay.p)
        lat, lon = self.orbit.latitude_and_longitude(t)
        cut = cutoff(lat, lon, p.theta(), p.phi())
        if (cut >= 1.e300):
            # Check for Earth occultation. If not occulted,
            # we may have a problem just due to insufficient
            # sampling of the solid angle by the rigidity maps
            # near the occultation cutoff.
            if (EarthOccultationFilter.__call__(self, t, cosmicRay) == 0):
                # Cosmic ray passes the Earth occultation filter.
                # Try again by changing the zenith angle a little.
                zangle = p.theta() + \
                         config.rigidity_map_covering_radius_deg/180.0*math.pi
                cut = cutoff(lat, lon, zangle, p.phi())
        if (cosmicRay.rigidity() > cut):
            return 0
        else:
            return 1