"""Obsplus utilities for working with times."""

from __future__ import annotations

import warnings
from collections.abc import Generator, Sequence
from functools import singledispatch
from typing import Any, TypeVar

import numpy as np
import obspy
import pandas as pd
from obspy.core import event as ev

from obsplus.constants import (
    DEFAULT_TIME,
    LARGEDT64,
    SMALLDT64,
    TIME_COLUMNS,
    event_time_type,
    relative_time_types,
    utc_able_type,
    wave_type,
)
from obsplus.exceptions import TimeOverflowWarning
from obsplus.utils.docs import compose_docstring

rtype = relative_time_types


@singledispatch
def get_reference_time(
    obj: event_time_type | wave_type | None,
) -> obspy.UTCDateTime | None:
    """
    Get a reference time inferred from an object.

    Parameters
    ----------
    obj
        The argument that will indicate a start time. Types and corresponding
        behavior are as follows:

        - ``float``: convert to ``UTCDateTime`` and interpret as a timestamp.
        - ``UTCDateTime``: return a new ``UTCDateTime`` copy.
        - ``catalog``: return the earliest reference time of all events.
        - ``event``: use preferred origin time, or earliest pick time if no
          origin is available.
        - ``stream``: return the earliest reference time of all traces.
        - ``trace``: return ``trace.stats.starttime``.

    Raises
    ------
    TypeError if the type is not supported.
    ValueError if the type is supported but no reference time could be determined.

    Returns
    -------
    obspy.UTCDateTime
    """
    if obj is None:
        return None
    try:
        return to_utc(obj)
    except TypeError:
        msg = f"get_reference_time does not support type {type(obj)}"
        raise TypeError(msg)


@get_reference_time.register(ev.Event)
def _get_event_origin_time(event):
    """Get the time from preferred origin from the event"""
    # try to get origin
    try:
        por = event.preferred_origin() or event.origins[-1]
    except IndexError:
        por = None
    if por is not None:
        if por.time is None:
            msg = f"no origin time found on {por}"
            raise ValueError(msg)
        return get_reference_time(por.time)
    # else try using picks
    elif event.picks:
        return get_reference_time(event.picks)
    else:
        msg = f"could not get reference time for {event}"
        raise ValueError(msg)


@get_reference_time.register(ev.Origin)
def _get_origin_time(origin):
    return get_reference_time(origin.time)


@get_reference_time.register(ev.Pick)
def _get_first_pick(pick):
    """Get the reference time from a single pick."""
    return get_reference_time(pick.time)


@get_reference_time.register(list)
def _from_list(input_list):
    """Get reference time as minimum in a list."""
    outs = [get_reference_time(x) for x in input_list]
    return min([x for x in outs if x is not None])


@get_reference_time.register(obspy.Catalog)
def _get_first_event(catalog):
    """Ensure the events is length one, return event"""
    assert len(catalog) == 1, f"{catalog} has more than one event"
    return _get_event_origin_time(catalog[0])


@get_reference_time.register(obspy.Stream)
def _get_stream_time(st):
    """Return the earliest start time for stream."""
    return min([get_reference_time(tr) for tr in st])


@get_reference_time.register(obspy.Trace)
def _get_trace_time(tr):
    """Return starttime of trace."""
    return tr.stats.starttime


@singledispatch
def to_datetime64(
    value: utc_able_type | Sequence[utc_able_type] | None, default=DEFAULT_TIME
) -> np.datetime64 | np.ndarray | pd.Series:
    """
    Convert time value to a numpy datetime64, or array of such.

    This function essentially uses the same logic as ~:class:`obspy.UTCDateTime`
    but also accounts for native numpy types.

    Parameters
    ----------
    value
        Any value that can be interpreted as a time. If a sequence is passed
        an ndarray of type "datetime64[ns]" is returned.
    default
        A value for missing data. pandas.NaT is used by default.

    Examples
    --------
    >>> import obspy
    >>> import pandas as pd
    >>> timestr = '2020-01-03T11:00:00'
    >>> inputs = [obspy.UTCDateTime(timestr), None, 2525, np.datetime64(timestr)]
    >>> out = to_datetime64(inputs)
    >>> for t in out:
    ...     assert isinstance(t, np.datetime64) or pd.isnull(t)
    >>> assert pd.isnull(out[1])

    """

    def _warn_too_large():
        msg = (
            f"time is too large to represent with a int64 with ns precision,"
            f" downgrading to {new}"
        )
        warnings.warn(msg, category=TimeOverflowWarning)

    # null values return default (usually NaT)
    if pd.isnull(value):
        if not pd.isnull(default):
            return to_datetime64(default)
        return default
    elif isinstance(value, np.datetime64):
        # The '.astype('datetime64[ns]') is necessary to make sure it really
        # does get converted to nanoseconds
        new = value.astype("datetime64[ns]")
        # NOTE: This should get removed after numpy fixes their datetime issue
        # (See #224), but this makes sure that casting to a higher precision
        # datetime didn't cause an int64 overflow
        if not value.dtype == np.dtype("M8[ns]"):
            if new.astype("M8[Y]") < value.astype("M8[Y]"):
                # The year changed, so it has this problem
                _warn_too_large()
                return LARGEDT64
            else:
                return new
        else:
            return new
    elif isinstance(value, pd.Timestamp):
        # The '.astype('datetime64[ns]') is necessary to make sure it really
        # does get converted to nanoseconds
        return value.to_datetime64().astype("datetime64[ns]")
    try:
        utc = obspy.UTCDateTime(value)
        return np.datetime64(utc._ns, "ns")
    # the UTCDateTime is too big or small, use biggest/smallest values instead
    except (SystemError, OverflowError):
        new = LARGEDT64 if np.sign(utc._ns) > 0 else SMALLDT64
        _warn_too_large()
        return new


@to_datetime64.register(str)
def _from_string(time_str: str, default=DEFAULT_TIME):
    """Convert to a string."""
    if not time_str:
        return default
    return to_datetime64(obspy.UTCDateTime(time_str))


@to_datetime64.register(pd.Series)
def _series_to_datetime(value, default=DEFAULT_TIME):
    """Convert a series to datetimes"""
    return value.apply(to_datetime64, default=default)


@to_datetime64.register(np.ndarray)
def _ndarray_to_datetime64(value, default=DEFAULT_TIME):
    """Convert an array to datetimes."""
    ns = np.array([to_datetime64(x, default=default) for x in value])
    return pd.to_datetime(ns, unit="ns").values


@to_datetime64.register(list)
@to_datetime64.register(tuple)
def _sequence_to_datetime64(value, default=DEFAULT_TIME):
    out = [to_datetime64(x, default=default) for x in value]
    return np.array(out)


def to_utc(
    value: utc_able_type | Sequence[utc_able_type],
) -> obspy.UTCDateTime | np.ndarray:
    """
    Convert an object or sequence to a UTCDateTime object.

    Parameters
    ----------
    value
        Any value readable by ~:class:`obspy.UTCDateTime`,
        ~:class:`numpy.datetime64` or a sequence of such.

    Examples
    --------
    >>> import numpy as np
    >>> import obspy
    >>> inputs = [1, '2017-09-18', np.datetime64('2020-01-03')]
    >>> out = to_utc(inputs)
    >>> for utc in out:
    ...     assert isinstance(utc, obspy.UTCDateTime)
    """

    def _dt64_to_utc(dt64):
        if isinstance(dt64, pd.Timestamp):
            dt64 = dt64.to_datetime64()
        ns = dt64.astype("datetime64[ns]").astype(np.int64)
        return obspy.UTCDateTime(ns=int(ns))

    # just use to_datetime64 for flexible handling of types
    dt64ish = to_datetime64(value)
    # return
    if isinstance(dt64ish, np.datetime64):
        return _dt64_to_utc(dt64ish)
    # else assume a sequence of some sort and every element is a dt64
    seq = [_dt64_to_utc(x) for x in dt64ish]
    return np.array(seq)


@singledispatch
def to_timedelta64(
    value: rtype | Sequence[rtype] | np.ndarray | pd.Series | None,
    default: Any = np.timedelta64(0, "s"),
) -> np.timedelta64 | np.ndarray | pd.Series:
    """
    Convert a value to a timedelta[ns].

    Numpy does not gracefully handle non-ints so we need to do some rounding
    first.

    Parameters
    ----------
    value
        A float or an int to convert to datetime.
    default
        The default to return if the input value is not truthy.

    Examples
    --------
    >>> import numpy as np
    >>> # get array of timedelta64[ns] from various input types
    >>> inputs = [1, 2.3232, np.timedelta64(5, 'Y')]
    >>> out = to_timedelta64(inputs)
    >>> assert isinstance(out, np.ndarray)
    """
    if pd.isnull(value):
        return np.timedelta64(default)
    if isinstance(value, np.timedelta64):
        return value
    if isinstance(value, pd.Timedelta):
        return value.to_timedelta64()
    ns = to_utc(value)._ns
    return np.timedelta64(ns, "ns")


@to_timedelta64.register(tuple)
@to_timedelta64.register(list)
def _list_tuple_to_timedelta(value, default=None):
    """Convert sequences to timedeltas."""
    out = [to_datetime64(x, default=default) for x in value]
    return np.array(out)


@to_timedelta64.register(pd.Series)
def _series_to_timedelta(ser) -> pd.Series:
    """Convert a series to a timedelta."""
    if np.issubdtype(ser.dtype, np.timedelta64):
        return ser.astype("timedelta64[ns]")  # convert to ns precision
    # or iterate all elements of series.
    return ser.apply(to_timedelta64)


@to_timedelta64.register(np.ndarray)
def _array_to_timedelta(obj):
    """Convert a series to a timedelta."""
    if np.issubdtype(obj.dtype, np.timedelta64):
        array = obj
    else:
        array = np.array([to_timedelta64(x) for x in obj])
    return array.astype("timedelta64[ns]")


@compose_docstring(time_keys=str(TIME_COLUMNS))
def _dict_times_to_npdatetimes(
    input_dict: dict[str, Any], time_keys: Sequence[str] = TIME_COLUMNS
) -> dict[str, Any]:
    """
    Ensure time values in input_dict are converted to np.datetime64.

    Parameters
    ----------
    input_dict
        A dict that may contain time representations.
    time_keys
        A sequence of keys to search for and convert to np.datetime64.
        Defaults are:
        {time_keys}
    """
    out = dict(input_dict)
    for time_key in set(input_dict) & set(time_keys):
        out[time_key] = to_datetime64(out[time_key])
    return out


@compose_docstring(time_keys=str(TIME_COLUMNS))
def _dict_times_to_ns(
    input_dict: dict[str, Any], time_keys: Sequence[str] = TIME_COLUMNS
) -> dict[str, Any]:
    """
    Ensure time values in input_dict are converted to ints (ns).

    Parameters
    ----------
    input_dict
        A dict that may contain time representations.
    time_keys
        A sequence of keys to search for and convert to np.datetime64.
        Defaults are:
        {time_keys}
    """
    out = dict(input_dict)
    for time_key in set(input_dict) & set(time_keys):
        if not isinstance(out[time_key], int):  # assume ints are ns
            out[time_key] = to_datetime64(out[time_key]).astype(np.int64)
    return out


utc_var = TypeVar("utc_var", utc_able_type, Sequence[utc_able_type])


def make_time_chunks(
    utc1: utc_able_type,
    utc2: utc_able_type,
    duration: float | int,
    overlap: float | int = 0.0,
) -> Generator[tuple[obspy.UTCDateTime, obspy.UTCDateTime], None, None]:
    """
    Yield time intervals fitting in given datetime range.

    Function takes two utc date time objects and create a generator to
    yield all time in between by intervals of duration. Overlap is number
    of seconds segment n will overlap into segment n + 1.

    Parameters
    ----------
    utc1 : obspy.UTCDateTime compatible object
        The start time of the entire interval
    utc2 : obspy.UTCDateTime compatible object
        The end time of the entire interval
    duration : float
        The duration of each chunk
    overlap : float
        The overlap each chunk should have (added at end)

    Yields
    ------
    (time1, time2)

    Examples
    --------
    >>> t1 = obspy.UTCDateTime('2016-01-01')
    >>> t2 = t1 + 2 * 3600
    >>> out = list(make_time_chunks(t1, t2, 3600))
    >>> assert out == [(t1, t1 + 3600), (t1 + 3600, t2)]
    """
    utc1 = to_utc(utc1)
    utc2 = to_utc(utc2)
    overlap = overlap or 0.0
    while utc1 < utc2:
        t2 = utc1 + duration + overlap
        if t2 >= utc2 + overlap:
            t2 = utc2 + overlap
        yield utc1, t2
        utc1 += duration  # add duration
