#!/usr/bin/env python

"""
MODULE:    t.rast.aggregate.update

AUTHOR(S): Stefan Blumentrath < stefan.blumentrath AT nina.no>

PURPOSE:   Update a STRDS generated by t.rast.aggregate from updated input STRDS

COPYRIGHT: (C) 2018 by the GRASS Development Team

           This program is free software under the GNU General Public
           License (>=v2). Read the file COPYING that comes with GRASS
           for details.
"""

"""
To Dos:
- use proper cleanup routine, esp if using csv + vrt (copy from other modules)
- handle layers in mask input
- add progress bar
- make date_from and date_to dependent on each other or use today as date_to if not specified

"""

#%module
#% description: Update a STRDS generated by t.rast.aggregate from updated input STRDS
#% keyword: temporal
#% keyword: aggregate
#%end

#%option G_OPT_STRDS_INPUT
#%end

#%option
#% key: completeness
#% type: integer
#% description: Percentage of completeness for granules to add/update
#% answer: 100
#% required: no
#%end

#%flag
#% key: n
#% description: Do not check for potenitally reprocessed input
#%end


import sys
import os
import grass.script as grass
import grass.pygrass.modules.interface as mod_iface
from datetime import datetime
from grass.temporal import datetime_math
from grass.temporal.core import get_current_mapset

from copy import deepcopy

import grass.temporal as tgis


if not "GISBASE" in os.environ.keys():
    grass.message("You must be in GRASS GIS to run this program.")
    sys.exit(1)


def cmd2dict(command):
    """Parse GRASS shell command string (e.g. from history) into Python dict

    :param command: GRASS shell command string
    :returns: Module name as string and Python dict with command options and flags
    :type command: string
    :rtype: string, dict

    :Example:

    >>> cmd2dict("v.db.select -r map=roadsmajor where=\"ROAD_NAME = 'NC-98'\")
    ('t.rast.aggregate',
     {'basename': 'snow_days_seNorge_1km_years',
      'granularity': '1 years',
      'input': 'snow_bin_seNorge_1km_days',
      'method': 'sum',
      'nprocs': '10',
      'output': 'snow_days_seNorge_1km_years',
      'overwrite': True,
      'quiet': True,
      'where': 'start_time >=  1958-01-01  AND start_time <=  2016-12-31 '})
    """ 

    import shlex
    command_dict = {}

    name = command.split(' ')[0]

    if command.find('--v') > -1 or command.find('--verbose') > -1:
        command_dict['verbose'] = True
    elif command.find('--q') > -1 or command.find('--quiet') > -1:
        command_dict['quiet'] = True

    if command.find('--o') > -1 or command.find('--overwrite') > -1:
        command_dict['overwrite'] = True

    for opt in shlex.split(command):
        if opt.find('=') > 0:
            optlist = opt.split('=', 1)
            opt_str = optlist[1].lstrip(' ').rstrip(' ')
            #if opt_str.find(' ') > 0:
            #    command_dict[optlist[0]] = '"{}"'.format(opt_str)
            #else:
            command_dict[optlist[0]] = opt_str
        if opt.startswith('-') and opt.find('--') <0:
            command_dict['flags'] = opt.lstrip('-')

    return name, command_dict

def genRandomName(length):
    """Generate a random name of length "length" starting with a letter

    :param length: length of the random name to generate
    :returns: String with a random name of length "length" starting with a letter
    :type length: int
    :rtype: string

    :Example:

    >>> genRandomName(12)
    'MxMa1kAS13s9'
    """

    import string
    import random
    chars = string.ascii_uppercase + string.ascii_lowercase + string.digits
    randomname = '1'
    while randomname[0].isdigit():
        randomname = ''.join(random.choice(chars) for _ in range(length))

    return randomname

randname =  genRandomName(21)

def integrateUpdates(agg_strds, input_strds, tempfile):
    """Integrates updates into input STRDS from t.rast.aggregate

    :param agg_strds: STRDS produced by t.rast.aggregate
    :param input_strds: STRDS with updated granules
    :param tempfile: Path to a tempfile
    :returns: None
    :type agg_strds: string
    :type input_strds: string
    :type tempfile: string
    :rtype: None

    :Example:

    >>> integrateUpdates(strds, randname, tmpf)
    """

    with open(tempfile, 'w') as new_maps:
        #new_maps = StringIO()
        new_maps.write(grass.read_command('t.rast.list',
                              input=input_strds, flags='u',
                              columns='name,start_time,end_time').rstrip('\n'))
                              
    grass.run_command('t.remove', inputs=input_strds)
    grass.run_command('t.register', input=agg_strds, file=tempfile, separator='|', overwrite=True)

def gran2where(datetime_start, granularity):
    end = datetime_math.increment_datetime_by_string(datetime_start, granularity, mult=1)
    start_str = datetime.strftime(datetime_start, "%Y-%m-%dT%H:%M:%S.%s")
    end_str = datetime.strftime(end, "%Y-%m-%dT%H:%M:%S.%s")

    where_str = "(end_time >= '{0}'  AND start_time <= '{1}')".format(start_str, end_str)
    return where_str

def checkGranule(instrds, datetime_start, completeness, from_gran, to_gran):
    end = datetime_math.increment_datetime_by_string(datetime_start, to_gran, mult=1)
    start_str = datetime.strftime(datetime_start, "%Y-%m-%dT%H:%M:%S.%s")
    end_str = datetime.strftime(end, "%Y-%m-%dT%H:%M:%S.%s")

    where_str = "(end_time >= '{0}'  AND start_time <= '{1}')".format(start_str, end_str)

    n = 0
    inc = datetime_start
    while inc < end:
        n = n + 1
        inc = datetime_math.increment_datetime_by_string(inc, from_gran, mult=1)

    inmaps = grass.read_command('t.rast.list',
                              input=instrds,
                              where=where_str, flags='u',
                              columns='id,start_time,end_time').rstrip('\n').split('\n')

    if len(inmaps) / float(completeness) >= n / 100.0:
        return where_str
    


def main():

    # Parse input options
    strds = options['input']
    completeness = options['completeness']
    tmpf = grass.tempfile()
    no_checks = flags['n']

    tgis.init()

    dbif = tgis.SQLDatabaseInterfaceConnection()
    dbif.connect()

    stds = tgis.open_old_stds(strds, 'strds', dbif)

    try:
        strdsinfo = grass.parse_command('t.info', flags='g', input=strds.split('@')[0])
    except:
        grass.fatal('{} not in current mapset! Cannot modify.'.format(strds))

    hist = grass.read_command('t.info', flags='h', input=strds).rstrip('\n')

    for l in hist.replace('\\\n','').split('\n'):
         if l.split(' ')[0] == 't.rast.aggregate':
             command = l

    name, command_dict = cmd2dict(command)
    cmapset = get_current_mapset()

    if command_dict['input'].find('@') < 0:
        strds_names = []
        strds_list = grass.read_command('t.list', quiet = True).rstrip('\n').split('\n')
        for s in strds_list:
            if s.startswith('{}@{}'.format(command_dict['input'], cmapset)):
                strds_name = s
            elif s.startswith('{}@'.format(command_dict['input'])):
                strds_names.append(s)
            if len(strds_names) > 1:
                grass.warning('Found more than one STRDS with the same name as the aggregation input')
                grass.warning('Using: {}'.format(strds_names[0]))
                grass.warning('Found also: {}'.format(','.join(strds_names[1:])))

        strds_name = strds_names[0]
        command_dict['input'] = strds_name


    inputinfo = grass.parse_command('t.info', flags='g', input=command_dict['input'])

    # Get all maps for update
    #############################################################################################
    # get possibly reprocessed maps
    updated_granules = []
    if not no_checks:
        agg_maps = grass.read_command('t.rast.list', input=strds, flags='u',
                                      columns='id,name,creation_time,start_time,end_time').rstrip('\n').split('\n')
        for m in agg_maps:
            in_maps = grass.read_command('t.rast.list',
                                         input=command_dict['input'],
                                         where='creation_time > \'{0}\' AND \
                                         end_time >= \'{1}\' AND \
                                         end_time <= \'{2}\' '.format(m.split('|')[2],
                                                                      m.split('|')[3],
                                                                      m.split('|')[4]), flags='u',
                                         columns='id,name,creation_time,start_time,end_time').rstrip('\n').split('\n')

            if in_maps[0] != '':
                if completeness > 0:
                    gran = datetime_math.string_to_datetime(m.split('|')[3])
                    checked_gran = checkGranule(command_dict['input'], gran, completeness, inputinfo['granularity'], strdsinfo['granularity'])
                    if checked_gran:
                        updated_granules.append(checked_gran)
                    else:
                        start_str = datetime.strftime(gran, "%Y-%m-%dT%H:%M:%S.%s")
                        grass.warning('Found only too few maps for granule starting with {}.'.format(gran))
                        grass.warning('Granule is incomplete and therefor not updated.'.format(start_str))

                else:
                    updated_granules.append(gran)

    #############################################################################################
    # get new maps
    added_maps = grass.read_command('t.rast.list',
                                    input=command_dict['input'],
                                    where='start_time >= \'{0}\''.format(strdsinfo['end_time']), flags='u',
                                    columns='id,name,creation_time,start_time,end_time').rstrip('\n').split('\n')

    new_granules = []
    if not added_maps[0] == '':
        # Get unique list of start based on granularity of aggregated map
        added_granules = []
        for m in added_maps:
            dt = datetime_math.string_to_datetime(m.split('|')[3])
            start = datetime_math.adjust_datetime_to_granularity(dt, strdsinfo['granularity'])
            added_granules.append(start)

        added_granules = list(set(added_granules))

        if len(added_granules) > 0:
            if completeness > 0:
                for gran in added_granules:
                    checked_gran = checkGranule(command_dict['input'], gran, completeness, inputinfo['granularity'], strdsinfo['granularity'])
                    if checked_gran:
                        new_granules.append(checked_gran)
                    else:
                        start_str = datetime.strftime(gran, "%Y-%m-%dT%H:%M:%S.%s")
                        grass.warning('Found only too few maps for granule starting with {}.'.format(gran))
                        grass.warning('Granule is incomplete and therefor not updated.'.format(start_str))

            else:
                new_granules = [gran2where(start, strdsinfo['granularity']) for start in added_granules]

    #############################################################################################
    #len(added_granules) + len(updated_granules)
    # Compute number of required maps (n) within aggregated granule
    #twhere = updated_granules + new_granules

    """
    updated = 0
    twhere = []
    for gran_list in (updated_granules, added_granules):
        if len(gran_list) == 0:
            continue
        else:
            for start in gran_list:
                end = datetime_math.increment_datetime_by_string(start, strdsinfo['granularity'], mult=1)
                start_str = datetime.strftime(start, "%Y-%m-%dT%H:%M:%S.%s")
                end_str = datetime.strftime(end, "%Y-%m-%dT%H:%M:%S.%s")

                where_str = "(end_time >= '{0}'  AND start_time <= '{1}')".format(start_str, end_str)

                if check_complete:
                    n = 0
                    inc = start
                    while inc < end:
                        n = n + 1
                        inc = datetime_math.increment_datetime_by_string(inc, inputinfo['granularity'], mult=1)

                    inmaps = grass.read_command('t.rast.list',
                                              input=command_dict['input'],
                                              where=where_str, flags='u',
                                              columns='id,start_time,end_time').rstrip('\n').split('\n')

                    if n != len(inmaps):
                        grass.warning('Found only {0} of {1} maps for granule starting with {2}.'.format(len(inmaps), n, start_str))
                        grass.warning('Granule is incomplete and therefor not updated.'.format(start_str))
                        continue

                twhere.append(where_str)
                updated = updated + 1

                grass.message('Updating {0} of {1} granules.'.format(updated, len(start_dates)))
    """
    update_granule = deepcopy(command_dict)
    update_granule['output'] = randname
    update_granule['overwrite'] = True

    num_suffix = False
    if 'suffix' in command_dict.keys():
        if command_dict['suffix'] == 'num':
            num_suffix = True

    if not num_suffix:
        # Check if it is possible to use OR in where conditions!!!
        #update_granule['where'] = '"{}"'.format(' OR '.join(twhere))
        twhere = updated_granules + new_granules
        if len(twhere) > 0:
            update_granule['where'] = ' OR '.join(twhere)
            #update_granule['where'] = where_str
            if len(twhere) < int(command_dict['nprocs']):
               update_granule['nprocs'] = len(twhere)

            # run module with modified paramters
            mod = mod_iface.module.Module(name, run_=False, **update_granule)
            mod.run()
            integrateUpdates(strds, randname, tmpf)
    elif not no_checks and len(updated_granules) > 0:
        update_granule['nprocs'] = 1
        for w in updated_granules:
            update_granule['where'] = w

            cur_map = grass.read_command('t.rast.list', input=strds, where=w, flags='u',
                                      columns='id,start_time,end_time').rstrip('\n').split('\n')

            mapid = cur_map[0].split('@')[0].split('_')[-1]
            if not mapid.isdigit():
                grass.warning('Could not identify the numerical suffix for granule starting with {}!'.format(w.split("'")[1]))
            else:
                update_granule['offset'] = int(mapid) - 1

            # run module with modified paramters
            mod = mod_iface.module.Module(name, run_=False, **update_granule)
            mod.run()
            integrateUpdates(strds, randname, tmpf)
    else:
        if len(new_granules) > 0:
            if len(new_granules) < int(command_dict['nprocs']):
               update_granule['nprocs'] = len(new_granules)
            update_granule['where'] = ' OR '.join(new_granules)

            old_offset = 0 if not 'offset' in command_dict.keys() else int(command_dict['offset'])
            offset = int(strdsinfo['number_of_maps']) + old_offset
            update_granule['offset'] = offset

            # run module with modified paramters
            mod = mod_iface.module.Module(name, run_=False, **update_granule)
            mod.run()
            integrateUpdates(strds, randname, tmpf)
        

    grass.verbose('Done. Updated {} granules.'.format(len(new_granules) + len(updated_granules)))


# Run the module
# ToDo: Add an atexit procedure which closes and removes the current map
if __name__ == "__main__":
    options, flags = grass.parser()
    sys.exit(main())

    
"""
    
    
def checkGranule(datetime_start, completeness):
    end = datetime_math.increment_datetime_by_string(datetime_start, strdsinfo['granularity'], mult=1)
    start_str = datetime.strftime(datetime_start, "%Y-%m-%dT%H:%M:%S.%s")
    end_str = datetime.strftime(end, "%Y-%m-%dT%H:%M:%S.%s")

    where_str = "(end_time >= '{0}'  AND start_time <= '{1}')".format(start_str, end_str)

    n = 0
    inc = datetime_start
    while inc < end:
        n = n + 1
        inc = datetime_math.increment_datetime_by_string(inc, inputinfo['granularity'], mult=1)

    inmaps = grass.read_command('t.rast.list',
                              input=command_dict['input'],
                              where=where_str, flags='u',
                              columns='id,start_time,end_time').rstrip('\n').split('\n')

    if len(inmaps) / float(completeness) >= n / 100.0:
        return where_str
    
    
agg_maps = grass.read_command('t.rast.list', input=strds, flags='u', columns='id,name,creation_time,start_time,end_time').rstrip('\n').split('\n')
                          
                          
                          
                          
                          
                          
                          
                          
inmaps = grass.read_command('t.rast.list',
                          input=command_dict['input'],
                          where=where_str, flags='u',
                          columns='id,start_time,end_time').rstrip('\n').split('\n')


for i in range(len(inmaps)):
    if i + 1 < len(inmaps):
        cur = datetime_math.string_to_datetime(inmaps[i].split('|')[2])
        nd = datetime_math.string_to_datetime(inmaps[i + 1].split('|')[2])
        if nd - cur > timedelta(1):
            print(cur)
            print(inmaps[i])
            print(inmaps[i + 1])

# run module with modified paramters
mod = mod_iface.module.Module(name, run_=False, **update_granule)

# Removes temporary strds

# (Re-) Register 

time_string = added_maps.split('\n')[1].split('|')[3]

dt = datetime_math.string_to_datetime(time_string)
start = datetime_math.adjust_datetime_to_granularity(dt, strdsinfo['granularity'])

end = datetime_math.increment_datetime_by_string(start, strdsinfo['granularity'], mult=1)


datetime_math.datetime_to_grass_datetime_string(dt)

tgis.init()

t.rast.aggregate input=snow_bin_seNorge_1km_days basename=snow_days_seNorge_1km_years suffix=gran granularity=1 years method=sum offset=0 nprocs=10 file_limit=1000 sampling=contains where="end_time >= '2017-01-01T00:00:00.1483225200'  AND start_time <= '2018-01-01T00:00:00.1514761200'" output=suJnMuk0VEXUOF5BzZNnE --o --q

"""