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

#
# scsn_compGrid.py
# $Id: scsn_compGrid.py 157 2009-02-16 11:15:52Z 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 sys
import getopt
import os
import glob
import re
import random
import datetime

import pickle
import cPickle

from mx.DateTime     import Date
from mx.DateTime     import DateTime
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 pmc.PMC      import *
from pmc.PMC_SCSN import *
from pmc.PMCData  import *

from QPGrid           import *

# os.path.dirname( path )  = os.path.split( path )[0]
# os.path.basename( path ) = os.path.split( path )[1]
# NOTE: os.path.split( '/foo/bar/' ) = ( '/foo/bar', '' )
#       os.path.split( '/foo/bar' )  = ( '/foo', 'bar' )
# NOTE: dirname has NO trailing slash

# 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] ) )
basepath   = os.path.dirname( os.path.abspath( sys.argv[0] ) ) # no trailing slash

# directory where catalog files in stp phase format are
catdir   = os.path.join( basepath, 'data/phase/' )

# location of station list file
stationfile_path = os.path.join( basepath, 'data/station/' )
stationfile      = None

# location of station alias file
aliasfile = None

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

# directory where pick info files will be written to
pickInfoDir = None

# directory where yearly probability distributions will be written to
distroDir = None

# PMC grid file prefix
gridDir = None

gridfile_prefix   = 'scsn'

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

# distStyle for probdistros
useDistStyle = 3

# max stations used for combined probability
useMaxStationCnt = 200  

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

# stuff for result XML
maggrid       = frange( 0.0, 4.0, 0.1 )
mp_prob_list  = [ 0.9, 0.95, 0.99, 0.999 ]

# set geographical region for completeness evaluation
geobox = { 'lonmin': -122.0, 'lonmax': -113.5, 
           'latmin':   31.5, 'latmax':   38.0,
           'londelta':  0.1, 'latdelta':  0.1,
           'depthmin':  7.5, 'depthmax':  7.5 }

# location of combinations pickle
combi_path   = os.path.join( basepath, '../jma/data/combinations' )
combi_pickle = os.path.join( combi_path, 'combinations-3-200.numpy.pickle' )

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

useDate   = None
startDate = None
endDate   = None

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

def main():

    setUp()
    compPickInfos()
    compProbDistros()
    compGrids()

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

def setUp():

    global rundir_base
    
    global useDate
    global startDate
    global endDate
    
    global stationfile_path
    global stationfile
    global aliasfile

    global runPath
    global pickInfoDir
    global distroDir
    global gridDir

    aliasfile_name    = None
    date_in           = None
    griddir_in        = None
    stationfile_name  = None
    earlier_in        = None
    later_in          = None
    workingdir_in     = None
    
    # Read commandline arguments
    
    cmdParams = sys.argv[1:]
    opts, args = getopt.gnu_getopt( cmdParams,
                                    'ha:d:g:s:t:T:w:',
                                    [ 'help', 'aliasfile=', 'date=', 'griddir=',
                                      'stationfile=', 'timebefore=', 'timeafter=',
                                      'workingdir=' ] )

    for option, parameter in opts:

        if option == '-a' or option == '--aliasfile':
            aliasfile_name = parameter

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

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

        if option == '-s' or option == '--stationfile':
            stationfile_name = parameter

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

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

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

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

    try:
        # useDate as mxDateTime instance
        useDate     = ParseDateTimeUTC( date_in )
        earlierDays = int( earlier_in )
        laterDays   = int( later_in )
    except:
        error_str = " illegal date format %s or non-integer time interval" % date_in
        raise ValueError, error_str

    # get start and end dates
    startDate = useDate - DateTimeDelta( float( earlierDays ) )
    endDate   = useDate + DateTimeDelta( float( laterDays ) )

    # set working dircetories
    if workingdir_in is not None:
        runPath = os.path.join( rundir_base, workingdir_in )

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

        # if grid dir is given, create directly below output dir
        # no date in dir name, this is to collect all grid files of a sequence of runs
        if griddir_in is not None:
            gridDir = os.path.join( runPath, griddir_in )

            if not os.path.isdir( gridDir ):
                os.makedirs( gridDir )

    else: 
        runPath = os.path.join( rundir_base, str( os.getpid() ) )

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

        # create grid directory 
        gridDir = os.path.join( runPath, mxDateTime2ISO( useDate, showtime = False ), 'grid/' )
        os.makedirs( gridDir )

    # create pickInfo and distro dirs
    pickInfoDir = os.path.join( runPath, mxDateTime2ISO( useDate, showtime = False ), 'pickInfo/' )
    distroDir   = os.path.join( runPath, mxDateTime2ISO( useDate, showtime = False ), 'distro/' )

    # if working directory was explicitly set, directories may already exist
    if not os.path.isdir( pickInfoDir ):
        os.makedirs( pickInfoDir )

    if not os.path.isdir( distroDir ):
        os.makedirs( distroDir )

    if stationfile_name is not None:
        stationfile = os.path.join( stationfile_path, stationfile_name )
    else:
        error_str = "error: no station file given"
        raise IOError, error_str

    if aliasfile_name is not None:
        aliasfile = os.path.join( stationfile_path, aliasfile_name )
    
    print
    print "========================================================================================="
    print "SCSN: starting PMC-QPGrid computation for %s, from %s to %s" % \
           ( mxDateTime2ISO( useDate, showtime = False ),
             mxDateTime2ISO( startDate, showtime = False ),
             mxDateTime2ISO( endDate, showtime = False ) )
    print "========================================================================================="
    print
    print " operating on paths:"
    print "  %s" % pickInfoDir
    print "  %s" % distroDir
    print "  %s" % gridDir
    print
    
def compPickInfos():

    global useDate
    global startDate
    global endDate
    
    global stationfile
    global aliasfile

    global pickInfoDir
    
    print " ========== computing pickInfos =========="

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

    print " importing PMC station file: %s" % stationfile
    pmc = PMC_SCSN( stationfile, 0, smooth = True, encoding = 'html' )

    if aliasfile is not None:
        print " importing PMC alias file: %s" % aliasfile
        pmc.importAliases( aliasfile )

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

    # loop over yeary catalog chunks
    for time_year in xrange( startDate.year, endDate.year+1 ):
        
        # get catalog decks for monthly chunk
        print " ----- processing year %s" % time_year

        catfile = os.path.join( catdir, ''.join( ( 'phase', str(time_year), '.dat.gz' ) ) )

        print "   ----- importing catalog chunk from file: ", catfile

        qpc = QPCatalog( idstyle='numeric' )
        qpc.importSTPPhase( catfile, compression='gz', minimumDataset=True )
        print "   catalog chunk has %s events" % qpc.size()

        pmc.preprocessCatalog( qpc, [ startDate, endDate ] )
        print "   after catalog preprocess: catalog chunk has %s events" % qpc.size()

        print "   assigning picks ... "
        pmc.assignPicks( qpc, reduce_catalog = False, verbose_level = 2 )

        del qpc

    pmc.mergeAliasStations()
    print " after merging aliases: inventory has %s stations" %  len( pmc.stations )
    
    print " ----- now writing station data"

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

    for curr_sta_idx in reversed( xrange( len( pmc.stations ) ) ):

        curr_sta = pmc.stations[curr_sta_idx]

        print "   ----- processing station %s: %s %s" % ( curr_sta_idx, curr_sta.networkCode, curr_sta.stationCode )

        # write pickInfo to disk if not empty
        if curr_sta.distribution.pickInfo is not None:

            sta_filename_xml = '.'.join( ( curr_sta.networkCode, curr_sta.stationCode, 'pickInfo.xml.bz2' ) )
            sta_file_xml     = os.path.join( pickInfoDir, sta_filename_xml )

            print "   ----- writing to XML"

            curr_stream = cStringIO.StringIO()
                
            curr_stream.write( '<?xml version="1.0" encoding="utf-8"?>' )
            curr_stream.write( '<PMC>' )
            curr_sta.toXML( 'PMCStation', curr_stream )
            curr_stream.write( '</PMC>' )

            xmlPrettyPrint( curr_stream, writeQPData( sta_file_xml, compression='bz2' ) )

        else:
            print "   ----- EMPTY"

        # delete station just processed
        del pmc.stations[curr_sta_idx]

def compProbDistros():

    global stationfile
    
    global pickInfoDir
    global distroDir

    global startDate
    global endDate

    global useDistStyle

    print " ========== computing probability distros =========="
    
    print " importing PMC inventory file: ", stationfile
    pmc = PMC_SCSN( stationfile )
    print " full inventory has %s stations" % len( pmc.stations )

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

    # create output dir
    if not os.path.isdir( distroDir ):
        print "   creating path %s" % distroDir
        os.makedirs( distroDir )
                
    # loop over stations
    for curr_sta_idx in reversed( xrange( len( pmc.stations ) ) ):

        curr_sta = pmc.stations[curr_sta_idx]

        print " processing station %s: %s %s" % ( curr_sta_idx, curr_sta.networkCode, curr_sta.stationCode )

        # create PMC for current station
        curr_pmc = PMC_SCSN()

        sta_filename_xml = '.'.join( ( curr_sta.networkCode, curr_sta.stationCode, 'pickInfo.xml.bz2' ) )
        sta_file_xml     = os.path.join( pickInfoDir, sta_filename_xml )
          
        if not os.path.isfile( sta_file_xml ):
            print "  EMPTY" 

        else:
            curr_pmc.readXML( getQPDataSource( sta_file_xml, compression='bz2' ) )

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

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

            # delete pick info
            del curr_pmc.stations[0].distribution.pickInfo

            # get filename for probdistro XML
            distro_filename_xml = '.'.join( ( curr_sta.networkCode, curr_sta.stationCode, 'distro.xml.bz2' ) )
            distro_file_out     = os.path.join( distroDir, distro_filename_xml )

            # write PMC as XML
            print " writing distro XML file: ", distro_file_out
            curr_pmc.writeXML( distro_file_out, compression='bz2', prettyPrint = True )

        # delete station just processed
        del curr_pmc

def compGrids():
    """
    Note: output grid file is prettyfied and gzipped
          has name <prefix>.YYYY-MM-DDTHH-MM-SS.xml.gz
    """

    global distroDir
    global gridDir

    global useDistStyle

    global gridfile_prefix
    
    global useDate
    global startDate
    global endDate

    global maggrid
    global mp_prob_list

    global geobox

    global combi_path
    global combi_pickle

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

    # create output dir
    if not os.path.isdir( gridDir ):
        print "   creating path %s" % gridDir
        os.makedirs( gridDir )
        
    print " load Combinations from pickle"
    combi = unpickleObj( combi_pickle )
        
    pmc = PMC_SCSN( combinations = combi )

    # import PMC distro data: loop over distro files for stations
    distrofiles = sorted( glob.glob( os.path.join( distroDir, '*.distro.xml.bz2' ) ) )

    for curr_sta_file in distrofiles:

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

        curr_pmc = PMC_SCSN()
        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

    g = QPGrid()
    g.setGridParameter( { 'lonDelta': geobox['londelta'], 'latDelta': geobox['latdelta'] } )
    g.annotation = pmc.annotation

    g.setupBox( geobox['lonmin'], geobox['lonmax'], 
                geobox['latmin'], geobox['latmax'], 
                geobox['depthmin'], geobox['depthmax'] )

    # mapdate = Date( useDate[0], useDate[1], useDate[2] )
    mapdate_str = mxDateTime2ISO( [ useDate ], timesepreplacechar = '-', showsecfrac = False )

    # set grid annotation
    g.annotation.setProperty( date = utc(),
                              starttime = useDate,
                              endtime   = useDate,
                              latmin = geobox['latmin'],
                              latmax = geobox['latmax'],
                              lonmin = geobox['lonmin'],
                              lonmax = geobox['lonmax'],
                              observationStartTime = startDate,
                              observationEndTime   = endDate )
                                
    print " compute probability map for Date: %s" % ( mapdate_str )
    pmc.getProbabilityMap( g, maggrid, useDate, useMaxStationCnt, mp_prob_list, verbose = False )

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

    grid_file_out = os.path.join( gridDir, '.'.join( ( gridfile_prefix, mapdate_str, 'xml.gz' ) ) )

    print " write grid file %s" % grid_file_out
    g.writeXML( grid_file_out, compression='gz', prettyPrint = True )

def PrintHelp():
    global scriptname
    
    print 'computes PMC grid'
    print 'Usage: %s [OPTION]' % scriptname
    print '  Options'
    print '   -a VALUE, --aliasfile=<value>    Alias file name'
    print '   -d VALUE, --date=<value>         Date for which PMC grid is computed (YYYY-MM-DD)'
    print '   -g VALUE, --griddir=<value>      Directory to place PMC grid file'
    print '   -s VALUE, --stationfile=<value>  Station file name'
    print '   -t VALUE, --timebefore=<value>   Days before date'
    print '   -T VALUE, --timeafter=<value>    Days after date'
    print '   -w VALUE, --workingdir=<value>   Working directory'
    print '   -h, --help                       print this information'

if __name__ == "__main__":
    main()