from hec.data.tx             import QualityTx
from hec.heclib.dss          import HecDss
from hec.heclib.util         import HecTime
from hec.heclib.util         import Heclib
from hec.hecmath             import TimeSeriesMath
from hec.io                  import TimeSeriesContainer
from hec.script              import Constants
from hec.script              import MessageBox
from java.lang               import System
from java.text               import SimpleDateFormat
from java.util               import TimeZone
from javax.swing             import JFileChooser
from javax.swing             import JOptionPane
from javax.swing.filechooser import FileNameExtensionFilter
from urllib2                 import urlopen
import os, re, string, sys, threading
#--------------------------------------------#
# conditional imports and global definitions #
#--------------------------------------------#
try :
	from usace.cavi.client import CAVI as cCAVI
	from usace.cavi.script import CAVI as sCAVI
	has_cavi      = True
	cavi_frame    = cCAVI.getBrowserFrame()
	proj_dir      = sCAVI.getWatershed().getProjectDirectory()
	cwms_home     = os.path.normpath(os.path.join(proj_dir, "..", ".."))
	dss_filename  = os.path.join(cwms_home, "database", "Russian_TimeSeries.dss")
	cdec_filename = os.path.join(cwms_home, "database", "Russian.cdec")
	script_dir    = os.path.join(proj_dir, "scripts")
	if not script_dir in sys.path :
		sys.path.append(script_dir)
	try :
		from MessageOutput import Outputter
		outputter = Outputter()
		has_outputter = True
	except ImportError :
		has_outputter = False
except ImportError :
	has_cavi      = False
	has_outputter = False
	dss_filename  = None
	cdec_filename = None
#---------------------------#
# static global definitions #
#---------------------------#
script_name     = "2-CDEC"
log_filename    = None
sensors         = [] # modified in main()
sensor_names    = [] # modified in main()
sensors_by_name = [] # modified in main()
#-------------------------#
# setup output() function #
#-------------------------#
if has_outputter :
	#--------------------------------------------------------------------------#
	# Use the following properties to control the default behavior of output() #
	#--------------------------------------------------------------------------#
	outputter.messenger     = cavi_frame
	outputter.level         = "INFO"
	outputter.beep          = False
	outputter.acknowledge   = False
	outputter.title         = script_name
	outputter.logmode       = "w"
	outputter.make_log_dir  = False
	outputter.logfile       = log_filename
	output = outputter.output
else :
	#-------------------------------------------------#
	# Create variable log_filename to enable logging. #
	#                                                 #
	# If logging is enabled, use the output_to_log    #
	# variable to turn logging on/off at as desired.  #
	#                                                 #
	# Default behavior of output() is :               #
	#   level       = "NONE"                          #
	#   beep        = False                           #
	#   acknowledge = False                           #
	#   title       = None (defaults to script_name)  #
	#   log         = output_to_log                   #
	#-------------------------------------------------#
	logfile = None
	output_to_log = False
	if log_filename : logfile = open(log_filename, "w")
	def _output(text, level=None, beep=False, acknowledge=False, title=None, log=None) :
		"""
		default output method in lieu of Outputter.output()
		"""
		global logfile, output_to_log
		if beep : print("\a")
		if   not level  or level.upper() == "NONE"    : prefix = ""
		elif level == 1 or level.upper() == "INFO"    : prefix = "INFO: "
		elif level == 2 or level.upper() == "WARNING" : prefix = "WARNING: "
		elif level == 3 or level.upper() == "ERROR"   : prefix = "ERROR: "
		else :
			raise ValueError("Invalid level: %s" % level)
		if prefix :
			for i in range(len(text)) :
				if not text[i].isspace() :
					to_output = "{0}{1}{2}".format(text[:i], prefix, text[i:])
					break
			else :
				to_output = "{0}{1}".format(prefix, text)
		else :
			to_output = text
		print(to_output)
		if logfile and output_to_log if log is None else bool(log) :
			timestr = str(datetime.now())[:-7]
			logfile.write("{0} {1}\n".format(timestr, to_output))
		if acknowledge :
			if not title : title = script_name
			if not level                                  : func = MessageBox.showInformation
			elif level == 0 or level.upper() == "NONE"    : func = MessageBox.showInformation
			elif level == 1 or level.upper() == "INFO"    : func = MessageBox.showInformation
			elif level == 2 or level.upper() == "WARNING" : func = MessageBox.showWarning
			elif level == 3 or level.upper() == "ERROR"   : func = MessageBox.showError
			func(text, title)
	output = _output
#---------------------------#
# function get_cdec_sensors #
#---------------------------#
def get_cdec_sensors() :
	"""
	Retrieves the list of sensor numbers with associated names, descriptions, and units from CDEC
	"""
	sensor_data = []
	sensor_web_page = urlopen("https://cdec.water.ca.gov/misc/senslist.html").read()
	sensor_table = re.search(r"<table.*?>.+?</table>", sensor_web_page, flags=re.M | re.S).group(0)
	match = re.search(r"<tr.*?>\s*<th.+?</tr>", sensor_table, flags=re.M | re.S)
	sensor_table_header_row = match.group(0)
	sensor_table_header_row_end = match.end(0)
	sensor_table_header_cols = [
		s.split("<")[0].strip() \
		for s in re.split(r"<th.*?>(?:\s*<.+?>)*", sensor_table_header_row, flags=re.M | re.S) \
		if not s.startswith("<")
	]
	assert "Sensor No"   in sensor_table_header_cols
	assert "Sensor"      in sensor_table_header_cols
	assert "Description" in sensor_table_header_cols
	assert "Units"       in sensor_table_header_cols
	column_index = dict(zip(sensor_table_header_cols, range(len(sensor_table_header_cols))))
	sensor_table_data_rows = re.split(r"<tr.*?>", sensor_table[sensor_table_header_row_end:], flags=re.M | re.S)
	for i in range(len(sensor_table_data_rows)) :
		sensor_table_row_cols = [
			s.split("<")[0].strip() \
			for s in re.split(r"<td.*?>(?:\s*<.+?>)*", sensor_table_data_rows[i], flags=re.M | re.S)
		][1:]
		if len(sensor_table_row_cols) != 5 : continue
		sensor_data.append([
			sensor_table_row_cols[j]
			for j in [
				column_index[name] \
				for name in ["Sensor No", "Sensor", "Description", "Units"]
			]
		])
	for i in range(len(sensor_data)) :
		sensor_data[i][column_index["Sensor No"]] = int(sensor_data[i][column_index["Sensor No"]])
	sensors = [s[column_index["Sensor No"]] for s in sensor_data] # sensor numbers
	sensor_names = dict([[s[column_index["Sensor No"]], s[column_index["Sensor"]]] for s in sensor_data])
	sensors_by_name = dict([[s[column_index["Sensor"]], s[column_index["Sensor No"]]] for s in sensor_data] + \
		               [[s[column_index["Description"]], s[column_index["Sensor No"]]] for s in sensor_data])
	return sensors, sensor_names, sensors_by_name
#--------------------------#
# function validate_sensor #
#--------------------------#
def validate_sensor(sensor) :
	"""
	Validate a sensor specified by number, name, or description, returning the number
	"""
	global sensors, sensor_names, sensors_by_name
	if isinstance(sensor, int) :
		if sensor not in sensors :
			raise ValueError("No such sensor: %s" % sensor)
		return sensor
	elif isinstance(sensor, basestring) :
		try :
			return sensors_by_name[sensor.upper()]
		except KeyError :
			raise ValueError('No such sensor: "%s"' % sensor)
	else :
		raise TypeError("Expected integer or string, got %s" % sensor.__class__.__name__)
#------------------------#
# function get_cdec_data #
#------------------------#
def get_cdec_data(station, end_datetime, num_hours, sensors_to_get=None) :
	"""
	Retrieve data from CDEC for specified station, time window, and sensor (parameter)
	"""
	url_template = "http://cdec.water.ca.gov/dynamicapp/QueryF?s={0}&d={1}&span={2}hours"
	# example       http://cdec.water.ca.gov/dynamicapp/QueryF?s=ORO&d=22-Apr-2021+08:30&span=72hours, time in Pacific Local
	columns_by_parameter = {}
	parameters_to_get = None
	columns_to_get = None
	parameters = []
	units = []
	times = []
	values = []
	#--------------------------------------#
	# process the list sensors to retrieve #
	#--------------------------------------#
	if sensors_to_get :
		if isinstance(sensors_to_get, (list, tuple)) :
			for i in range(len(sensors_to_get)) :
				sensors_to_get[i] = validate_sensor(sensors_to_get[i])
		else :
			sensors_to_get = [validate_sensor(sensors_to_get)]
		parameters_to_get = [sensor_names[s] for s in sensors_to_get]
	#----------------------------#
	# retrieve data from the url #
	#----------------------------#
	url  = url_template.format(station, end_datetime, num_hours)
	data = urlopen(url).read()
	#------------------------#
	# extract the data table #
	#------------------------#
	match = re.search(r"<table.+?>\s*<thead>.+?</thead>\s*<tbody>.+?</tbody>\s*</table>", data, flags=re.M | re.S)
	if not match :
		output("\tNo data found for station {0}".format(station), level="WARNING")
	else :
		data_table = match.group(0)
		#----------------------------------------------------#
		# get the parameters and units from the table header #
		#----------------------------------------------------#
		header_row = re.search(r"<th.*?>.+</th>\s*</tr>", data_table, flags=re.M | re.S).group(0)
		columns = re.split(r"(<th.*?>.+?</th>)", header_row, flags=re.M | re.S)
		for i in range(len(columns)) :
			if i % 2 == 1 :
				parts = re.split(r"<br>", columns[i], flags=re.M | re.S)
				parameter = re.search(r"<th.*?>(?:<.+?>)*(.+?)(?:<|$)", parts[0], flags=re.M | re.S).group(1)
				if not parameters_to_get or parameter in parameters_to_get :
					parameters.append(parameter)
					units.append(re.search(r"(?:<.+?>)*\s*(.+?)\s*(?:<|$)", parts[1], flags=re.M | re.S).group(1))
					columns_by_parameter[parameter] = i / 2
		units_by_parameter = dict(zip(parameters, units))
		#------------------------------------------------------------#
		# organize the returned columns by specified parameter order #
		#------------------------------------------------------------#
		if parameters_to_get :
			parameters = []
			units = []
			for parameter in parameters_to_get :
				if parameter in columns_by_parameter :
					parameters.append(parameter)
					units.append(units_by_parameter[parameter])
				else :
					output("\tStation {0} does not have parameter {1}".format(station, parameter), level="WARNING")
			parameters_by_column = dict(zip(columns_by_parameter.values(), columns_by_parameter.keys()))
			columns_to_get = [0] + [columns_by_parameter[p] for p in parameters]
		else :
			columns_to_get = sorted(columns_by_parameter.values())
		val_cols = columns_to_get[1:] # ignore time column
		column_count = len(val_cols)
		column_range = range(column_count)
		column_index = dict(zip([val_cols[j] for j in column_range], column_range))
		for i in range(column_count) : values.append([])
		if parameters :
			#--------------------------------------------------------#
			# get the times and parameter values from the table rows #
			#--------------------------------------------------------#
			table_body = re.search(r"<tbody>\s*(<tr>.+</tr>)\s*</tbody>", data_table, flags=re.M | re.S).group(1)
			body_rows = re.split(r"\s*</tr>\s*", table_body, flags=re.M | re.S)
			for i in range(len(body_rows)) :
				columns = [c for c in re.split(r"\s*</td>\s*", body_rows[i], flags=re.M | re.S) if c.startswith("<t")]
				if columns :
					for j in columns_to_get :
						value = re.search(r"<td.*?>(?:<.+?>)*?([^<>]+)(?:<|$)", columns[j], flags=re.M | re.S).group(1)
						value = value.replace("&nbsp;", "").strip()
						if j == 0 :
							times.append(value)
						else :
							try    : value = float(value)
							except : value = Constants.UNDEFINED
							values[column_index[j]].append(value)
	return parameters, units, times, values
#-----------------------------------#
# function make_TimeSeriesContainer #
#-----------------------------------#
def make_TimeSeriesContainer(times, values, unit, watershed, location, parameter, version) :
	"""
	Creates a 1HOUR TimeSeriesContainer from data returned from CDEC
	"""
	def hecTimeValue(timestr, hecTime=None) :
		if not hecTime : hecTime = HecTime()
		hecTime.set(timestr)
		return hecTime.value()
	if len(times) < 2 : return None
	#-------------------------------#
	# build the TimeSeriesContainer #
	#-------------------------------#
	tsc = TimeSeriesContainer()
	ht = HecTime()
	tsc.times = [hecTimeValue(ts, ht) for ts in times]
	tsc.values = values
	tsc.numberValues = len(values)
	tsc.watershed = watershed
	tsc.location = location
	tsc.parameter = parameter
	tsc.version = version
	tsc.units = unit
	if parameter.upper().find("PRECIP") != -1 :
		if parameter.upper().find("CUM") != -1 :
			tsc.type = "INST-CUM"
		elif parameter.upper().find("INC") != -1 :
			tsc.type = "PER-CUM"
		else :
			output("Don't know whether precip is cumulative or incremental", level="ERROR")
			return None
	else :
		tsc.type = "INST-VAL"
	tsc.interval = tsc.times[1] - tsc.times[0]
	if tsc.interval in (15, 30) :
		ePart = "{0}Minute".format(tsc.interval)
	else :
		ePart = "{0}Hour".format(tsc.interval / 60)
	tsc.fullName = "/{0}/{1}/{2}//{3}/{4}/".format(watershed, location, parameter, ePart, version)
	#------------------------------------------------#
	# transform the TimeSeriesContainer if necessary #
	#------------------------------------------------#
	if tsc.interval != 60 :
		tsc = TimeSeriesMath(tsc).transformTimeSeries("1HOUR", "0MIN", "INT").getData()
	if tsc.type == "INST-CUM" :
		tsc = TimeSeriesMath(tsc).successiveDifferences().screenWithMaxMin(0.0, 4.0, 4.0).getData()
		for i in range(tsc.numberValues) :
			if tsc.quality :
				if QualityTx.isReject_int(tsc.quality[i]) \
				or QualityTx.isMissing_int(tsc.quality[i]) \
				or tsc.values[i] == Constants.UNDEFINED :
					tsc.values[i] = 0
					tsc.quality[i] = QualityTx.setRevisedAutomatically_int(tsc.quality[i])
			else :
				if tsc.values[i] == Constants.UNDEFINED :
					tsc.values[i] = 0
	return tsc
#---------------------------#
# function get_dss_timezone #
#---------------------------#
def get_dss_timezone() :
	"""
	gets the working time zone for the watershed
	"""
	tz = None
	if has_cavi :
		projectFilename = sCAVI.getWatershed().getProjectFile().getCanonicalPath()
		with open(projectFilename) as f : data = f.read()
		tzlines = [line for line in data.strip().split("\n") if line.startswith("TimeZone=")]
		if len(tzlines) != 1 :
			raise Exception("No time zone specified in watershed file {0}".format(projectFilename))
		tz = TimeZone.getTimeZone(tzlines[0].split("=")[1])
	else :
		# set global variable "dssTz" to desired time zone to prevent the following interactive action
		try :
			dssTz
		except :
			time_zone_ids = TimeZone.getAvailableIDs()
			user_input = JOptionPane.showInputDialog(
				None,
				"Enter time zone for DSS file",
				script_name,
				JOptionPane.INFORMATION_MESSAGE,
				None,
				time_zone_ids,
				"US/Central")
			try :
				tz = TimeZone.getTimeZone(user_input)
				if tz.getID() != user_input :
					tz = None
					raise ValueError()
			except :
				output('"{0}" is not a valid time zone'.format(user_input), level="ERROR")
		else :
			if dssTz :
				try :
					tz = TimeZone.getTimeZone(dssTz)
				except :
					output('"{0}" is not a valid time zone'.format(dssTz), level="ERROR")
	return tz
#----------------------#
# function choose_file #
#----------------------#
def choose_file(frame, title, starting_dir, description=None, extensions=None, allow_multi_select=False) :
	"""
	Interactively choose file from file system
	"""
	jfc = JFileChooser(starting_dir)
	jfc.setDialogTitle(title)
	jfc.setMultiSelectionEnabled(allow_multi_select)
	if extensions : jfc.setFileFilter(FileNameExtensionFilter(description, extensions))
	if jfc.showOpenDialog(frame) == JFileChooser.APPROVE_OPTION :
		if allow_multi_select :
			return [f.getCanonicalPath() for f in jfc.getSelectedFiles()]
		else :
			return jfc.getSelectedFile().getCanonicalPath()
	return (None, [])[allow_multi_select]
#---------------#
# function main #
#---------------#
def main() :
	"""
	main function
	"""
	global dss_filename, cdec_filename, sensors, sensor_names, sensors_by_name
	max_lookback_days  = 60
	lookback_days      = None
	locations          = {}
	tzCdec             = TimeZone.getTimeZone("US/Pacific")
	tzDss              = get_dss_timezone()
	sdfCdec            = SimpleDateFormat("dd-MMM-yyyy+HH:mm")
	sdfHecCdec         = SimpleDateFormat("ddMMMyyyy, HH:mm")
	sdfHecDss          = SimpleDateFormat("ddMMMyyyy, HH:mm")
	t1                 = HecTime()
	t2                 = HecTime()
	sdfCdec.setTimeZone(tzCdec)
	sdfHecCdec.setTimeZone(tzCdec)
	if not tzDss :
		tzDss = TimeZone.getDefault()
		message = "No valid time zone has been entered for the DSS time zone.\n" + \
		          "Do you want to continue with the default of {0}?".format(tzDss.getID())
		if "no" == MessageBox.showYesNo(message, script_name) :
			return
	sdfHecDss.setTimeZone(tzDss)
	end_date = sdfCdec.format(System.currentTimeMillis())
	if has_cavi :
		cavi = cavi_frame
		starting_dir = cwms_home
	else :
		cavi = None
		starting_dir = "."
	#------------------------------------------------------------#
	# interactively get the DSS filenme if not already specified #
	#------------------------------------------------------------#
	if not dss_filename :
		dss_filename = choose_file(cavi, "Choose HEC-DSS File", starting_dir, "HEC-DSS Files", ["dss"])
		if not dss_filename :
			output("Nothing to do - no DSS file selected", level="WARNING")
			return
		starting_dir = os.path.dirname(dss_filename)
	#------------------------------------------------------------------#
	# interactively get the CDEC data filenme if not already specified #
	#------------------------------------------------------------------#
	if not cdec_filename :
		cdec_filename = choose_file(cavi, "Choose CDEC Input File", starting_dir, "CDEC Input Files", ["cdec"])
		if not cdec_filename :
			output("Nothing to do - no input file selected", level="WARNING")
			return
	#---------------------------------------------------------------------------#
	# interactively get the number of days to retrieve if not already specified #
	#---------------------------------------------------------------------------#
	if not lookback_days :
		user_input = JOptionPane.showInputDialog(
			cavi,
			"Enter number of days back to retrieve",
			script_name,
			JOptionPane.QUESTION_MESSAGE)
		if user_input :
			try :
				lookback_days = int(user_input)
			except :
				output('Invalid entry for number of days to retrieve: "{0}"'.format(user_input), level="ERROR")
				return
			if not 0 <= lookback_days <= max_lookback_days :
				output("Number of lookback days must be between 0 and {0}".format(max_lookback_days), level="ERROR")
				return
			if lookback_days == 0 :
				output("WARNING: Nothing to do - zero lookback days specified", level="ERROR")
				return
		else :
			output("Nothing to do - Number of lookback days not specified", level="WARNING")
			return
	#-----------------------------------------------------------------#
	# read the CDEC data file and parse into the locations dictionary #
	#-----------------------------------------------------------------#
	output("Reading CDEC file {0}".format(cdec_filename))
	with open(cdec_filename) as f :
		cdec_station_data = f.read()
	for line in cdec_station_data.strip().split("\n") :
		line = line.split("#")[0].strip()
		if not line : continue
		parts = map(string.strip, line.split("="))
		if len(parts) != 2 : continue
		key, value = parts[0].upper(), parts[1]
		if key == "STATION NAME" :
			location = value
			locations.setdefault(location, {
				"basin"     : "",
				"station "  : "",
				"parameter" : [],
				"version"   : "",
				"sensors"   : []})
		elif key == "BASIN NAME" :
			locations[location]["basin"] = value
		elif key == "STATION ID" :
			locations[location]["station"] = value
		elif key == "SENSOR NAME" :
			locations[location]["parameter"].append(value)
		elif key == "VERSION NAME" :
			locations[location]["version"] = value
		elif key == "SENSOR NUMBER" :
			locations[location]["sensors"].append(int(value))
		elif key == "DURATION CODE" :
			pass #meh
	#-------------------------------#
	# retrieve the CDEC sensor list #
	#-------------------------------#
	output("Retrieving sensor information from CDEC")
	sensors, sensor_names, sensors_by_name = get_cdec_sensors()
	#-----------------------------------------------------------#
	# retrieve the CDEC data for each location and store to DSS #
	#-----------------------------------------------------------#
	Heclib.zset("MLVL", "", 0)
	dss = HecDss.open(dss_filename)
	output("Retrieving data for {0} locations".format(len(locations)))
	i = 0
	for location in locations :
		i += 1
		output("{0}/{1} Processing location {2}".format(i, len(locations), location))
		#-------------------------------#
		# retrieve data for the station #
		#-------------------------------#
		parameters, units, times, values = get_cdec_data(
			locations[location]["station"],
			end_date,
			lookback_days*24,
			locations[location]["sensors"])
		#-------------------------------------------#
		# process data for each parameter retrieved #
		#-------------------------------------------#
		for j in range(len(parameters)) :
			value_count = len([v for v in values[j] if v != Constants.UNDEFINED])
			output("\t{0} valid values retrieved for {1}/{2}".format(
				value_count,
				locations[location]["station"],
				parameters[j]))
			if value_count > 1 :
				#------------------------------------------#
				# make a TimeSeriesContainer from the data #
				#------------------------------------------#
				tsc = make_TimeSeriesContainer(
					times,
					values[j],
					units[j],
					locations[location]["basin"],
					location,
					locations[location]["parameter"][j],
					locations[location]["version"])
				if tsc :
					#----------------------------------------------------------#
					# convert times from CDEC time zone to Watershed time zone #
					#----------------------------------------------------------#
					times = tsc.numberValues * [None]
					for k in range(tsc.numberValues) :
						t1.set(tsc.times[k])
						t2.set(sdfHecDss.format(sdfHecCdec.parse(t1.dateAndTime(4))))
						times[k] = t2.value()
					#---------------------------------#
					# (finally) store the data to DSS #
					#---------------------------------#
					dss.put(tsc)
					value_count = len([v for v in tsc.values if v != Constants.UNDEFINED])
					output("\t{0} valid values written to {1}:{2}".format(value_count, dss_filename, tsc.fullName))
			else :
				output("\tToo few values retrieved", level="WARNING")
	dss.close()
	output('Done', beep=True, acknowledge=True)
#
if __name__ == "__main__" :
	#--------------------------------------------------#
	# run the script in another thread to not block UI #
	#--------------------------------------------------#
	threading.Thread(target=main).start()
