"""Utilisation functions."""

from __future__ import annotations

import pandas as pd
import reverse_geocoder as rg

CONTINENT_ZONES = {
    "europe": [6136, 22095, 22096, 22190, 22194, 43153, 43170, 46171, 46172],
    "north_america": [6121, 6137, 6191, 14104, 33117, 37176, 46173],
    "south_america": [2154, 20087, 44149],
    "africa": [7091, 9028, 9030, 10162, 12048, 32114, 32121, 34120, 36124, 38127],
    "asia": [11047, 17135, 17166, 25143, 25147, 43134],
    "australia": [],  # TODO: What zones for Australia?
}
CONTINENT_ZONES_REV = {x: k for k, v in CONTINENT_ZONES.items() for x in v}


def transform_countries(df: pd.DataFrame, cols: list[str] = ["lat", "lon"]) -> None:
    """Filter results based on countries."""
    coordinates = [(x[0], x[1]) for x in df[cols].values]
    df["country"] = _predict_countries(coordinates)


def _predict_countries(coordinates: list[tuple[float, float]]) -> str:
    """Predict the country based on the latitude and longitude."""
    return [x.get("cc", "Other") for x in rg.search(coordinates, verbose=False)]


def transform_continents(df: pd.DataFrame) -> None:
    """Transform the zone-ID to a continent."""
    if "aez_zoneid" in df.columns:
        df["continent"] = df[["aez_zoneid"]].apply(
            lambda x: CONTINENT_ZONES_REV[x["aez_zoneid"]],
            axis=1,
        )
    elif ("lat" in df.columns) and ("lon" in df.columns):
        df["continent"] = df[["lat", "lon"]].apply(
            lambda x: _predict_continent(lat=x["lat"], lon=x["lon"]),
            axis=1,
        )
    else:
        df["continent"] = "Other"


def _predict_continent(lat: float, lon: float) -> str:
    """Predict the continent based on the latitude and longitude."""
    if (lat > 34) and (lon < 45) and (lon > -20):
        return "europe"
    elif (lat > 25) and (lon < -20):
        return "north_america"
    else:
        return "Other"
