#!/usr/bin/python3

import sys, getopt
from datetime import datetime, date, time, tzinfo
import re
import math
import matplotlib.pyplot as plt

#sys.path.append ('lib64')
#sys.path.append ('lib64/python3.12/site-packages/pytrademgen')

#------------------------------------------------------------------------------	
def usage (script_name):
    print
    print " That script extracts only the booking request data from the log"
    print " file generated by the 'trademgen_generateDemand' binary."
    print " It then plots the print corresponding to booking arrival patterns."
    print
    print " Usage: %s [options]" % script_name
    print
    print "  -h, --help           : outputs help and quits"
    print "  -o, --output <path>  : path to output file (if blank, stdout)"
    print "  <path>               : input file"
    print "                         Default: 'trademgen_generateDemand.log'"
    print "  -p, --plot           : plot (if blank, no plot is produced)"
    print

#------------------------------------------------------------------------------	
def handle_opt():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "ho:p", 
                                   ["help", "output", "plot"])
    except getopt.GetoptError, err:
        # Will print something like "option -a not recognized"
        print str (err)
        usage (sys.argv[0])
        sys.exit (2)
	
    # Default options
    plotFlag = False
    input_filename = 'trademgen_generateDemand.log'
    output_filename = ''
    input_file = sys.stdin
    output_file = sys.stdout

    # Input stream/file
    if len (args) != 0:
        input_filename = args[0]

    # Handling
    for o, a in opts:
        if o in ("-h", "--help"):
            usage (sys.argv[0])
            sys.exit()
        elif o in ("-o", "--output"):
            output_filename = a
        elif o in ("-p", "--plot"):
            plotFlag = True
        else:
            assert False, "Unhandled option"

    # Input file (bookign requests). That file is normally compressed
    # with GNU Zip (gzip)
    if (input_filename != ''):
	flag_gz = True
	if len (input_filename) < 2:
            flag_gz = False
	elif input_filename[-2:] != 'gz':
            flag_gz = False
	if flag_gz:
            input_file = gzip.open (input_filename, 'rb')
	else:
            input_file = open (input_filename, 'r')

    if (output_filename != ''):
        output_file = open (output_filename, 'w')

    #
    print "Produces a plot: " + str(plotFlag)
    print "Input stream/file: '" + input_filename + "'"
    print "Output stream/file: '" + output_filename + "'"
    return plotFlag, input_file, output_file

#------------------------------------------
def plot_simple(yVals):
    plt.plot(yVals,'b-o')
    plt.show()

#--------------------------------------------
def main():
    # Parse command options
    plotFlag, input_file, output_file = handle_opt()

    # Iterate while there is something to read from the file
    bkgIdx = 0

    # The O&D statistics holder is a dictionary made of:
    #  - a key, itself made of the origin, destination and flight departure date
    #  - the list of booking dates corresponding to all the booking requests
    OnDStatsHolder = {}

    # The O&D statistics object is a dictionary made of:
    #  - a key, itself made of the origin, destination and flight departure date
    #  - the cumulative number of booking requests made at the corresponding
    #    booking date for that key
    OnDStats = {}

    # Iterate while there are booking requests in the TraDemGen log file
    for line in input_file:
        # The lines of interest contain the "Poped" key-word
        matchedLine = re.search ("BKG", line)

        # If there is no match, the loop will continue on to the next line
        if not matchedLine:
            continue

        # There is a match. Now, the booking request information will
        # be extracted
        extractor = re.compile ("^.*(20[0-9]{2}-[A-Z][a-z]{2}-[0-9]{2}).*([A-Z]{3}-[A-Z]{3}).*(20[0-9]{2}-[A-Z][a-z]{2}-[0-9]{2}).*$")
        extractedList = extractor.search (line)

        # DEBUG
        #print ("1: " + extractedList.group(1) 
        #       + ", 2: " + extractedList.group(2)
        #       + ", 3: " + extractedList.group(3))

        #
        bkgIdx += 1

        # Booking date
        bookingDateStr = extractedList.group(1)
        bookingDate = datetime.strptime (bookingDateStr, "%Y-%b-%d")

        # O&D
        OnD = extractedList.group(2).split('-')
        origin = OnD[0]
        destination = OnD[1]

        # Flight departure date
        departureDateStr = extractedList.group(3)
        departureDate = datetime.strptime (departureDateStr, "%Y-%b-%d")

        # Create an entry key, made of the origin, destination and
        # flight departure date
        entryKey = (origin, destination, departureDate.strftime ("%Y-%m-%d"))
        if entryKey in OnDStatsHolder:
            # There are already entries for that key (i.e., origin, destination
            # and departure date). So, just add the new booking request
            # characteristics for that key.
            #
            # The cumulative number of booking requests is incremented
            # (as well as the corresponding list, used only to ease
            # later plotting).
            nbOfBkgs = OnDStats[entryKey][-1]
            OnDStats[entryKey].append (nbOfBkgs+1)
            OnDStatsHolder[entryKey].append (bookingDate.toordinal())

        else:
            # There is no entry yet for that key (i.e., origin,
            # destination, departure date). So, create a new entry
            # with the booking request characteristics for that key.
            OnDStats[entryKey] = [1]
            OnDStatsHolder[entryKey] = [bookingDate.toordinal()]

        # DEBUG
        #print "BookingRequest[" + str(bkgIdx) + "]: ", bookingRequest,"\n"

    # Calculate the colors
    # Note that xFF == 256  
    nbOfStreams = len (OnDStatsHolder)
    if (nbOfStreams == 0):
        print ("No booking request information found in the '"+ input_file.name + "' file.")
        quit()  
        
    assert (nbOfStreams != 0)
    colorIncrementList = [0xFF / nbOfStreams, 0xFF / nbOfStreams, 
                          0xFF / nbOfStreams] 
    colorIdxList = [0, 0, 0]
    colorIdx = 0

    #print "Color increments: " + str(colorIncrementList)

    # Browse through the (O&D, departure date) statistics
    for (OnDEntryKey, OnDEntryList) in OnDStatsHolder.iteritems():
        #
        OnDEntryListRegular = [date.fromordinal(d).strftime ("%Y-%m-%d")
                               for d in OnDEntryList]
        output_file.write ("[" + str(OnDEntryKey) + "] " 
                           + str(OnDEntryListRegular) + "\n")
        output_file.write ("[" + str(OnDEntryKey) + "] " 
                           + str(OnDStats[OnDEntryKey]) + "\n")

        # Plot with MatplotLib
        if plotFlag:
            # Convert the decimal into the wanted hexadecimal
            # (see, e.g., http://matplotlib.sourceforge.net/api/pyplot_api.html#matplotlib.pyplot.plot)
            colorIdxStr = re.sub("0x", "", hex(colorIdx)).zfill(6)
            colorStr = "#" + colorIdxStr

            #print "Color string: " + colorStr

            # Plot with solid line ('-') and blue ('b') circles ('o')
            plt.plot_date (sorted(OnDEntryList), OnDStats[OnDEntryKey], 
                           color=colorStr, linestyle='solid', marker='o')

        # Increment each component separately
        colorIdxList = [(x[0]+x[1])
                        for x in zip(colorIdxList, colorIncrementList)]
        # print ("colorIdxList: '" + str(colorIdxList) + "'")

        # Recompose the global color number
        colorIdx = int(colorIdxList[0])*0x010000 + int(math.sqrt(colorIdxList[1]))*0x0100 + int(math.log10(colorIdxList[2]))*0x01

    #
    output_file.write ("\n")

    # Plot with MatplotLib
    if plotFlag:
        plt.show()

#--------------------------------------------
if __name__ == "__main__":
    main()

