#! /usr/bin/env python
"""
Time conversion utilities
"""

import os
from datetime import datetime, timedelta
import time

import numpy as np
import matplotlib.dates

#Seconds per year
spy = 86400.*365.25

#lon,lat = geolib.get_center(ds, t_srs=geolib.wgs_srs)
def getTimeZone(lat, lon):
    """Get timezone for a given lat/lon
    """
    #Need to fix for Python 2.x and 3.X support
    import urllib.request, urllib.error, urllib.parse
    import xml.etree.ElementTree as ET
    #http://api.askgeo.com/v1/918/aa8292ec06199d1207ccc15be3180213c984832707f0cbf3d3859db279b4b324/query.xml?points=37.78%2C-122.42%3B40.71%2C-74.01&databases=Point%2CTimeZone%2CAstronomy%2CNaturalEarthCountry%2CUsState2010%2CUsCounty2010%2CUsCountySubdivision2010%2CUsTract2010%2CUsBlockGroup2010%2CUsPlace2010%2CUsZcta2010
    req = "http://api.askgeo.com/v1/918/aa8292ec06199d1207ccc15be3180213c984832707f0cbf3d3859db279b4b324/query.xml?points="+str(lat)+"%2C"+str(lon)+"&databases=TimeZone"
    opener = urllib.request.build_opener()
    f = opener.open(req)
    tree = ET.parse(f)
    root = tree.getroot()
    #Check response
    tzid = None
    if root.attrib['code'] == '0':
        tz = list(root.iter('TimeZone'))[0]
        #shortname = tz.attrib['ShortName']
        tzid = tz.attrib['TimeZoneId']
    return tzid 

def getLocalTime(utc_dt, tz):
    """Return local timezone time
    """
    import pytz
    local_tz = pytz.timezone(tz)
    local_dt = utc_dt.replace(tzinfo=pytz.utc).astimezone(local_tz)
    return local_dt

def ul_time(utc_dt, lon):
    """Compute local time for input longitude
    """
    #return utc_dt + timedelta(hours=lon / np.pi * 12)
    offset = timedelta(hours=(lon*(24.0/360)))
    return utc_dt + offset

def solarTime(utc_dt, lat, lon):
    """Compute local solar time for given (lat, lon)
    """
    import ephem
    o = ephem.Observer()
    o.date = utc_dt
    o.lat = str(lat)
    o.lon = str(lon)
    sun = ephem.Sun()
    sun.compute(o)
    hour_angle = o.sidereal_time() - sun.ra
    rad = str(ephem.hours(hour_angle + ephem.hours('12:00')).norm)
    t = datetime.strptime(rad, '%H:%M:%S.%f')
    solar_dt = datetime.combine(utc_dt.date(), t.time()) 
    return solar_dt 

def strptime_fuzzy(s):
    """Fuzzy date string parsing

    Note: this returns current date if not found. If only year is provided, will return current month, day
    """
    import dateutil.parser
    dt = dateutil.parser.parse(str(s), fuzzy=True) 
    return dt 

def fn_getdatetime(fn):
    """Extract datetime from input filename
    """
    dt_list = fn_getdatetime_list(fn)
    if dt_list:
        return dt_list[0]
    else:
        return None

#Return datetime object extracted from arbitrary filename
def fn_getdatetime_list(fn):
    """Extract all datetime strings from input filename
    """
    #Want to split last component
    fn = os.path.split(os.path.splitext(fn)[0])[-1]
    import re
    #WV01_12JUN152223255-P1BS_R1C1-102001001B3B9800__WV01_12JUN152224050-P1BS_R1C1-102001001C555C00-DEM_4x.tif
    #Need to parse above with month name 
    #Note: made this more restrictive to avoid false matches:
    #'20130304_1510_1030010020770600_1030010020CEAB00-DEM_4x'
    #This is a problem, b/c 2015/17/00:
    #WV02_20130315_10300100207D5600_1030010020151700
    #This code should be obsolete before 2019 
    #Assume new filenames
    #fn = fn[0:13]
    #Use cascading re find to pull out timestamps
    #Note: Want to be less restrictive here - could have a mix of YYYYMMDD_HHMM, YYYYMMDD and YYYY in filename
    #Should probably search for all possibilities, then prune  
    #NOTE: these don't include seconds in the time
    #NOTE: could have 20130304_1510__20130304__whatever in filename
    #The current approach will only catch the first datetime 
    dstr = None
    out = None
    #20180101_1200 or 20180101T1200
    dstr = re.findall(r'(?:^|_|-)(?:19|20)[0-9][0-9](?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01])[_T](?:0[0-9]|1[0-9]|2[0-3])[0-5][0-9]', fn)
    #201801011200
    if not dstr:
        dstr = re.findall(r'(?:^|_|-)(?:19|20)[0-9][0-9](?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01])(?:0[0-9]|1[0-9]|2[0-3])[0-5][0-9]', fn)
    #20180101
    if not dstr:
        dstr = re.findall(r'(?:^|_|-)(?:19|20)[0-9][0-9](?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01])(?:$|_|-)', fn)
        #This should pick up dates separated by a dash
        #dstr = re.findall(r'(?:^|_|-)(?:19|20)[0-9][0-9](?:0[1-9]|1[012])(?:0[1-9]|[12][0-9]|3[01])', fn)
    #2018.609990
    if not dstr:
        dstr = re.findall(r'(?:^|_|-)(?:19|20)[0-9][0-9]\.[0-9][0-9][0-9]*(?:$|_|-)', fn)
        dstr = [d.lstrip('_').rstrip('_') for d in dstr]
        dstr = [d.lstrip('-').rstrip('-') for d in dstr]
        out = [decyear2dt(float(s)) for s in dstr]
        dstr = None
    #2018
    if not dstr:
        dstr = re.findall(r'(?:^|_|-)(?:19|20)[0-9][0-9](?:$|_|-)', fn)
    #This is for USGS archive filenames
    if not dstr:
        dstr = re.findall(r'[0-3][0-9][a-z][a-z][a-z][0-9][0-9]', fn)
        #This is USGS archive format
        if dstr:
            out = [datetime.strptime(s, '%d%b%y') for s in dstr]
            dstr = None
    if dstr:
        #This is a hack to remove peripheral underscores and dashes
        dstr = [d.lstrip('_').rstrip('_') for d in dstr]
        dstr = [d.lstrip('-').rstrip('-') for d in dstr]
        #This returns an empty list of nothing is found
        out = [strptime_fuzzy(s) for s in dstr]
    return out

def get_t_factor(t1, t2):
    """Time difference between two datetimes, expressed as decimal year 
    """
    t_factor = None
    if t1 is not None and t2 is not None and t1 != t2:  
        dt = t2 - t1
        year = timedelta(days=365.25)
        t_factor = abs(dt.total_seconds() / year.total_seconds()) 
    return t_factor

def get_t_factor_fn(fn1, fn2, ds=None):
    t_factor = None
    #Extract timestamps from input filenames
    t1 = fn_getdatetime(fn1)
    t2 = fn_getdatetime(fn2)
    t_factor = get_t_factor(t1,t2)
    #Attempt to load timestamp arrays (for mosaics with variable timestamps)
    t1_fn = os.path.splitext(fn1)[0]+'_ts.tif'
    t2_fn = os.path.splitext(fn2)[0]+'_ts.tif'
    if os.path.exists(t1_fn) and os.path.exists(t2_fn) and ds is not None:
        print("Preparing timestamp arrays")
        from pygeotools.lib import warplib
        t1_ds, t2_ds = warplib.memwarp_multi_fn([t1_fn, t2_fn], extent=ds, res=ds)
        print("Loading timestamps into masked arrays")
        from pygeotools.lib import iolib
        t1 = iolib.ds_getma(t1_ds)
        t2 = iolib.ds_getma(t2_ds)
        #This is a new masked array
        t_factor = (t2 - t1) / 365.25
    return t_factor

def sort_fn_list(fn_list):
    """Sort input filename list by datetime
    """
    dt_list = get_dt_list(fn_list)
    fn_list_sort = [fn for (dt,fn) in sorted(zip(dt_list,fn_list))]
    return fn_list_sort

def fix_repeat_dt(dt_list, offset_s=0.001):
    """Add some small offset to remove duplicate times
    Needed for xarray interp, which expects monotonically increasing times
    """
    idx = (np.diff(dt_list) == timedelta(0))
    while np.any(idx):
        dt_list[idx.nonzero()[0] + 1] += timedelta(seconds=offset_s)
        idx = (np.diff(dt_list) == timedelta(0))
    return dt_list

def get_dt_list(fn_list):
    """Get list of datetime objects, extracted from a filename
    """
    dt_list = np.array([fn_getdatetime(fn) for fn in fn_list])
    return dt_list

#Pad must be timedelta
#pad=timedelta(days=30)
def filter_fn_list(dt, fn_list, pad):
    dt_list = get_dt_list(fn_list)
    #These should be sorted by time
    #This pulls fixed number on either side of dt
    #idx = timelib.get_closest_dt_idx(cdt, v_dt_list)
    #idx = 'idx-2:idx+2' 
    #This pulls from a fixed time interval on either side of dt
    idx = get_closest_dt_padded_idx(dt, dt_list, pad)
    fn_list_sel = fn_list[idx]
    return fn_list_sel

def get_closest_dt_fn(fn, fn_list):
    dt = fn_getdatetime(fn)
    dt_list = np.array([fn_getdatetime(fn) for fn in fn_list])
    idx = get_closest_dt_idx(dt, dt_list) 
    return fn_list[idx]

def get_closest_dt_idx(dt, dt_list):
    """Get indices of dt_list that is closest to input dt
    """
    from pygeotools.lib import malib
    dt_list = malib.checkma(dt_list, fix=False)
    dt_diff = np.abs(dt - dt_list)
    return dt_diff.argmin()

def get_closest_dt_padded_idx(dt, dt_list, pad=timedelta(days=30)):
    """Get indices of dt_list that is closest to input dt +/- pad days
    """
    #If pad is in decimal days
    if not isinstance(pad, timedelta):
        pad = timedelta(days=pad)
    from pygeotools.lib import malib
    dt_list = malib.checkma(dt_list, fix=False)
    dt_diff = np.abs(dt - dt_list)
    valid_idx = (dt_diff.data < pad).nonzero()[0]
    return valid_idx

def get_unique_monthyear(dt_list):
    my = [(dt.year,dt.month) for dt in dt_list]
    return np.unique(my)

def get_dt_bounds_monthyear(dt_list):
    my_list = get_unique_monthyear(dt_list)
    out = []
    for my in my_list:
        dt1 = datetime(my[0], my[1], 1)
        dt2 = datetime(my[0], my[1], 1)
    out.append((dt1, dt2))
    return out

def get_unique_years(dt_list):
    years = [dt.year for dt in dt_list]
    return np.unique(years)

def dt_filter_rel_annual_idx(dt_list, min_rel_dt=(1,1), max_rel_dt=(12,31)):
    """Return dictionary containing indices of timestamps that fall within relative month/day bounds of each year
    """
    dt_list = np.array(dt_list)
    years = get_unique_years(dt_list)
    from collections import OrderedDict
    out = OrderedDict() 
    for year in years:
        #If within the same year
        if min_rel_dt[0] < max_rel_dt[1]:
            dt1 = datetime(year, min_rel_dt[0], min_rel_dt[1])
            dt2 = datetime(year, max_rel_dt[0], max_rel_dt[1])
        #Or if our relative values include Jan 1
        else:
            dt1 = datetime(year, min_rel_dt[0], min_rel_dt[1])
            dt2 = datetime(year+1, max_rel_dt[0], max_rel_dt[1])
        idx = np.logical_and((dt_list >= dt1), (dt_list <= dt2))
        if np.any(idx):
            out[year] = idx
    return out

#Use this to get datetime bounds for annual mosaics
def get_dt_bounds(dt_list, min_rel_dt=(1,1), max_rel_dt=(12,31)):
    years = get_unique_years(dt_list)
    out = []
    for year in years:
        #If within the same year
        if min_rel_dt[0] < max_rel_dt[1]:
            dt1 = datetime(year, min_rel_dt[0], min_rel_dt[1])
            dt2 = datetime(year, max_rel_dt[0], max_rel_dt[1])
        else:
            dt1 = datetime(year, min_rel_dt[0], min_rel_dt[1])
            dt2 = datetime(year+1, max_rel_dt[0], max_rel_dt[1])
        if dt2 > dt_list[0] and dt1 < dt_list[-1]:
            out.append((dt1, dt2))
    return out

def get_dt_bounds_fn(list_fn, min_rel_dt=(5,31), max_rel_dt=(6,1)):
    f = open(list_fn, 'r')
    fn_list = []
    for line in f:
        fn_list.append(line)
    fn_list = np.array(fn_list)
    fn_list.sort()
    dt_list = [fn_getdatetime(fn) for fn in fn_list]
    dt_list = np.array(dt_list)
    bounds = get_dt_bounds(dt_list, min_rel_dt, max_rel_dt)
    for b in bounds:
        print(b)
        c_date = center_date(b[0], b[1])
        #c_date = datetime(b[1].year,1,1)
        idx = (dt_list >= b[0]) & (dt_list < b[1])
        #out_fn = os.path.splitext(list_fn)[0]+'_%s_%s-%s_fn_list.txt' % \
        #        (c_date.strftime('%Y%m%d'), b[0].strftime('%Y%m%d'), b[1].strftime('%Y%m%d'))
        out_fn = '%s_%s-%s_fn_list.txt' % \
                (c_date.strftime('%Y%m%d'), b[0].strftime('%Y%m%d'), b[1].strftime('%Y%m%d'))
        out_f = open(out_fn, 'w')
        for fn in fn_list[idx]:
            out_f.write('%s\n' % fn)
        out_f = None

#parallel 'dem_mosaic -l {} --count -o {.}' ::: 2*fn_list.txt

#This checks to see if input dt is between the given relative month/day interval 
def rel_dt_test(dt, min_rel_dt=(1,1), max_rel_dt=(12,31)):
    if dt_check(dt): 
        dt_doy = dt2doy(dt)
        min_doy = dt2doy(datetime(dt.year, *min_rel_dt))
        max_doy = dt2doy(datetime(dt.year, *max_rel_dt))
        #If both relative dates are in the same year
        if min_doy < max_doy:
            min_dt = datetime(dt.year, min_rel_dt[0], min_rel_dt[1])
            max_dt = datetime(dt.year, max_rel_dt[0], max_rel_dt[1])
        else:
            #If relative dates span Jan 1
            if dt_doy >= min_doy:
                min_dt = datetime(dt.year, min_rel_dt[0], min_rel_dt[1])
                max_dt = datetime(dt.year + 1, max_rel_dt[0], max_rel_dt[1])
            else:
                min_dt = datetime(dt.year - 1, min_rel_dt[0], min_rel_dt[1])
                max_dt = datetime(dt.year, max_rel_dt[0], max_rel_dt[1])
        out = (dt >= min_dt) & (dt <= max_dt)
    else:
        out = False