from __future__ import annotations

import math
from datetime import datetime, timedelta, tzinfo
from typing import TYPE_CHECKING, overload
from zoneinfo import ZoneInfo

import numpy as np
from dateutil.rrule import rrule

from ..utils import get_timezone, isclose_abs
from .date_utils import Interval, align_limits, expand_datetime_limits
from .types import DateFrequency, date_breaks_info

if TYPE_CHECKING:
    from typing import Generator, Optional, Sequence

    from mizani.typing import (
        Datetime,
        DatetimeBreaksUnits,
        FloatArrayLike,
        NDArrayDatetime,
        NDArrayFloat,
        SeqDatetime,
        TzInfo,
    )


EPOCH = datetime(1970, 1, 1, tzinfo=None)
EPOCH64 = np.datetime64("1970", "Y")
SECONDS_PER_DAY = 24 * 60 * 60
MICROSECONDS_PER_DAY = SECONDS_PER_DAY * (10**6)

NaT_int = np.datetime64("NaT").astype(np.int64)
MIN_DATETIME64 = np.datetime64("0001-01-01")
MAX_DATETIME64 = np.datetime64("10000-01-01")
UTC = ZoneInfo("UTC")
DF = DateFrequency


def _from_ordinalf(x: float, tz: tzinfo | None) -> datetime:
    """
    Convert float array to datetime
    """
    dt64 = EPOCH64 + np.timedelta64(
        int(np.round(x * MICROSECONDS_PER_DAY)), "us"
    )
    if not (MIN_DATETIME64 < dt64 <= MAX_DATETIME64):
        raise ValueError(
            f"Date ordinal {x} converts to {dt64} (using "
            f"epoch {EPOCH}). The supported dates must be  "
            "between year 0001 and 9999."
        )

    # convert from datetime64 to datetime:
    dt: datetime = dt64.astype(object)

    # but maybe we are working in a different timezone so move.
    if tz:
        # datetime64 is always UTC:
        dt = dt.replace(tzinfo=UTC)
        dt = dt.astimezone(tz)

    # fix round off errors
    if np.abs(x) > 70 * 365:
        # if x is big, round off to nearest twenty microseconds.
        # This avoids floating point roundoff error
        ms = round(dt.microsecond / 20) * 20
        if ms == 1000000:
            dt = dt.replace(microsecond=0) + timedelta(seconds=1)
        else:
            dt = dt.replace(microsecond=ms)

    return dt


_from_ordinalf_np_vectorized = np.vectorize(_from_ordinalf, otypes="O")


def get_tzinfo(tz: Optional[str | TzInfo] = None) -> TzInfo | None:
    """
    Generate `~datetime.tzinfo` from a string or return `~datetime.tzinfo`.

    If argument is None, return None.
    """
    if tz is None:
        return None

    if isinstance(tz, str):
        return ZoneInfo(tz)

    if isinstance(tz, tzinfo):
        return tz

    raise TypeError("tz must be string or tzinfo subclass.")


@overload
def datetime_to_num(x: SeqDatetime) -> NDArrayFloat: ...


@overload
def datetime_to_num(x: Datetime) -> float: ...


def datetime_to_num(x: SeqDatetime | Datetime) -> NDArrayFloat | float:
    """
    Convery any datetime sequence to float array
    """
    iterable = np.iterable(x)
    _x = x if iterable else [x]
    try:
        x0 = next(iter(_x))
    except StopIteration:
        return np.array([], dtype=float)

    if isinstance(x0, datetime) and x0.tzinfo:
        _x = [
            dt.astimezone(UTC).replace(tzinfo=None)  # type: ignore
            for dt in _x
        ]
    res = datetime64_to_num(np.asarray(_x, dtype="datetime64"))
    return res if iterable else res[0]


def datetime64_to_num(x: NDArrayDatetime) -> NDArrayFloat:
    """
    Convery any numpy datetime64 array to float array
    """
    x_secs = x.astype("datetime64[s]")
    diff_ns = (x - x_secs).astype("timedelta64[ns]")

    # In seconds + nanoseconds, in days
    res = (
        (x_secs - EPOCH64).astype(np.float64)
        + diff_ns.astype(np.float64) / 1.0e9
    ) / SECONDS_PER_DAY

    x_int = x.astype(np.int64)
    res[x_int == NaT_int] = np.nan
    return res


def num_to_datetime(
    x: FloatArrayLike, tz: Optional[str | TzInfo] = None
) -> NDArrayDatetime:
    """
    Convert any float array to numpy datetime64 array
    """
    tz = get_tzinfo(tz) or UTC
    return _from_ordinalf_np_vectorized(x, tz)


WIDTHS: dict[DateFrequency, Sequence[int]] = {
    DF.YEARLY: (1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000),
    DF.MONTHLY: (1, 2, 3, 4, 6),
    DF.DAILY: (1, 2, 4, 7, 14),
    DF.HOURLY: (1, 2, 3, 4, 6, 12),
    DF.MINUTELY: (1, 5, 10, 15, 30),
    DF.SECONDLY: (1, 5, 10, 15, 30),
    DF.MICROSECONDLY: (
        1,
        2,
        5,
        10,
        20,
        50,
        100,
        200,
        500,
        1000,
        2000,
        5000,
        10000,
        20000,
        50000,
        100000,
        200000,
        500000,
        1000000,
    ),
}

MAX_BREAKS = dict(zip(DateFrequency, (11, 12, 11, 12, 11, 11, 8)))


def _viable_freqs(
    min_breaks: int,
    unit_durations: tuple[int, int, int, int, int, int, int],
) -> Generator[tuple[DateFrequency, int], None, None]:
    """
    Find viable frequency, duration pairs

    A pair is viable if it can yeild a suitable number of breaks
    For example:
        - YEARLY frequency, 3 year unit_duration and
          8 min_breaks is not viable
        - MONTHLY frequency, 36 month unit_duration and
          8 min_breaks is viable
    """
    for freq, duration in zip(DateFrequency, unit_durations):
        max_width = WIDTHS[freq][-1]
        max_breaks = max(min_breaks, MAX_BREAKS.get(freq, 11))
        if duration <= max_width * max_breaks - 1:
            yield (freq, duration)


def calculate_date_breaks_info(
    limits: tuple[datetime, datetime], n: int = 5
) -> date_breaks_info:
    """
    Calculate information required to generate breaks

    Parameters
    ----------
    limits:
        Datetime limits for the breaks
    n:
        Desired number of breaks.
    """
    tz = get_timezone(limits)
    _max_breaks_lookup = {f: max(b, n) for f, b in MAX_BREAKS.items()}

    # Widen the duration at each granularity
    itv = Interval(*limits)
    unit_durations = (
        itv.y_wide,
        itv.M_wide,
        itv.d_wide,
        itv.h_wide,
        itv.m_wide,
        itv.s,
        itv.u,
    )
    # Search frequencies from longest (yearly) to the smallest
    # for one that would a good width between the breaks

    # Defaults
    freq = DF.YEARLY  # makes pyright happy
    break_width = 1
    duration = n

    for freq, duration in _viable_freqs(n, unit_durations):
        # Search for breaks in the range
        _max_breaks = _max_breaks_lookup[freq]
        for mb in range(n, _max_breaks + 1):
            # There are few breaks at this frequency
            # e.g. (freq=YEARLY, duration=2) but mb = 5
            if duration < mb:
                continue

            for break_width in WIDTHS[freq]:
                if duration <= break_width * mb - 1:
                    break
            else:  # pragma: no cover
                continue
            break
        else:
            continue
        break

    num_breaks = duration // break_width
    limits = itv.limits_for_frequency(freq)
    res = date_breaks_info(
        freq,
        num_breaks,
        break_width,
        start=limits[0].replace(tzinfo=tz),
        until=limits[1].replace(tzinfo=tz),
        tz=tz,
    )
    return res


def calculate_date_breaks_auto(limits, n: int = 5) -> Sequence[datetime]:
    """
    Calcuate date breaks using appropriate units
    """
    info = calculate_date_breaks_info(limits, n=n)
    lookup = {
        DF.YEARLY: yearly_breaks,
        DF.MONTHLY: monthly_breaks,
        DF.DAILY: daily_breaks,
        DF.HOURLY: hourly_breaks,
        DF.MINUTELY: minutely_breaks,
        DF.SECONDLY: secondly_breaks,
        DF.MICROSECONDLY: microsecondly_breaks,
    }
    return lookup[info.frequency](info)


def calculate_date_breaks_byunits(
    limits,
    units: DatetimeBreaksUnits,
    width: int,
    max_breaks: Optional[int] = None,
    tz: Optional[TzInfo] = None,
) -> Sequence[datetime]:
    """
    Calcuate date breaks using appropriate units
    """
    timely_name = f"{units.upper()}LY"
    if timely_name in ("DAYLY", "WEEKLY"):
        timely_name = "DAILY"

    freq = getattr(DF, timely_name)

    # Appropriate start and end dates
    start, until = expand_datetime_limits(limits, width, units)

    if units == "week":
        width *= 7

    info = date_breaks_info(
        freq,
        n=-1,
        width=width,
        start=start,
        until=until,
        tz=tz,
    )

    lookup = {
        "year": rrulely_breaks,
        "month": rrulely_breaks,
        "week": rrulely_breaks,
        "day": rrulely_breaks,
        "hour": rrulely_breaks,
        "minute": rrulely_breaks,
        "second": rrulely_breaks,
        "microsecond": microsecondly_breaks,
    }
    return lookup[units](info)


def rrulely_breaks(info: date_breaks_info) -> Sequence[datetime]:
    r = rrule(
        info.frequency,
        interval=info.width,
        dtstart=info.start,
        until=info.until,
    )
    return list(r)


def yearly_breaks(info: date_breaks_info) -> Sequence[datetime]:
    """
    Calculate yearly breaks
    """
    # New limits so that breaks fall on multiples of
    # the width
    limits = info.start.year, info.until.year
    l, h = align_limits(limits, info.width)
    l, h = math.floor(l), math.ceil(h)

    _replace_d = {
        "month": 1,
        "day": 1,
        "hour": 0,
        "minute": 0,
        "second": 0,
        "tzinfo": info.tz,
    }

    start = info.start.replace(year=l, **_replace_d)
    until = info.until.replace(year=h, **_replace_d)
    r = rrule(
        info.frequency,
        interval=info.width,
        dtstart=start,
        until=until,
    )
    return list(r)


def monthly_breaks(info: date_breaks_info) -> Sequence[datetime]:
    """
    Calculate monthly breaks
    """
    r = rrule(
        info.frequency,
        interval=info.width,
        dtstart=info.start,
        until=info.until,
    )
    return list(r)


def daily_breaks(info: date_breaks_info) -> Sequence[datetime]:
    """
    Calculate daily breaks
    """
    if info.width == 7:
        interval = 1
        bymonthday = (1, 8, 15, 22)
    elif info.width == 14:
        interval = 1
        bymonthday = (1, 15)
    else:
        interval = info.width
        bymonthday = None

    r = rrule(
        info.frequency,
        interval=interval,
        dtstart=info.start,
        until=info.until,
        bymonthday=bymonthday,
    )
    return r.between(info.start, info.until, True)


def hourly_breaks(info: date_breaks_info) -> Sequence[datetime]:
    """
    Calculate hourly breaks
    """
    r = rrule(
        info.frequency,
        interval=1,
        dtstart=info.start,
        until=info.until,
        byhour=range(0, 24, info.width),
    )
    return r.between(info.start, info.until, True)


def minutely_breaks(info: date_breaks_info) -> Sequence[datetime]:
    """
    Calculate minutely breaks
    """
    r = rrule(
        info.frequency,
        interval=1,
        dtstart=info.start,
        until=info.until,
        byminute=range(0, 60, info.width),
    )
    return r.between(info.start, info.until, True)


def secondly_breaks(info: date_breaks_info) -> Sequence[datetime]:
    """
    Calculate minutely breaks
    """
    r = rrule(
        info.frequency,
        interval=1,
        dtstart=info.start,
        until=info.until,
        bysecond=range(0, 60, info.width),
    )
    return r.between(info.start, info.until, True)


def microsecondly_breaks(info: date_breaks_info) -> NDArrayDatetime:
    """
    Calculate breaks at microsecond intervals
    """
    # TODO: A little too complicated, could use some refactoring
    nmin: float
    nmax: float
    width = info.width

    nmin, nmax = datetime_to_num((info.start, info.until))
    day0: float = np.floor(nmin)

    # difference in microseconds
    umax = (nmax - day0) * MICROSECONDS_PER_DAY
    umin = (nmin - day0) * MICROSECONDS_PER_DAY

    # Ensure max is a multiple of the width
    width = info.width
    h, m = divmod(umax, width)
    if not isclose_abs(m / width, 0):
        h += 1
    umax = h * width

    # Generate breaks at the multiples of the width
    n = (umax - umin + 0.001 * width) // width
    ubreaks = umin - width + np.arange(n + 3) * width
    breaks = day0 + ubreaks / MICROSECONDS_PER_DAY
    return num_to_datetime(breaks, info.tz)
