import argparse
import os

import szilagyi
from geopy.geocoders import Nominatim
from metpy import calc
from metpy.units import units
from windy import point_forecast
from windy import Windy


def main():
	parser = argparse.ArgumentParser()
	parser.add_argument("location")
	args = parser.parse_args()
	key = os.getenv("WINDY_KEY")
	if not key:
		raise RuntimeError("expected windy's api key")
	locator = Nominatim(user_agent="waterspout-radar")
	location = locator.geocode(args.location)
	coords = (location.latitude, location.longitude)
	registry = units
	model = point_forecast.Model.ICONEU
	levels = (point_forecast.Level.SURFACE, point_forecast.Level.H850, point_forecast.Level.H700, point_forecast.Level.H500)
	windy = Windy(registry)
	forecast = windy.point_forecast(key, *coords, model, ("temp", "dewpoint", "wind", "pressure"), levels)
	for entry in forecast:
		dt = abs(entry["temp-850h"] - entry["temp-surface"])
		pressure, _ = calc.lcl(entry["pressure-surface"], entry["temp-surface"], entry["dewpoint-surface"])
		lcl = calc.pressure_to_height_std(pressure)
		pressure, _ = calc.el(
			[entry["pressure-surface"].m_as(units.hPa), 850, 700, 500] * units.hPa,
			[entry["temp-" + str(x)].m_as(units.degK) for x in levels] * units.degK,
			[entry["dewpoint-" + str(x)].m_as(units.degK) for x in levels] * units.degK)
		el = calc.pressure_to_height_std(pressure)
		ccd = (el - lcl).to(units.ft)
		try:
			swi = szilagyi.calculate_swi(dt, ccd)
		except ValueError:
			swi = -10
		print(entry.timestamp, dt, ccd, swi)


if __name__ == '__main__':
	main()