#!/usr/bin/env python
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# Functionality provided: Prepare input 3D-wind fields in hybrid coordinates + surface fields for FLEXPART runs
#
# Creation: October  2014 - Anne Fouilloux - University of Oslo
# Extension November 2015 - Leopold Haimberger - University of Vienna for:
# - using the WebAPI also for general MARS retrievals
# - job submission on ecgate and cca
# - job templates suitable for twice daily operational dissemination
# - dividing retrievals of longer periods into digestable chunks
# - retrieve also longer term forecasts, not only analyses and short term forecast data
# - conversion into GRIB2
# - conversion into .fp format for faster execution of FLEXPART
#
#
# Further documentation may be obtained from www.flexpart.eu
#
# Requirements:
# in addition to a standard python 2.6 or 2.7 installation the following packages need to be installed
# ECMWF WebMARS, gribAPI with python enabled, emoslib, ecaccess web toolkit, all available from https://software.ecmwf.int/
# dateutils
# matplotlib (optional, for debugging)
#
# Get MARS GRIB fields from ECMWF for FLEXPART
#

#import socket

#hostname=socket.gethostname()
#ecapi= 'ecmwf' not in hostname
try:
    ecapi=True
    import ecmwfapi
except ImportError:
    ecapi=False

import calendar
import shutil
import datetime
import time
import os,glob,sys,inspect
#from string import strip
from argparse import ArgumentParser,ArgumentDefaultsHelpFormatter
# add path to submit.py to pythonpath so that python finds its buddies
localpythonpath=os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
if localpythonpath not in sys.path:
    sys.path.append(localpythonpath)

from FlexpartTools import ECFlexpart,  \
                          Control, myerror, normalexit, \
                          interpret_args_and_control


def getMARSdata(args, c):


    if not os.path.exists(c.inputdir):
        os.makedirs(c.inputdir)
    print("start date %s " % (c.start_date))
    print("end date %s " % (c.end_date))

    if ecapi:
        server = ecmwfapi.ECMWFService("mars")
    else:
        server = False

    c.ecapi = ecapi
    print 'ecapi:', c.ecapi

# Retrieve ERA interim data for running flexpart
#AP change this variant to correct format conversion with datetime
#AP import datetime and timedelta explicitly
    syear = int(c.start_date[:4])
    smonth = int(c.start_date[4:6])
    sday = int(c.start_date[6:])
    start = datetime.date(year=syear, month=smonth, day=sday)
    startm1 = start - datetime.timedelta(days=1)
    if c.basetime == '00':
        start = startm1
    eyear = int(c.end_date[:4])
    emonth = int(c.end_date[4:6])
    eday = int(c.end_date[6:])
    end = datetime.date(year=eyear, month=emonth, day=eday)
    if c.basetime == '00' or c.basetime == '12':
        endp1 = end + datetime.timedelta(days=1)
    else:
        endp1 = end + datetime.timedelta(days=2)

    datechunk = datetime.timedelta(days=int(c.date_chunk))
    print 'removing content of ' + c.inputdir
    tobecleaned = glob.glob(c.inputdir + '/*_acc_*.' + str(os.getppid()) + '.*.grb')
    for f in tobecleaned:
        os.remove(f)

    times=None
    if c.maxstep<24:
        day=startm1
        while day<endp1:
                # we need to retrieve MARS data for this period (maximum one month)
                flexpart = ECFlexpart(c,fluxes=True)
                if day+datechunk-datetime.timedelta(days=1)<endp1:
                    dates= day.strftime("%Y%m%d") + "/to/" + (day+datechunk-datetime.timedelta(days=1)).strftime("%Y%m%d")
                else:
                    dates= day.strftime("%Y%m%d") + "/to/" + end.strftime("%Y%m%d")

                print "retrieve " + dates + " in dir " + c.inputdir
                try:
                    flexpart.retrieve(server, dates, times, c.inputdir)
                except IOError:
                    myerror(c,'MARS request failed')

                day+=datechunk
    else:
        day=start
        while day<=end:
                # we need to retrieve MARS data for this period (maximum one month)
                flexpart = ECFlexpart(c,fluxes=True)
                if day+datechunk-datetime.timedelta(days=1)<end:
                    dates= day.strftime("%Y%m%d") + "/to/" + (day+datechunk-datetime.timedelta(days=1)).strftime("%Y%m%d")
                else:
                    dates= day.strftime("%Y%m%d") + "/to/" + end.strftime("%Y%m%d")

                print "retrieve " + dates + " in dir " + c.inputdir
                flexpart.retrieve(server, dates, times, c.inputdir)
                day+=datechunk



    tobecleaned=glob.glob(c.inputdir+'/*__*.'+str(os.getppid())+'.*.grb')
    for f in tobecleaned:
        os.remove(f)
    day=start
    times=None
    while day<=end:

               # we need to retrieve MARS data for this period (maximum one month)
            flexpart = ECFlexpart(c)
            if day+datechunk-datetime.timedelta(days=1)<end:
                dates= day.strftime("%Y%m%d") + "/to/" + (day+datechunk-datetime.timedelta(days=1)).strftime("%Y%m%d")
            else:
                dates= day.strftime("%Y%m%d") + "/to/" + end.strftime("%Y%m%d")
            print "retrieve " + dates + " in dir " + c.inputdir

            flexpart.retrieve(server, dates, times, c.inputdir)
            day+=datechunk


if __name__ == "__main__":

    args,c=interpret_args_and_control()
    getMARSdata(args,c)
    normalexit(c)