import dataclasses
import datetime
import typing

import pint
import szilagyi
from geopy import geocoders
from metpy import calc
from metpy import units as metunits
from windy import point_forecast, Windy

_L = point_forecast.Level


@dataclasses.dataclass
class Prediction:
	time: datetime.datetime
	latitude: float
	longitude: float
	temperature_difference: pint.Quantity
	convective_cloud_depth: pint.Quantity
	wind: pint.Quantity
	low_clouds: float
	swi: float


def calculate(config) -> typing.List[Prediction]:
	units = metunits.units
	windy = Windy(units)

	def _calculate(latitude, longitude):
		forecasts = windy.point_forecast(
			config.key,
			latitude, longitude,
			point_forecast.Model.ICONEU,
			("temp", "dewpoint", "wind", "pressure", "lclouds"),
			tuple(_L))
		for cast in forecasts:
			dt = abs(cast.at("temp", _L.H850) - cast.at("temp", _L.SURFACE))
			pressure, _ = calc.lcl(
				cast.at("pressure", _L.SURFACE),
				cast.at("temp", _L.SURFACE),
				cast.at("dewpoint", _L.SURFACE))
			lcl = calc.pressure_to_height_std(pressure)
			pressure, _ = calc.el(cast["pressure"], cast["temp"], cast["dewpoint"])
			el = calc.pressure_to_height_std(pressure)
			ccd = (el - lcl).to(units.ft)
			clouds = cast["lclouds"].magnitude / 100
			try:
				swi = szilagyi.calculate_swi(dt, ccd)
			except ValueError:
				swi = -10
			yield Prediction(cast.timestamp, latitude, longitude, dt, ccd, 0, clouds, swi)

	predictions = []
	locator = geocoders.Nominatim(user_agent="waterspout-radar")
	for location in config.locations:
		found = locator.geocode(location)
		predictions.extend(_calculate(found.latitude, found.longitude))
	return predictions