import dataclasses
import datetime
import typing

import numpy as np
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
	made: datetime.datetime
	latitude: float
	longitude: float
	temperature_difference: pint.Quantity
	convective_cloud_depth: pint.Quantity
	wind: pint.Quantity
	low_clouds: float
	swi: float

	def json(self):
		return {
			"time": self.time.isoformat(),
			"made": self.made.isoformat(),
			"latitude": self.latitude,
			"longitude": self.longitude,
			"dt": self.temperature_difference.m_as("kelvin"),
			"ccd": self.convective_cloud_depth.m_as("foot"),
			"wind": self.wind.m_as("knot"),
			"clouds": self.low_clouds,
			"swi": self.swi,
		}

	@classmethod
	def from_json(cls, json):
		return cls(
			datetime.datetime.fromisoformat(json["time"]),
			datetime.datetime.fromisoformat(json["made"]),
			json["latitude"],
			json["longitude"],
			json["dt"] * metunits.units.kelvin,
			json["ccd"] * metunits.units.foot,
			json["wind"] * metunits.units.knot,
			json["clouds"],
			json["swi"])



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

	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"][0].magnitude / 100
			try:
				swi = szilagyi.calculate_swi(dt, ccd)
			except ValueError:
				swi = -10
			wind = [cast.at("wind_u", _L.H850), cast.at("wind_v", _L.H850)]
			wind = np.array([x.m_as("kts") for x in wind])
			wind = np.linalg.norm(wind) * units.kts
			yield Prediction(cast.timestamp, now, latitude, longitude, dt, ccd, wind, 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