from __future__ import annotations

import functools as ft
from fractions import Fraction
from typing import Literal

import equinox as eqx
import jax.numpy as jnp
import numpy as np
from jaxtyping import Array, ArrayLike, Float

from ._geo import EARTH_RADIUS


ft.total_ordering


class Unit(eqx.Module):
    """
    Base class representing [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`].

    Methods
    -------
    __eq__(other)
        Checks if two [`pastax.utils.Unit`][] are equal.
    __lt__(other)
        Checks if one [`pastax.utils.Unit`][] is less than another (using their name).
    __hash__()
        Returns the hash of the [`pastax.utils.Unit`].
    __repr__()
        Returns the string representation of the [`pastax.utils.Unit`].
    convert_to(unit, value, exp, *args)
        Converts the `value` to the specified [`pastax.utils.Unit`].
    """

    name: str = eqx.field(static=True, default_factory=lambda: "")

    def __eq__(self, other):
        if isinstance(other, Unit):
            return self.name == other.name
        return False

    def __lt__(self, other):
        if isinstance(other, Unit):
            return self.name < other.name
        return NotImplemented

    def __hash__(self):
        return hash(self.name)

    def __repr__(self):
        return self.name

    @staticmethod
    def _pre_convert(value: Array, exp: int | float) -> Array:
        """
        Prepares the value for conversion between "base" [`pastax.utils.Unit`][] by raising it to the power of the
        reciprocal of the exponent.

        Parameters
        ----------
        value : Array
            The value to be converted.
        exp : int or float
            The exponent to use for conversion.

        Returns
        -------
        Array
            The prepared value.
        """
        if exp == 1:
            return value
        return value ** (1 / exp)

    @staticmethod
    def _post_convert(value: Array, exp: int | float) -> Array:
        """
        Finalizes the conversion between "base" [`pastax.utils.Unit`][] by raising the value to the power of the
        exponent.

        Parameters
        ----------
        value : Array
            The value to be converted.
        exp : int or float
            The exponent to use for conversion.

        Returns
        -------
        Array
            The converted value.
        """
        if exp == 1:
            return value
        return value**exp

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the `value` to the specified [`pastax.utils.Unit`].

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        NotImplementedError
            If not implemented by subclasses.
        """
        raise NotImplementedError


class Meters(Unit):
    """
    Class representing meters as a [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`], set to `"m"`.
    """

    name: str = eqx.field(static=True, default_factory=lambda: "m")

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the value to the specified [`pastax.utils.Unit`].

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        ValueError
            If the conversion is not possible.
        """
        if isinstance(unit, Meters):
            return value

        value = self._pre_convert(value, exp)

        if isinstance(unit, Kilometers):
            value = meters_to_kilometers(value)
        elif isinstance(unit, LatLonDegrees):
            value = meters_to_degrees(value, *args)
        else:
            raise ValueError(f"Cannot convert {self} to {unit}")

        return self._post_convert(value, exp)


class Kilometers(Unit):
    """
    Class representing kilometers as a [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`], set to `"km"`.
    """

    name: str = eqx.field(static=True, default_factory=lambda: "km")

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the value to the specified [`pastax.utils.Unit`].

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        ValueError
            If the conversion is not possible.
        """
        if isinstance(unit, Kilometers):
            return value

        value = self._pre_convert(value, exp)

        if isinstance(unit, Meters):
            value = kilometers_to_meters(value)
        elif isinstance(unit, LatLonDegrees):
            value = kilometers_to_degrees(value, *args)
        else:
            raise ValueError(f"Cannot convert {self} to {unit}")

        return self._post_convert(value, exp)


class LatLonDegrees(Unit):
    """
    Class representing latitude and longitude degrees as a [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`], set to `"°"`.
    """

    name: str = eqx.field(static=True, default_factory=lambda: "°")

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the value to the specified [`pastax.utils.Unit`].

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        ValueError
            If the conversion is not possible.
        """
        if isinstance(unit, LatLonDegrees):
            return value

        value = self._pre_convert(value, exp)

        if isinstance(unit, Meters):
            value = degrees_to_meters(value, *args)
        elif isinstance(unit, Kilometers):
            value = degrees_to_kilometers(value, *args)
        else:
            raise ValueError(f"Cannot convert {self} to {unit}")

        return self._post_convert(value, exp)


class Seconds(Unit):
    """
    Class representing seconds as a [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`], set to `"s"`.
    """

    name: str = eqx.field(static=True, default_factory=lambda: "s")

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the value to the specified unit.

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        ValueError
            If the conversion is not possible.
        """
        if isinstance(unit, Seconds):
            return value

        value = self._pre_convert(value, exp)

        if isinstance(unit, Minutes):
            value = seconds_to_minutes(value)
        elif isinstance(unit, Hours):
            value = seconds_to_hours(value)
        elif isinstance(unit, Days):
            value = seconds_to_days(value)
        else:
            raise ValueError(f"Cannot convert {self} to {unit}")

        return self._post_convert(value, exp)


class Minutes(Unit):
    """
    Class representing minutes as a [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`], set to `"min"`.
    """

    name: str = eqx.field(static=True, default_factory=lambda: "min")

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the value to the specified unit.

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        ValueError
            If the conversion is not possible.
        """
        if isinstance(unit, Minutes):
            return value

        value = self._pre_convert(value, exp)

        if isinstance(unit, Seconds):
            value = minutes_to_seconds(value)
        elif isinstance(unit, Hours):
            value = minutes_to_hours(value)
        elif isinstance(unit, Days):
            value = minutes_to_days(value)
        else:
            raise ValueError(f"Cannot convert {self} to {unit}")

        return self._post_convert(value, exp)


class Hours(Unit):
    """
    Class representing hours as a [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`], set to `"h"`.
    """

    name: str = eqx.field(static=True, default_factory=lambda: "h")

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the value to the specified unit.

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        ValueError
            If the conversion is not possible.
        """
        if isinstance(unit, Hours):
            return value

        value = self._pre_convert(value, exp)

        if isinstance(unit, Seconds):
            value = hours_to_seconds(value)
        elif isinstance(unit, Minutes):
            value = hours_to_minutes(value)
        elif isinstance(unit, Days):
            value = hours_to_days(value)
        else:
            raise ValueError(f"Cannot convert {self} to {unit}")

        return self._post_convert(value, exp)


class Days(Unit):
    """
    Class representing days as a [`pastax.utils.Unit`][] of measurement.

    Attributes
    ----------
    name : str
        The name of the [`pastax.utils.Unit`], set to `"d"`.
    """

    name: str = eqx.field(static=True, default_factory=lambda: "d")

    def convert_to(self, unit: Unit, value: Array, exp: int | float = 1, *args) -> Array:
        """
        Converts the value to the specified unit.

        Parameters
        ----------
        unit : Unit
            The [`pastax.utils.Unit`][] to convert to.
        value : Array
            The value to convert.
        exp : int or float, optional
            The exponent to use for conversion, defaults to 1.
        *args
            Additional arguments for conversion.

        Returns
        -------
        Array
            The converted value.

        Raises
        ------
        ValueError
            If the conversion is not possible.
        """
        if isinstance(unit, Days):
            return value

        value = self._pre_convert(value, exp)

        if isinstance(unit, Seconds):
            value = days_to_seconds(value)
        elif isinstance(unit, Minutes):
            value = days_to_minutes(value)
        elif isinstance(unit, Hours):
            value = days_to_hours(value)
        else:
            raise ValueError(f"Cannot convert {self} to {unit}")

        return self._post_convert(value, exp)


UNIT = {
    "m": Meters(),
    "km": Kilometers(),
    "°": LatLonDegrees(),
    "s": Seconds(),
    "min": Minutes(),
    "h": Hours(),
    "d": Days(),
}
"""
A dictionary mapping unit symbols to their corresponding [`pastax.utils.Unit`][] objects.

Keys
----
"m" : Meters
    Represents meters as a [`pastax.utils.Unit`][] of measurement.
"km" : Kilometers
    Represents kilometers as a [`pastax.utils.Unit`][] of measurement.
"°" : LatLonDegrees
    Represents latitude and longitude degrees as a [`pastax.utils.Unit`][] of measurement.
"s" : Seconds
    Represents seconds as a [`pastax.utils.Unit`][] of measurement.
"min" : Minutes
    Represents minutes as a [`pastax.utils.Unit`][] of measurement.
"h" : Hours
    Represents hours as a [`pastax.utils.Unit`][] of measurement.
"d" : Days
    Represents days as a [`pastax.utils.Unit`][] of measurement.

Values
------
Unit
    The corresponding [`pastax.utils.Unit`][] object for each unit symbol.
"""


def units_to_str(unit: dict[Unit, int | float]) -> str:
    """
    Converts a dictionary of [`pastax.utils.Unit`][] with their exponents to a string representation.

    Parameters
    ----------
    unit : dict[Unit, int or float]
        A dictionary of [`pastax.utils.Unit`][] with their exponents.

    Returns
    -------
    str
        A string representation of the [`pastax.utils.Unit`][] with their exponents.
    """

    def get_exp_str(exp: int | float) -> str:
        if exp == 1:
            return ""
        else:
            return f"^{{{Fraction(exp).limit_denominator()}}}"

    def get_dim_str(dim: Unit, exp: int | float) -> str:
        if exp == 0:
            return ""
        else:
            return f"{dim}{get_exp_str(exp)}"

    return " ".join(get_dim_str(dim, exp) for dim, exp in unit.items()).strip()


def compose_units(
    unit1: dict[Unit, int | float],
    unit2: dict[Unit, int | float] | None,
    mul: Literal[-1, 1],
) -> dict[Unit, int | float]:
    """
    Compose two [`pastax.utils.Unit`][] dictionaries by combining their values,
    optionally multiplying the second dictionary's values by a factor to account for multiplication or division.

    Parameters
    ----------
    unit1 : dict[Unit, int | float]
        The first [`pastax.utils.Unit`][] dictionary.
    unit2 : dict[Unit, int | float]
        The second [`pastax.utils.Unit`][] dictionary.
    mul : Literal[-1, 1]
        The multiplier for the second [`pastax.utils.Unit`][] dictionary's values.
        Should be either 1 in case of multiplication or -1 in case of division.

    Returns
    -------
    dict[Unit, int | float]
        The composed [`pastax.utils.Unit`][] dictionary, or an empty dictionary if both input dictionaries are empty.
    """
    if (not unit1) and (not unit2):
        return {}
    if not unit1:
        return unit2  # type: ignore
    if not unit2:
        return unit1

    unit = unit1.copy()

    for k, v in unit2.items():
        v *= mul
        if k in unit:
            unit[k] += v
        else:
            unit[k] = v

    return unit


def meters_to_kilometers(arr: Array) -> Array:
    """
    Converts an array of distances from meters to kilometers.

    Parameters
    ----------
    arr : Array
        An array of distances in meters.

    Returns
    -------
    Array
        An array of distances in kilometers.
    """
    return arr / 1000


def meters_to_degrees(arr: Float[Array, "... 2"], latitude: Float[Array, "..."]) -> Float[Array, "... 2"]:
    """
    Converts an array of latitude/longitude distances from meters to degrees.

    Parameters
    ----------
    arr : Float[Array, "... 2"]
        An array of latitude/longitude distances in meters.
    latitude : Float[Array, "..."]
        The latitude in degrees at which the conversion is to be performed.

    Returns
    -------
    Float[Array, "... 2"]
        An array of latitude/longitude distances in degrees.
    """
    arr = jnp.degrees(arr / EARTH_RADIUS)
    arr = arr.at[..., 1].divide(jnp.cos(jnp.radians(latitude)))
    return arr


def kilometers_to_meters(arr: Array) -> Array:
    """
    Converts an array of distances from kilometers to meters.

    Parameters
    ----------
    arr : Array
        An array of distances in kilometers.

    Returns
    -------
    Array
        An array of distances in meters.
    """
    return arr * 1000


def kilometers_to_degrees(arr: Float[Array, "... 2"], latitude: Float[Array, "..."]) -> Float[Array, "... 2"]:
    """
    Converts an array of latitude/longitude distances from kilometers to degrees.

    Parameters
    ----------
    arr : Float[Array, "... 2"]
        An array of latitude/longitude distances in kilometers.
    latitude : Float[Array, "..."]
        The latitude in degrees at which the conversion is to be performed.

    Returns
    -------
    Float[Array, "... 2"]
        An array of latitude/longitude distances in degrees.
    """
    return meters_to_degrees(kilometers_to_meters(arr), latitude)


def degrees_to_meters(arr: Float[Array, "... 2"], latitude: Float[Array, "..."]) -> Float[Array, "... 2"]:
    """
    Converts an array of latitude/longitude distances from degrees to meters.

    Parameters
    ----------
    arr : Float[Array, "... 2"]
        An array of latitude/longitude distances in degrees.
    latitude : Float[Array, "..."]
        The latitude in degrees at which the conversion is to be performed.

    Returns
    -------
    Float[Array, "... 2"]
        An array of latitude/longitude distances in meters.
    """
    arr = jnp.radians(arr) * EARTH_RADIUS
    arr = arr.at[..., 1].multiply(jnp.cos(jnp.radians(latitude)))
    return arr


def degrees_to_kilometers(arr: Float[Array, "... 2"], latitude: Float[Array, "..."]) -> Float[Array, "... 2"]:
    """
    Converts an array of latitude/longitude distances from degrees to kilometers.

    Parameters
    ----------
    arr : Float[Array, "... 2"]
        An array of latitude/longitude distances in degrees.
    latitude : Float[Array, "..."]
        The latitude in degrees at which the conversion is to be performed.

    Returns
    -------
    Float[Array, "... 2"]
        An array of latitude/longitude distances in kilometers.
    """
    return meters_to_kilometers(degrees_to_meters(arr, latitude))


def seconds_to_minutes(arr: Array) -> Array:
    """
    Converts an array of time durations from seconds to minutes.

    Parameters
    ----------
    arr : Array
        An array of time durations in seconds.

    Returns
    -------
    Array
        An array of time durations in minutes.
    """
    return arr / 60


def seconds_to_hours(arr: Array) -> Array:
    """
    Converts an array of time durations from seconds to hours.

    Parameters
    ----------
    arr : Array
        An array of time durations in seconds.

    Returns
    -------
    Array
        An array of time durations in hours.
    """
    return minutes_to_hours(seconds_to_minutes(arr))


def seconds_to_days(arr: Array) -> Array:
    """
    Converts an array of time durations from seconds to days.

    Parameters
    ----------
    arr : Array
        An array of time durations in seconds.

    Returns
    -------
    Array
        An array of time durations in days.
    """
    return minutes_to_days(seconds_to_minutes(arr))


def minutes_to_seconds(arr: Array) -> Array:
    """
    Converts an array of time durations from minutes to seconds.

    Parameters
    ----------
    arr : Array
        An array of time durations in minutes.

    Returns
    -------
    Array
        An array of time durations in seconds.
    """
    return arr * 60


def minutes_to_hours(arr: Array) -> Array:
    """
    Converts an array of time durations from minutes to hours.

    Parameters
    ----------
    arr : Array
        An array of time durations in minutes.

    Returns
    -------
    Array
        An array of time durations in hours.
    """
    return arr / 60


def minutes_to_days(arr: Array) -> Array:
    """
    Converts an array of time durations from minutes to days.

    Parameters
    ----------
    arr : Array
        An array of time durations in minutes.

    Returns
    -------
    Array
        An array of time durations in days.
    """
    return hours_to_days(minutes_to_hours(arr))


def hours_to_seconds(arr: Array) -> Array:
    """
    Converts an array of time durations from hours to seconds.

    Parameters
    ----------
    arr : Array
        An array of time durations in hours.

    Returns
    -------
    Array
        An array of time durations in seconds.
    """
    return minutes_to_seconds(hours_to_minutes(arr))


def hours_to_minutes(arr: Array) -> Array:
    """
    Converts an array of time durations from hours to minutes.

    Parameters
    ----------
    arr : Array
        An array of time durations in hours.

    Returns
    -------
    Array
        An array of time durations in minutes.
    """
    return arr * 60


def hours_to_days(arr: Array) -> Array:
    """
    Converts an array of time durations from hours to days.

    Parameters
    ----------
    arr : Array
        An array of time durations in hours.

    Returns
    -------
    Array
        An array of time durations in days.
    """
    return arr / 24


def days_to_seconds(arr: Array) -> Array:
    """
    Converts an array of time durations from days to seconds.

    Parameters
    ----------
    arr : Array
        An array of time durations in days.

    Returns
    -------
    Array
        An array of time durations in seconds.
    """
    return hours_to_seconds(days_to_hours(arr))


def days_to_minutes(arr: Array) -> Array:
    """
    Converts an array of time durations from days to minutes.

    Parameters
    ----------
    arr : Array
        An array of time durations in days.

    Returns
    -------
    Array
        An array of time durations in minutes.
    """
    return hours_to_minutes(days_to_hours(arr))


def days_to_hours(arr: Array) -> Array:
    """
    Converts an array of time durations from days to hours.

    Parameters
    ----------
    arr : Array
        An array of time durations in days.

    Returns
    -------
    Array
        An array of time durations in hours.
    """
    return arr * 24


def time_in_seconds(arr: Array | ArrayLike) -> Array | ArrayLike:
    """
    Converts an array of datetime64 values to seconds since the Unix epoch.

    Parameters
    ----------
    arr : Array | ArrayLike
        An array of datetime64 values or a single datetime64 value.

    Returns
    -------
    Array | ArrayLike
        An array of integers representing the number of seconds since the Unix epoch.
    """
    if isinstance(arr, np.datetime64) or (isinstance(arr, np.ndarray) and np.issubdtype(arr.dtype, np.datetime64)):
        arr = arr.astype("datetime64[s]").astype(int)

    return arr
