#!/usr/bin/env python

import psycopg2 as pg
from sys import argv, exit
from os import environ, popen
from string import lower, upper, split, strip

program = 'db2wx'
author = 'Walter Brisken and Daniel Lyons'
version = '0.4'
verdate = '20151021'



# This program is to mine the database over a specified mjd time range and 
# extract all of the weather data.
# 
# The output is one file per antenna with rows of weather information sorted 
# in time order.  Every so often the weather station should populate the 
# database with a set of monitor points with the same (or nearly so) 
# timetag.  These should be collected into one record and written like:
# 
# # Ant MJD temp(C) pressure(mbar) dewpoint(C) windspeed(m/s) winddir(degrees) rain(cm) gust(m/s)
# BR 55162.4000000 11.8 978.4 6.1 0.5 329.0 0.00 1.4
# BR 55162.4059259 11.6 978.5 6.2 0.0 247.0 0.00 1.1
# BR 55162.4118519 11.5 978.5 6.0 0.0 201.0 0.00 0.7
# BR 55162.4177778 10.8 978.5 6.0 0.0 305.0 0.00 0.2
# BR 55162.4237037 10.8 978.6 5.9 0.0 260.0 0.00 0.5
# BR 55162.4296296 10.9 978.6 5.8 0.0 320.0 0.00 0.7
# BR 55162.4355440 10.5 978.7 6.0 0.5 214.0 0.00 0.7
# 
# 
# There are 9 columns
# 1. Antenna name (upper case, the prefix to -ws in the hostname field)
# 2. MJD of the record
# 3. "temp" mon point
# 4. "hdewpt" mon point
# 5. "barometer" mon point
# 6. "windspeed" mon point
# 7. "winddir" mon point
# 8. "rain" mon point
# 9. "windgust" mon point
# 
# Comment lines begin with #.  You can just copy the one shown above and put 
# that as the first line to keep things uniform.
# 
# For mon points that don't have data in the interval a value of -999 should 
# be substituted.  The "rain" mon point probably doesn't exist for the GBT 
# but will for the VLBA once it starts populating the database with weather 
# info.
# 
# That should be all.  I have a similar program that mines the database for 
# Mark5 module information: /home/swc/DiFX-2.2/src/difxdb/mark5c2db .  I 
# hope to use what I learn from you to improve that program.

# For VLA (from Bryan):
# temperature = emr-m352.HMT337.Temperature
# pressure = emr-m352.WXT520.Pressure
# dew point temperature = emr-m352.HMT337.Dewpoint_Temperature
# wind speed = emr-m352.WXT520.Wind_Speed_Average
# wind speed max (gust) = emr-m352.WXT520.Wind_Speed_Maximum
# wind direction = emr-m352.WXT520.Wind_Direction_Average
# rain accumulation = emr-m352.WXT520.Rain_Accumulation
# the rain accumulation isn't done in the way you want - you'll have to
# interpret that.

# fetch the weather data for the observation

def fetchVLBAWeather(start, stop):
	
	dbAccess = environ['VLBAMPTS_DB']

	arguments = { 'from': '%14.8f' % (start - 0.1), 'to': '%14.8f' % (stop + 0.1) }

	db = pg.connect(dbAccess).cursor()
	db.execute("""
	SELECT
	  split_part(hostname_timestamp, ',', 1) as hostname,
	  split_part(hostname_timestamp, ',', 2)::double precision as timestamp,
	  COALESCE(Temperature, -999) AS Temperature,
	  COALESCE(DewPointTemperature, -999) AS DewPointTemperature,
	  COALESCE(Pressure, -999) AS Pressure,
	  COALESCE(WindSpeed, -999) AS WindSpeed,
	  COALESCE(WindDirection, -999) AS WindDirection,
	  COALESCE(Rain, 0.0) AS Rain,
	  COALESCE(WindGust, -999) AS WindGust
	FROM crosstab(
	  $$ SELECT
	       UPPER(SUBSTR(hostname,1,2)) || ',' || timestamp,
	       monpointname,
	       monpointvalue
	     FROM mcdata
	     WHERE devicename = 'WEATHER' AND
		   monpointname IN ('Temperature', 'DewPointTemperature', 'Pressure', 'WindSpeed',
				    'WindDirection', 'Rain', 'WindGust') AND
		   timestamp BETWEEN %(from)s AND %(to)s
	     ORDER BY 1,2 $$,
	  $$ VALUES ('Temperature'), ('DewPointTemperature'), ('Pressure'), ('WindSpeed'), ('WindDirection'),
		    ('Rain'), ('WindGust') $$
	  ) AS ct(hostname_timestamp text,
		  Temperature double precision,
		  DewPointTemperature double precision,
		  Pressure double precision,
		  WindSpeed double precision,
		  WindDirection double precision,
		  Rain double precision,
		  WindGust double precision)""", arguments)

	wxData = db.fetchall()

	return wxData

def getVLBAWeather(wxData, start, stop, station, filename):

	out = open(filename, 'w')
	print 'Collecting %s Weather data from vlbampts db (may take a couple minutes)' % upper(station)

	print 'Writing %s' % filename
	out.write('# VLBA Weather data produced by %s ver. %s\n' % (program, version) )
	out.write('# Ant MJD temp(C) pressure(mbar) dewpoint(C) WindSpeed(m/s) winddir(degrees) rain(cm) gust(m/s)\n')
	for host, time, Temperature, DewPointTemperature, Pressure, WindSpeed, WindDirection, Rain, WindGust in wxData:
		if host != station:
			continue
		if time < start or time > stop:
			continue
		out.write('%s %f %f %f %f %f %f %f %f\n' % (host, time, Temperature, Pressure, DewPointTemperature, WindSpeed, WindDirection, Rain, WindGust) )
	
	out.close()

def getGBWeather(start, stop, filename):

	arguments = { 'from': '%14.8f' % (start - 0.1), 'to': '%14.8f' % (stop + 0.1) }
	dbAccess = environ['VLBAMPTS_DB']

	out = open(filename, 'w')
	print 'Collecting GB Weather data from vlbampts db (may take a couple minutes)'

	db = pg.connect(dbAccess).cursor()
	db.execute("""
	SELECT
	  split_part(hostname_timestamp, ',', 1) as hostname,
	  split_part(hostname_timestamp, ',', 2)::double precision as timestamp,
	  COALESCE(temp, -999) AS temp,
	  COALESCE(hdewpt, -999) AS hdewpt,
	  COALESCE(barometer, -999) AS barometer,
	  COALESCE(windspeed, -999) AS windspeed,
	  COALESCE(winddir, -999) AS winddir,
	  COALESCE(rain, 0.0) AS rain,
	  COALESCE(windgust, -999) AS windgust
	FROM crosstab(
	  $$ SELECT
	       UPPER(SUBSTR(hostname,1,2)) || ',' || timestamp,
	       monpointname,
	       monpointvalue
	     FROM mcdata
	     WHERE hostname = 'gb-ws' AND
		   devicename = 'WS' AND
		   monpointname IN ('temp', 'hdewpt', 'barometer', 'windspeed',
				    'winddir', 'rain', 'windgust') AND
		   timestamp BETWEEN %(from)s AND %(to)s
	     ORDER BY 1,2 $$,
	  $$ VALUES ('temp'), ('hdewpt'), ('barometer'), ('windspeed'), ('winddir'),
		    ('rain'), ('windgust') $$
	  ) AS ct(hostname_timestamp text,
		  temp double precision,
		  hdewpt double precision,
		  barometer double precision,
		  windspeed double precision,
		  winddir double precision,
		  rain double precision,
		  windgust double precision)""", arguments)

	print 'Writing %s' % filename
	out.write('# GB Weather data produced by %s ver. %s\n' % (program, version) )
	out.write('# Ant MJD temp(C) pressure(mbar) dewpoint(C) windspeed(m/s) winddir(degrees) rain(cm) gust(m/s)\n')
	for host, time, temp, hdewpt, barometer, windspeed, dir, rain, gust in db.fetchall():
		out.write('%s %f %f %f %f %f %f %f %f\n' % (host, time, temp, barometer, hdewpt, windspeed, dir, rain, gust) )
	
	out.close()

def getVLAWeather(start, stop, filename):

	dbAccess = environ['EVLAMPTS_DB']

	out = open(filename, 'w')
	print 'Collecting VLA Weather data from evlampts db (may take a couple minutes)'

	db = pg.connect(dbAccess).cursor()
	query = """
	SELECT
	  split_part(hostname_timestamp, ',', 1) as hostname,
	  split_part(hostname_timestamp, ',', 2)::double precision as timestamp,
	  COALESCE(Temperature, -999) AS Temperature,
	  COALESCE(Pressure, -999) AS Pressure,
	  COALESCE(Dewpoint_Temperature, -999) AS Dewpoint_Temperature,
	  COALESCE(Wind_Speed_Average, -999) AS Wind_Speed_Average,
	  COALESCE(Wind_Speed_Maximum, -999) AS Wind_Speed_Maximum,
	  COALESCE(Wind_Direction_Average, -999) AS Wind_Direction_Average,
	  COALESCE(Rain_Accumulation, -999) AS Rain_Accumulation
	FROM crosstab(
	  $$ SELECT
	       UPPER(SUBSTR(hostname,1,2)) || ',' || timestamp,
	       monpointname,
	       monpointvalue
	     FROM mcdata
	     WHERE hostname = 'evla-m352' AND
		timestamp BETWEEN %14.8f AND %14.8f AND
	        ( 
		   ( 
		     devicename = 'HMT337' AND
		     monpointname IN ('Temperature', 'Dewpoint_Temperature') 
		   ) 
		 OR
		   ( 
		     devicename = 'WXT520' AND
		     monpointname IN ('Pressure', 'Wind_Speed_Average', 'Wind_Speed_Maximum', 'Wind_Direction_Average', 'Rain_Accumulation')
		   ) 
		 ) 
	     ORDER BY timestamp $$,
	  $$ VALUES ('Temperature'), ('Pressure'), ('Dewpoint_Temperature'), ('Wind_Speed_Average'), ('Wind_Speed_Maximum'), ('Wind_Direction_Average'), ('Rain_Accumulation') $$
	  ) AS ct(hostname_timestamp text,
	          Temperature double precision,
		  Pressure double precision,
		  Dewpoint_Temperature double precision,
		  Wind_Speed_Average double precision,
		  Wind_Speed_Maximum double precision,
		  Wind_Direction_Average double precision,
		  Rain_Accumulation double precision)""" % (start, stop)

	db.execute(query)

	lastP  = -999
	lastDP = -999
	lastWS = -999
	lastWM = -999
	lastWD = -999
	lastR  = -999
	n = 0

	print 'Writing %s' % filename
	out.write('# VLA Weather data produced by %s ver. %s\n' % (program, version) )
	out.write('# Note: Temperature measurements below were made at the associated timestamp.\n')
	out.write('# Other mon points were not collected at exactly the same time.  The previous\n')
	out.write('# measured values for the other monpoints are used.  The maximum timestamp\n')
	out.write('# error is 1 minute.\n')
	out.write('# Ant MJD temp(C) pressure(mbar) dewpoint(C) windspeed(m/s) winddir(degrees) rain(cm) gust(m/s)\n')

	for h, t, T, P, DP, WS, WM, WD, R in db.fetchall():
		if P > -999:
			lastP = P
		if DP > -999:
			lastDP = DP
		if WS > -999:
			lastWS = WS
		if WM > -999:
			lastWM = WM
		if WD > -999:
			lastWD = WD
		if R > -999:
			lastR = R
		if T > -999:
			n += 1
			if n > 2:
				# drop first 2 records in case not all mon points have been seen yet
				out.write('Y %f %f %f %f %f %f %f %f\n' % (t, T, lastP, lastDP, lastWS, lastWD, lastR, lastWM) )
	
	out.close()
				

# Returns obsCode, obsSeg, { ANT : [startmjd, stopmjd] }
def vexPeek(vexFile, verbose):
	
	cmd = 'vexpeek %s' % vexFile
	if verbose > 0:
		print 'Executing command: %s' % cmd

	p = popen(cmd)
	data = p.readlines()

	if len(data) == 0:
		return 'Error', 'Error', 'Error'
	
	obsCode = upper(strip(data[0]))
	obsSeg = ''

	if obsCode[0:5] == 'ERROR':
		return 'Error', 'Error', 'Error'

	if len(obsCode) > 3:
		if obsCode[0].isalpha() and obsCode[1].isalpha() and obsCode[2].isdigit():
			for i in range(3, len(obsCode)):
				if obsCode[i].isalpha():
					obsSeg = obsCode[i:]
					obsCode = obsCode[0:i]
					break;
		if obsCode[0].isalpha() and obsCode[1].isdigit():
			for i in range(2, len(obsCode)):
				if obsCode[i].isalpha():
					obsSeg = obsCode[i:]
					obsCode = obsCode[0:i]
					break;
		

	antennaTimes = {}
	for d in data[1:]:
		s = split(strip(d))
		antennaTimes[upper(s[0])] = [float(s[1]), float(s[2])]

	if verbose > 0:
		print 'This is experiment %s %s' % (obsCode, obsSeg)

	return obsCode, obsSeg, antennaTimes

def exptTimeRange(antennaTimes):
	start = 1e9
	stop = -1e9
	for a in antennaTimes.keys():
		if antennaTimes[a][0] < start:
			start = antennaTimes[a][0]
		if antennaTimes[a][1] > stop:
			stop = antennaTimes[a][0]
	return [start, stop]

def usage(pgm):
	print '\n%s ver. %s  %s  %s' % (program, version, author, verdate)
	print '\nUsage:  %s [options] <vexFile> [ant1 [ant2 ... ] ]' % pgm
	print '\nOptions can be:'
	print '  -h or --help    : print this help information and quit.'
	print '  -f or --force   : force execution, even if output files exist.'
	print '  -v or --verbose : be more verbose in execution.'
	print ''

vexFile = None
antennas = []
verbose = 0
force = False

for a in argv[1:]:
	if a[0] == '-':
		if a in ['-h', '--help']:
			usage(argv[0])
			exit(0)
		if a in ['-v', '--verbose']:
			verbose += 1
		if a in ['-f', '--force']:
			force = True
		else:
			print 'Unknown option: %s' % a
			exit(0)
	elif vexFile == None:
		vexFile = a
	else:
		antennas.append(upper(a))

if vexFile == None:
	print 'Incomplete command line.  Run with --help for info.'
	exit(0)

obsCode, obsSeg, antennaTimes = vexPeek(vexFile, verbose)

if not force:
	nWrongAntenna = 0
	for a in antennas:
		if a not in antennaTimes.keys():
			print 'Error: %s is not in this vex file' % a
			nWrongAntenna += 1
	if nWrongAntenna > 0:
		exit(0)
else:
	for a in antennas:
		if a not in antennaTimes.keys():
			antennaTimes[a] = exptTimeRange(antennaTimes)

if len(antennas) == 0:
	antennas = antennaTimes.keys()


# different antennas have their data stored in different DBs / mon points...

ignored = ''

print '\n%s is now collecting weather data...' % program

start = 1.0e12
stop = -1.0e12
for a in antennas:
	if not a in ['BR', 'FD', 'HN', 'KP', 'LA', 'MK', 'NL', 'OV', 'PT', 'SC']:
		continue
	if antennaTimes[a][0] < start:
		start = antennaTimes[a][0]
	if antennaTimes[a][1] > stop:
		stop = antennaTimes[a][1]

if start < stop:
	vlbaWxData = fetchVLBAWeather(start-0.04, stop+0.04)

for a in antennas:
	if a == 'GB':
		wxFile = '%s%s.%s.weather' % (lower(obsCode), lower(obsSeg), lower(a))
		getGBWeather(antennaTimes[a][0]-0.04, antennaTimes[a][1]+0.04, wxFile)
	elif a == 'Y':
		wxFile = '%s%s.%s.weather' % (lower(obsCode), lower(obsSeg), lower(a))
		getVLAWeather(antennaTimes[a][0]-0.04, antennaTimes[a][1]+0.04, wxFile)
	elif a in ['BR', 'FD', 'HN', 'KP', 'LA', 'MK', 'NL', 'OV', 'PT', 'SC']:
		wxFile = '%s%s.%s.weather' % (lower(obsCode), lower(obsSeg), lower(a))
		getVLBAWeather(vlbaWxData, antennaTimes[a][0]-0.04, antennaTimes[a][1]+0.04, a, wxFile)
	else:
		ignored += ' %s' % a

if len(ignored) > 0:
	print '\nThe following antennas are not (yet) supported by %s: %s' % (program, ignored)

print ''