#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

#
# PMCmaster.py
# $Id: PMCmaster.py 252 2009-11-06 14:36:39Z fab $
#

############################################################################
#    Copyright (C) 2008-2009 by Fabian Euchner and Danijel Schorlemmer     #
#    fabian@fabian-euchner.de                                              #
#                                                                          #
#    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 2 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, write to the                         #
#    Free Software Foundation, Inc.,                                       #
#    59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             #
############################################################################

import cProfile

import sys
import os
import getopt
import glob
import re
import random
import datetime

import pickle
import cPickle

from mx.DateTime     import Date, DateTime, utc
from mx.DateTime.ISO import ParseDateTimeUTC

sys.path.append('../../..')
sys.path.append('../..')
sys.path.append('..')

from QPCore    import *
from QPUtils   import *
from QPCatalog import *
from QPPolygon import *

from pmc.PMC         import *
from pmc.PMCFactory  import *
from pmc.PMCMetadata import *

from pmc.PMCGridN    import *

# metadata for this computation run
metadata = PMCMetadata()

## set a few paths that do not go into metadata

# this gets the absolute base path  == path where this script resides
# subsequent directory hierarchy will be below this directory
scriptname = os.path.basename( os.path.abspath( sys.argv[0] ) )

# get working directory - subsequent directory hierarchy will be below
basepath = os.getcwd()

### location of station list and station alias files
stationdir = os.path.join( basepath, 'data/station/' )

### directory where data from this run will be written to
rundir = os.path.join( basepath, 'data/runs/' )

### location of catalog files
catalogdir = os.path.join( basepath, 'data/catalog/' )

### location of combinations pickle
combidir = os.path.join( basepath, 'data/combinations' )

### location of catalog files
colormapdir = os.path.join( basepath, 'data/colormaps/' )

# ----------------------------------------------------------------------------------

metadata.smoothDist   = True

# metadata file prefix
metadata_file_prefix = 'pmcinfo'

# verbosity level
verbose_level = 2
verbose_grid  = True

# profiling
profileFlag   = False
profileFile   = 'profiler.result'

# ----------------------------------------------------------------------------------

def main():

    global scriptname
    global metadata

    global verbose_level
    global profileFlag
    global profileFile

    programStartTime = utc()
    
    print " +++ program start time: %s" % ( mxDateTime2ISO( programStartTime, showsecfrac=False ) )
    
    setUp()

    if (     ( metadata.combinePickInfoProbDistro is False )
         and ( metadata.skipProbDistro is True )
         and (     metadata.skipPickInfo is False
               and metadata.skipGrid is False ) ):
        print "%s error: skipping only distro creation (-D) does not make sense. Combine with switches -P and -G" % ( scriptname )
        sys.exit()

    if (     ( metadata.discardPickInfo is True )
         and ( metadata.combinePickInfoProbDistro is False ) ):
        print "%s error: you cannot discard pickInfo if not combined with distro computation. Use switch -C" % ( scriptname )
        sys.exit()

    if profileFlag is False:
        runParts()
    else:
        cProfile.run( 'runParts()', profileFile )

    programEndTime = utc()
    print " +++ program end time: %s" % mxDateTime2ISO( programEndTime, showsecfrac=False )
    print " +++ execution time: %s" % str( programEndTime - programStartTime )

    
def runParts():

    global metadata
    global verbose_level
    
    if metadata.skipPickInfo is False:
        compPickInfos( verbose_level=verbose_level )

    if metadata.skipProbDistro is False:
        compProbDistros( verbose_level=verbose_level )

    if metadata.skipGrid is False:
        compGrids( verbose_level=verbose_level )

# ---------------------------------------------------------------------------

def PrintHelp():
    global scriptname
    
    print 'computes PMC grid'
    print 'Usage: %s [OPTION]' % scriptname
    print '  Options'
    print '   -a FILE, --aliasfile=<file>      Alias file name'
    print '   -A DIR, --stationdir=<dir>       Station directory'
    print '   -c DIR, --catalogdir=<dir>       Catalog directory'
    print '   -C, --combipickprob              Compute prob distro directly after pickInfo for each station'
    print '   -d VALUE, --date=<value>         Date for which PMC grid is computed (YYYY-MM-DD)'
    print '   -D, --skipprobdistro             Skip computation of probability distributions'
    print '   -f, --plotpickinfo               Plot pickInfos'
    print '   -F, --plotdistro                 Plot probability distros'
    print '   -g VALUE, --griddir=<value>      Directory to place PMC grid file'
    print '   -G, --skipgrid                   Skip computation of PMC Grid'
    print '   -i FILE, --pickcountfile=<file>  Pick count file name'
    print '   -l DIR, --colormapdir=<dir>      Color map directory'
    print '   -n VALUE, --network=<value>      Network identifier [SC,JP,NZ]'
    print '   -N, --maxstations                Maximum number of stations used for probability computation (default: 100)'
    print '   -o DIR, --combidir=<dir>         Combinations directory'
    print '   -O, --overwrite                  Write to already existing dirs, possibly overwrite existing files'
    print '   -P, --skippickinfo               Skip computation of pickInfo'
    print '   -p FILE, --polygonfile=<file>    Polygon boundary file name'
    print '   -r FILE, --regionfile=<file>     Region settings file name'
    print '   -R, --refine                     Refine pick info'
    print '   -s FILE, --stationfile=<file>    Station file name'
    print '   -S, --profile                    Profile run'
    print '   -t VALUE, --timebefore=<value>   Days before date'
    print '   -T VALUE, --timeafter=<value>    Days after date'
    print '   -w VALUE, --workingdir=<value>   Working directory (absolute path)'
    print '   -X, --discardpickinfo            Do not store pickInfo files, only together with -C option'
    print '   -h, --help                       print this information'

# ---------------------------------------------------------------------------

def setUp():

    global rundir
    global scriptname

    global metadata
    
    global stationdir
    global metadata_file_prefix

    global profileFlag
    global profileFile

    # command line variables
    in_skipdistro       = False
    in_skipgrid         = False
    in_skippick         = False
    in_combipickprob    = False
    in_discardpickinfo  = False
    in_refinepickinfo   = False
    in_overwrite        = False

    in_plotpickinfo     = False
    in_plotdistro       = False
    
    in_aliasfile_name   = None
    in_stationdir       = None
    in_polygonfile_name = None
    in_regionfile_name  = None
    in_pickcountfile_name  = None
    in_catalogdir       = None
    in_colormapdir      = None
    in_combidir         = None
    
    in_date             = None
    in_griddir          = None
    
    in_network          = None
    in_stationfile_name = None
    in_earlier          = None
    in_later            = None
    in_workingdir       = None

    in_max_stations_cnt = 100
               
    # Read commandline arguments
    cmdParams = sys.argv[1:]
    if len( cmdParams ) == 0:
        PrintHelp()
        sys.exit()
            
    opts, args = getopt.gnu_getopt( cmdParams,
                                    'hSDGPCXfFROa:A:c:d:g:i:l:n:N:o:p:r:s:t:T:w:',
                                    [ 'help',
                                      'profile',
                                      'skipprobdistro',
                                      'skipgrid',
                                      'skippickinfo',
                                      'combipickprob',
                                      'discardpickinfo',
                                      'plotpickinfo',
                                      'plotdistro',
                                      'refine',
                                      'overwrite',
                                      'aliasfile=',
                                      'stationdir=',
                                      'catalogdir=',
                                      'date=',
                                      'griddir=',
                                      'pickcountfile=',
                                      'colormapdir=',
                                      'network=',
                                      'maxstations=',
                                      'combidir=',
                                      'polygonfile=',
                                      'regionfile=',
                                      'stationfile=',
                                      'timebefore=',
                                      'timeafter=',
                                      'workingdir=' ] )

    for option, parameter in opts:

        if option == '-S' or option == '--profile':
            profileFlag = True

        if option == '-D' or option == '--skipprobdistro':
            in_skipdistro = True

        if option == '-G' or option == '--skipgrid':
            in_skipgrid = True
            
        if option == '-P' or option == '--skippickinfo':
            in_skippick = True

        if option == '-C' or option == '--combipickprob':
            in_combipickprob = True

        if option == '-X' or option == '--discardpickinfo':
            in_discardpickinfo = True

        if option == '-f' or option == '--plotpickinfo':
            in_plotpickinfo = True

        if option == '-F' or option == '--plotdistro':
            in_plotdistro = True

        if option == '-R' or option == '--refine':
            in_refinepickinfo = True
            
        if option == '-O' or option == '--overwrite':
            in_overwrite = True
            
        if option == '-a' or option == '--aliasfile':
            in_aliasfile_name = parameter

        if option == '-A' or option == '--stationdir':
            in_stationdir = parameter
            
        if option == '-c' or option == '--catalogdir':
            in_catalogdir = parameter

        if option == '-d' or option == '--date':
            in_date = parameter

        if option == '-g' or option == '--griddir':
            in_griddir = parameter

        if option == '-i' or option == '--pickcountfile':
            in_pickcountfile_name = parameter
            
        if option == '-l' or option == '--colormapdir':
            in_colormapdir = parameter
             
        if option == '-n' or option == '--network':
            in_network = parameter
            
        if option == '-N' or option == '--maxstations':
            in_max_stations_cnt = parameter

        if option == '-o' or option == '--combidir':
            in_combidir = parameter
             
        if option == '-p' or option == '--polygonfile':
            in_polygonfile_name = parameter

        if option == '-r' or option == '--regionfile':
            in_regionfile_name = parameter
            
        if option == '-s' or option == '--stationfile':
            in_stationfile_name = parameter

        if option == '-t' or option == '--timebefore':
            in_earlier = parameter

        if option == '-T' or option == '--timeafter':
            in_later = parameter

        if option == '-w' or option == '--workingdir':
            in_workingdir = parameter

        if option == '-h' or option == '--help':
            PrintHelp()
            sys.exit()

    # check if valid network identifier has been specified
    if in_network is not None and in_network in ( 'SC', 'JP', 'NZ' ):
        metadata.network = in_network
    else:
        error_str = "%s - no valid network identifier has been specified" % ( scriptname )
        raise ValueError, error_str

    # network-specific settings, hard-coded for the moment
    if metadata.network == 'SC':

        ### REGION SETTINGS
        
        metadata.catalogFilesPattern = '*.dat.gz'

        # CHECK:
        metadata.minimumPickCount = 10

        metadata.catalogStartDate = DateTime( 1980, 1, 1 )
        metadata.catalogEndDate   = DateTime( 2008, 1, 1 )

        # distStyle for prob distros
        metadata.useDistStyle = 3

        # stuff for resulting grid XML
        metadata.targetMagArray  = frange( 0.0, 4.0, 0.1 )
        metadata.targetProbArray = [ 0.9, 0.95, 0.99, 0.999 ]

        # coordinate limits: lon (-122.0,-113.5), lat (31.5,38.0)
        metadata.areaPolygon = { 'vertices': ( ( -122.0, 31.5 ),
                                               ( -113.5, 31.5 ),
                                               ( -113.5, 38.0 ),
                                               ( -122.0, 38.0 ),
                                               ( -122.0, 31.5 ) ),
                                               
                                 'metadata': { 'lonDelta'              : 0.1,
                                               'latDelta'              : 0.1,
                                               'lonAlign'              : 0.0,
                                               'latAlign'              : 0.0,
                                               'includePointOnBoundary': True,
                                               'shift'                 : False  }
                               } 

        metadata.areaDepth = 7.5

    elif metadata.network == 'JP':

        ### REGION SETTINGS

        # deck file names must end with '.Z.bz2'
        # works only for deck files with A,B,C part in names (starting from October 1997)
        metadata.catalogFilesPattern = '*.Z.bz2'

        metadata.minimumPickCount    = 100
        
        # catalog start and end date
        # Note: this is hard-coded at the moment, should be computed on the fly later
        metadata.catalogStartDate = DateTime( 2000, 1, 1 )
        metadata.catalogEndDate   = DateTime( 2009, 2, 1 )

        # distStyle for prob distros
        metadata.useDistStyle = 5

        # stuff for resulting grid XML
        metadata.targetMagArray  = frange( -1.0, 5.0, 0.1 )
        metadata.targetProbArray = [ 0.9, 0.95, 0.99, 0.999 ]

        # coordinate limits: lon (120,150), lat (20,49)
        metadata.areaPolygon = { 'vertices': ( ( 120.0, 20.0 ),
                                               ( 120.0, 36.0 ),
                                               ( 134.0, 49.0 ),
                                               ( 150.0, 49.0 ),
                                               ( 150.0, 20.0 ),
                                               ( 120.0, 20.0 ) ),
                                               
                                 'metadata': { 'lonDelta'              : 0.1,
                                               'latDelta'              : 0.1,
                                               'lonAlign'              : 0.0,
                                               'latAlign'              : 0.0,
                                               'includePointOnBoundary': True,
                                               'shift'                 : False  }
                               } 

        metadata.areaDepth = 30.0

    
    elif metadata.network == 'NZ':
        raise ValueError, "network NZ: not yet implemented"

    # PMC grid file prefix
    metadata.gridFilePrefix = metadata.network.lower()

    ## times: computation datetime and surrounding interval
    
    try:
        # useDate as mxDateTime instance
        metadata.useDate = ParseDateTimeUTC( in_date )
    except:
        error_str = "%s - illegal date format %s" % ( scriptname, in_date )
        raise ValueError, error_str

    # get start and end dates
    if in_earlier is not None:
        metadata.earlierDays = int( in_earlier )
        metadata.startDate   = metadata.useDate - DateTimeDelta( float( metadata.earlierDays ) )
    else:
        # use data from beginning of catalog
        metadata.startDate   = metadata.catalogStartDate
        metadata.earlierDays = ( metadata.useDate - metadata.startDate ).days
        
    if in_earlier is not None:
        metadata.laterDays = int( in_later )
        metadata.endDate   = metadata.useDate + DateTimeDelta( float( metadata.laterDays ) )
    else:
        # use data until end of catalog
        metadata.endDate   = metadata.catalogEndDate
        metadata.laterDays = ( metadata.endDate - metadata.useDate ).days
    
    # max stations used for combined probability
    metadata.useMaxStationCnt = in_max_stations_cnt

    # refine pickInfo using Michael's method
    metadata.refinePickInfo = in_refinepickinfo

    ## evaluate directory settings from commandline options

    if in_stationdir is not None:
        metadata.stationDir = in_stationdir
    else:
        metadata.stationDir = stationdir
    
    if in_catalogdir is not None:
        metadata.catalogBaseDir = in_catalogdir
    else:
        metadata.catalogBaseDir = catalogdir
        
    if in_colormapdir is not None:
        metadata.colormapDir = in_colormapdir
    else:
        metadata.colormapDir = colormapdir

    if in_combidir is not None:
        metadata.combiDir = in_combidir
    else:
        metadata.combiDir = combidir

    # combinations pickle file
    metadata.combinationsPickleFile = os.path.join( metadata.combiDir, 'combinations-3-200.numpy.pickle' )

    # color map file for prob distros
    metadata.probDistroColorMapFile = os.path.join( metadata.colormapDir, 'rastafari.pov' )

    # set working dircetories
    if in_workingdir is not None:

        # working directory has explicitly been given
        metadata.runPath = in_workingdir

        if not os.path.isdir( metadata.runPath ):
            os.makedirs( metadata.runPath )

        # if grid dir is given, create directly below output dir (one level above pickInfo and distro)
        # no date component in dir name
        # this is to collect all grid files of a sequence of runs
        # if no explicit grid dir is given, create std grid dir below working dir
        if in_griddir is not None:
            metadata.gridDir = os.path.join( metadata.runPath, in_griddir )

        # we know where to find pickInfo and probdistro
        # i.e., skipping of computation is possible
        # set flags for skipping computation
        metadata.skipPickInfo   = in_skippick
        metadata.skipProbDistro = in_skipdistro
        metadata.skipGrid       = in_skipgrid

    else:

        # no explicit working directory, use PID
        metadata.runPath = os.path.join( rundir, str( os.getpid() ) )

        if os.path.isdir( metadata.runPath ):
            error_str = "%s - output path already exists, %s" % ( scriptname, metadata.runPath )
            raise IOError, error_str
        else:
            os.makedirs( metadata.runPath )

    # skipping/combining options
    metadata.combinePickInfoProbDistro = in_combipickprob
    metadata.discardPickInfo           = in_discardpickinfo

    # setting combined pickInfo and prob distro switch overrides individual skip switches
    if metadata.combinePickInfoProbDistro is True:
        metadata.skipPickInfo   = False
        metadata.skipProbDistro = True
            
    # plotting of results
    metadata.plotPickInfo   = in_plotpickinfo
    metadata.plotProbDistro = in_plotdistro

    # set instance dir: runPath + date
    instanceDir = os.path.join( metadata.runPath,
                                mxDateTime2ISO( metadata.useDate, showtime=False ) )
     
    # if new pickInfos or distros are computed:
    # if instance path exists, stop program execution (don't overwrite)
    # if overwrite option is not explicitly set
    if (     os.path.isdir( instanceDir )
         and ( metadata.skipPickInfo is False or metadata.skipProbDistro is False )
         and in_overwrite is False ):
        error_str = "%s - instance directory already exists, %s" % ( scriptname, instanceDir )
        raise IOError, error_str
        
    # create pickInfo and distro dirs
    metadata.pickInfoDir = os.path.join( instanceDir, 'pickInfo/' )
    metadata.distroDir   = os.path.join( instanceDir, 'distro/' )

    metadata.pickInfoPlotDir = os.path.join( instanceDir, 'pickInfo-plot/' )
    metadata.distroPlotDir   = os.path.join( instanceDir, 'distro-plot/' )
                                         
    # create grid dir, if not explicitly set
    if metadata.gridDir is None:
        metadata.gridDir = os.path.join( instanceDir, 'grid/' )

    if in_stationfile_name is not None:
        metadata.stationfile = os.path.join( metadata.stationDir, in_stationfile_name )
    else:
        error_str = "%s - no station file name has been specified" % ( scriptname )
        raise ValueError, error_str

    if in_aliasfile_name is not None:
        metadata.aliasfile = os.path.join( metadata.stationDir, in_aliasfile_name )

    if in_pickcountfile_name is not None:
        metadata.pickCountFile = os.path.join( metadata.stationDir, in_pickcountfile_name )

    if in_polygonfile_name is not None:
        metadata.polygonfile = os.path.join( metadata.stationDir, in_polygonfile_name )
        
    # set metadata file
    metadataFilename = '.'.join( ( metadata_file_prefix,
                                   mxDateTime2ISO( metadata.useDate, showtime=False ),
                                   '00.dat' ) )
    metadata.metadataFile = os.path.join( metadata.runPath, metadataFilename )
    
    # check if metadata file exists, if yes, try another name
    metadataFileCtr = 0
    while ( os.path.isfile( metadata.metadataFile ) ):

        print " skipping metadata file %s" % metadata.metadataFile

        metadataFileCtr += 1
        metadataFilename = '.'.join( ( metadata_file_prefix,
                                       mxDateTime2ISO( metadata.useDate, showtime=False ),
                                       "%02d" % metadataFileCtr,
                                       'dat' ) )
        metadata.metadataFile = os.path.join( metadata.runPath, metadataFilename )
    
    # write metadata to file and screen
    print " write run metadata to file %s" % metadata.metadataFile
    writeMetadata( writeQPData( metadata.metadataFile ) )
    writeMetadata( sys.stdout )

    # set name of profiling file
    if profileFlag is True:
        profileFile = profileFile +'.' + mxDateTime2ISO(
            metadata.useDate, showtime=False ) + ".%02d" % ( metadataFileCtr ) +'.' + mxDateTime2ISO(
            metadata.creationDate, showsecfrac=False ) + '.dat'


def writeMetadata( stream ):
    """
    metadata is a dictionary with relevant metadata of computation run
    """

    global metadata
    
    stream.write( ''.join( ( "=========================================================================================", '\n' ) ) )
    stream.write( ''.join( ( "PMC computation for network %s, date %s" % (
                  metadata.network, mxDateTime2ISO( metadata.useDate, showtime=False ) ), '\n' ) ) )
    stream.write( ''.join( ( " phase data starts at %s (%s days earlier)" % \
                  ( mxDateTime2ISO( metadata.startDate, showtime=False ),
                    metadata.earlierDays ), '\n' ) ) )
    stream.write( ''.join( ( " phase data ends at %s (%s days later)" % \
                  ( mxDateTime2ISO( metadata.endDate, showtime=False ),
                    metadata.laterDays ), '\n' ) ) )
    stream.write( ''.join( ( " run started: %s" % \
                  mxDateTime2ISO( metadata.creationDate, showsecfrac=False ), '\n' ) ) )
    stream.write( ''.join( ( "=========================================================================================", '\n' ) ) )

    if metadata.combinePickInfoProbDistro is True:
        stream.write( ''.join( ( " COMPUTING combined pickInfo and probability distributions", '\n' ) ) )

        if metadata.discardPickInfo is True:
            stream.write( ''.join( ( " -> DISCARDING pickInfo files", '\n' ) ) )
    else:
        if metadata.skipPickInfo is True:
            stream.write( ''.join( ( " SKIPPING pickInfo", '\n' ) ) )
        else:
            stream.write( ''.join( ( " COMPUTING pickInfo", '\n' ) ) )
        if metadata.skipProbDistro is True:
            stream.write( ''.join( ( " SKIPPING probability distributions", '\n' ) ) )
        else:
            stream.write( ''.join( ( " COMPUTING probability distributions", '\n' ) ) )

    if metadata.skipGrid is True:
        stream.write( ''.join( ( " SKIPPING PMC Grid", '\n' ) ) )
    else:
        stream.write( ''.join( ( " COMPUTING PMC Grid", '\n' ) ) )
        
    stream.write( ''.join( ( "=========================================================================================", '\n' ) ) )
    stream.write( ''.join( ( " input files:", '\n' ) ) )
    stream.write( ''.join( ( "  station file: %s" % metadata.stationfile, '\n' ) ) )
    stream.write( ''.join( ( "  alias file: %s" % metadata.aliasfile, '\n' ) ) )
    stream.write( ''.join( ( "  pick count file: %s" % metadata.pickCountFile, '\n' ) ) )
    stream.write( ''.join( ( "  polygon file: %s" % metadata.polygonfile, '\n' ) ) )
    stream.write( ''.join( ( "=========================================================================================", '\n' ) ) )
    stream.write( ''.join( ( " operating on paths:", '\n' ) ) )
    stream.write( ''.join( ( "  pickInfo dir: %s" % metadata.pickInfoDir, '\n' ) ) )
    stream.write( ''.join( ( "  prob distro dir: %s" % metadata.distroDir, '\n' ) ) )
    stream.write( ''.join( ( "  grid dir: %s" % metadata.gridDir, '\n' ) ) )
    stream.write( ''.join( ( "=========================================================================================", '\n' ) ) )
    stream.write( ''.join( ( " computation parameters:", '\n' ) ) )
    stream.write( ''.join( ( "  max. number of stations for combined probabilities: %s" % metadata.useMaxStationCnt, '\n' ) ) )
    stream.write( ''.join( ( "  refine pickInfo: %s" % metadata.refinePickInfo, '\n' ) ) )
    stream.write( ''.join( ( "  minimumPickCount (ignore stations with fewer picks): %s" % metadata.minimumPickCount, '\n' ) ) )
    stream.write( ''.join( ( "=========================================================================================", '\n' ) ) )
    stream.write( ''.join( ( "the command line was:", '\n' ) ) )
    stream.write( ' '.join( sys.argv ) )
    stream.write( '\n' )
    stream.write( ''.join( ( "=========================================================================================", '\n' ) ) )


def compPickInfos( **kwargs ):

    global metadata
    
    print " ========== computing pickInfos =========="

    useDateStr = mxDateTime2ISO( metadata.useDate, showtime=False )
    
    # new inventory object for useDate

    # create output dir
    if not os.path.isdir( metadata.pickInfoDir ):
        print "   creating path %s" % metadata.pickInfoDir
        os.makedirs( metadata.pickInfoDir )

    if (     ( not os.path.isdir( metadata.pickInfoPlotDir ) )
         and metadata.plotPickInfo is True ):
        print "   creating path %s" % metadata.pickInfoPlotDir
        os.makedirs( metadata.pickInfoPlotDir )
        
    if metadata.combinePickInfoProbDistro is True:
        if not os.path.isdir( metadata.distroDir ):
            print "   creating path %s" % metadata.distroDir
            os.makedirs( metadata.distroDir )

            if (     ( not os.path.isdir( metadata.distroPlotDir ) )
                 and metadata.plotProbDistro is True ):
                print "   creating path %s" % metadata.distroPlotDir
                os.makedirs( metadata.distroPlotDir )

    print " importing PMC station file: %s" % metadata.stationfile
    pmc = PMCFactory().createPMC( metadata.network, **kwargs )
    
    pmc.distStyle = 0
    pmc.importStations( metadata.stationfile, encoding='html' )
    
    # set timeZoneShift in metadata
    metadata.timeZoneShift = pmc.timeZoneShift
    
    if metadata.aliasfile is not None:
        print " importing PMC alias file: %s" % metadata.aliasfile
        pmc.importAliases( metadata.aliasfile )

    print " before preprocess: PMC inventory has %s stations" % len( pmc.stations )
    pmc.preprocessInventory( [ metadata.startDate, metadata.endDate ] )
    print " after inventory preprocess: inventory has %s stations" %  len( pmc.stations )

    pmc.getCatalogFiles( metadata )

    print " ----- assign picks for %s stations, using %s catalog files" % ( len( pmc.stations ),
                                                                            len( metadata.catalogFiles ) )
                                                                            
    pmc.assignPicks( metadata=metadata, compression='bz2', verbose_level=kwargs['verbose_level'] )

    if metadata.aliasfile is not None:
        pmc.mergeAliasStations()
        print " after merging aliases: inventory has %s stations" %  len( pmc.stations )

    
def compProbDistros( **kwargs ):

    global metadata

    print " ========== computing probability distros =========="

    # create output dir
    if not os.path.isdir( metadata.distroDir ):
        print "   creating path %s" % metadata.distroDir
        os.makedirs( metadata.distroDir )

    if (     ( not os.path.isdir( metadata.distroPlotDir ) )
         and metadata.plotProbDistro is True ):
        print "   creating path %s" % metadata.distroPlotDir
        os.makedirs( metadata.distroPlotDir )

    # get pick info files
    pickinfofiles = sorted( glob.glob( os.path.join( metadata.pickInfoDir, '*.pickInfo.xml.bz2' ) ) )
    
    # loop over stations
    for pickinfofile in pickinfofiles:

        print " processing pick info file %s" % ( pickinfofile )

        # create PMC for current station
        curr_pmc = PMCFactory().createPMC( metadata.network )
        curr_pmc.distStyle = metadata.useDistStyle
        
        curr_pmc.readXML( getQPDataSource( pickinfofile, compression='bz2' ) )

        # fillup
        curr_pmc.stations[0].distribution.restoreDistStyle( metadata.useDistStyle )
        curr_pmc.stations[0].distribution.setSmooth( metadata.smoothDist )

        print " ----- now calling fillup ..."
        curr_pmc.stations[0].fillup( curr_pmc._calcMagnitudeFromDistance )

        # delete pick info
        if hasattr( curr_pmc.stations[0].distribution, 'pickInfo' ):
            del curr_pmc.stations[0].distribution.pickInfo

        # write distro to disk
        curr_pmc.writeProbDistros( metadata, verbose_level=kwargs['verbose_level']  )

        # plot prob distro
        if metadata.plotProbDistro is True:

            if verbose_level is not None and ( verbose_level > 1 ):
                print " --- creating plot of prob distro"

            curr_pmc.plotDistribution( curr_pmc.stations[0],
                                       metadata,
                                       colortable=metadata.probDistroColorMapFile )
        
        # delete station just processed
        del curr_pmc

        
def compGrids( **kwargs ):
    """
    Note: output grid file is gzipped
          do not use built-in pretty-print (uses too much memory)
          use 'xmllint --format' for pretty-print
          has name <prefix>.YYYY-MM-DDTHH-MM-SS.xml.gz
    """

    global metadata
    global verbose_grid

    print " ========== computing JMA PMC grid =========="
    
    print " ----- get PMC from XML -----  "

    # create output dir
    if not os.path.isdir( metadata.gridDir ):
        print "   creating path %s" % metadata.gridDir
        os.makedirs( metadata.gridDir )
        
    print " load Combinations from pickle"
    combi = unpickleObj( metadata.combinationsPickleFile )
        
    pmc = PMCFactory().createPMC( metadata.network, combinations=combi )
    
    # import PMC distro data: loop over distro files for stations
    distrofiles = sorted( glob.glob( os.path.join( metadata.distroDir, '*.distro.xml.bz2' ) ) )

    for curr_sta_file in distrofiles:

        print " importing PMC from XML file: %s" % curr_sta_file

        curr_pmc = PMCFactory().createPMC( metadata.network )
        curr_pmc.readXML( getQPDataSource( curr_sta_file, compression='bz2' ) )

        # copy over PMC for station to master PMC object
        pmc.stations.append( curr_pmc.stations[0] )

        del curr_pmc

    print " loaded PMC from XML: %s stations" % ( len( pmc.stations ) )

    print " ----- compute QPGrid -----  "

    ## compute grid

    # create polygon object for border of region
    if metadata.polygonfile is None:
        polygon = QPPolygon( metadata.areaPolygon['vertices'] )
    else:
        polygon = QPPolygon( metadata.polygonfile )
        metadata.areaPolygon['vertices'] = polygon.vertices
    
    g = PMCGridN( metadata )
    g.setGridParameter( metadata.areaPolygon['metadata'] )
    g.setupPolygon( polygon, metadata )

    # add annotation object to grid
    g.annotation = pmc.annotation

    g.annotation.setProperty( date                 = utc(),
                              starttime            = metadata.useDate,
                              endtime              = metadata.useDate,
                              lonmin               = g.extent['lonMin'],
                              lonmax               = g.extent['lonMax'],
                              latmin               = g.extent['latMin'],
                              latmax               = g.extent['latMax'],
                              boundary             = polygon.vertices,
                              observationStartTime = metadata.startDate,
                              observationEndTime   = metadata.endDate,
                              maxStationCnt        = metadata.useMaxStationCnt )
                                
    print " compute probability map for Date: %s" % ( metadata.useDate )
    pmc.getProbabilityMap( g, metadata, verbose=verbose_grid )

    # add station data to grid - delete not required fields
    g.stations = []
    for curr_sta in pmc.stations:
        del curr_sta.distribution
        del curr_sta.channels
        del curr_sta.onTime
        del curr_sta.onTimeRefined
        del curr_sta.interPickTime

        g.stations.append( curr_sta )

    mapdate_str = mxDateTime2ISO( [ metadata.useDate ], timesepreplacechar='-', showsecfrac=False )

    grid_file_temp = os.path.join( metadata.gridDir, '.'.join( ( metadata.gridFilePrefix, mapdate_str, 'short.xml' ) ) )
    grid_file_out  = os.path.join( metadata.gridDir, '.'.join( ( metadata.gridFilePrefix, mapdate_str, 'xml.gz' ) ) )

    print " write uncompressed, compact grid file %s" % grid_file_temp
    g.writeXML( grid_file_temp, prettyPrint=False, noStationChildElements=True )
    
    print " creating pretty-printed, compressed grid file: %s" % grid_file_out
    exec_str = "xmllint --format %s | gzip > %s" % ( grid_file_temp, grid_file_out )
    os.system( exec_str )

    # remove compact file
    os.remove( grid_file_temp )


if __name__ == "__main__":
    main()
