import numpy
import numpy as np
import datetime
import pytest
from numpy.testing import (
    assert_, assert_equal, assert_raises, assert_warns, suppress_warnings,
    assert_raises_regex,
    )
from numpy.compat import pickle

# Use pytz to test out various time zones if available
try:
    from pytz import timezone as tz
    _has_pytz = True
except ImportError:
    _has_pytz = False

try:
    RecursionError
except NameError:
    RecursionError = RuntimeError  # python < 3.5


class TestDateTime:
    def test_datetime_dtype_creation(self):
        for unit in ['Y', 'M', 'W', 'D',
                     'h', 'm', 's', 'ms', 'us',
                     'ns', 'ps', 'fs', 'as']:
            dt1 = np.dtype('M8[750%s]' % unit)
            assert_(dt1 == np.dtype('datetime64[750%s]' % unit))
            dt2 = np.dtype('m8[%s]' % unit)
            assert_(dt2 == np.dtype('timedelta64[%s]' % unit))

        # Generic units shouldn't add [] to the end
        assert_equal(str(np.dtype("M8")), "datetime64")

        # Should be possible to specify the endianness
        assert_equal(np.dtype("=M8"), np.dtype("M8"))
        assert_equal(np.dtype("=M8[s]"), np.dtype("M8[s]"))
        assert_(np.dtype(">M8") == np.dtype("M8") or
                np.dtype("<M8") == np.dtype("M8"))
        assert_(np.dtype(">M8[D]") == np.dtype("M8[D]") or
                np.dtype("<M8[D]") == np.dtype("M8[D]"))
        assert_(np.dtype(">M8") != np.dtype("<M8"))

        assert_equal(np.dtype("=m8"), np.dtype("m8"))
        assert_equal(np.dtype("=m8[s]"), np.dtype("m8[s]"))
        assert_(np.dtype(">m8") == np.dtype("m8") or
                np.dtype("<m8") == np.dtype("m8"))
        assert_(np.dtype(">m8[D]") == np.dtype("m8[D]") or
                np.dtype("<m8[D]") == np.dtype("m8[D]"))
        assert_(np.dtype(">m8") != np.dtype("<m8"))

        # Check that the parser rejects bad datetime types
        assert_raises(TypeError, np.dtype, 'M8[badunit]')
        assert_raises(TypeError, np.dtype, 'm8[badunit]')
        assert_raises(TypeError, np.dtype, 'M8[YY]')
        assert_raises(TypeError, np.dtype, 'm8[YY]')
        assert_raises(TypeError, np.dtype, 'm4')
        assert_raises(TypeError, np.dtype, 'M7')
        assert_raises(TypeError, np.dtype, 'm7')
        assert_raises(TypeError, np.dtype, 'M16')
        assert_raises(TypeError, np.dtype, 'm16')

    def test_datetime_casting_rules(self):
        # Cannot cast safely/same_kind between timedelta and datetime
        assert_(not np.can_cast('m8', 'M8', casting='same_kind'))
        assert_(not np.can_cast('M8', 'm8', casting='same_kind'))
        assert_(not np.can_cast('m8', 'M8', casting='safe'))
        assert_(not np.can_cast('M8', 'm8', casting='safe'))

        # Can cast safely/same_kind from integer to timedelta
        assert_(np.can_cast('i8', 'm8', casting='same_kind'))
        assert_(np.can_cast('i8', 'm8', casting='safe'))
        assert_(np.can_cast('i4', 'm8', casting='same_kind'))
        assert_(np.can_cast('i4', 'm8', casting='safe'))
        assert_(np.can_cast('u4', 'm8', casting='same_kind'))
        assert_(np.can_cast('u4', 'm8', casting='safe'))

        # Cannot cast safely from unsigned integer of the same size, which
        # could overflow
        assert_(np.can_cast('u8', 'm8', casting='same_kind'))
        assert_(not np.can_cast('u8', 'm8', casting='safe'))

        # Cannot cast safely/same_kind from float to timedelta
        assert_(not np.can_cast('f4', 'm8', casting='same_kind'))
        assert_(not np.can_cast('f4', 'm8', casting='safe'))

        # Cannot cast safely/same_kind from integer to datetime
        assert_(not np.can_cast('i8', 'M8', casting='same_kind'))
        assert_(not np.can_cast('i8', 'M8', casting='safe'))

        # Cannot cast safely/same_kind from bool to datetime
        assert_(not np.can_cast('b1', 'M8', casting='same_kind'))
        assert_(not np.can_cast('b1', 'M8', casting='safe'))
        # Can cast safely/same_kind from bool to timedelta
        assert_(np.can_cast('b1', 'm8', casting='same_kind'))
        assert_(np.can_cast('b1', 'm8', casting='safe'))

        # Can cast datetime safely from months/years to days
        assert_(np.can_cast('M8[M]', 'M8[D]', casting='safe'))
        assert_(np.can_cast('M8[Y]', 'M8[D]', casting='safe'))
        # Cannot cast timedelta safely from months/years to days
        assert_(not np.can_cast('m8[M]', 'm8[D]', casting='safe'))
        assert_(not np.can_cast('m8[Y]', 'm8[D]', casting='safe'))
        # Can cast datetime same_kind from months/years to days
        assert_(np.can_cast('M8[M]', 'M8[D]', casting='same_kind'))
        assert_(np.can_cast('M8[Y]', 'M8[D]', casting='same_kind'))
        # Can't cast timedelta same_kind from months/years to days
        assert_(not np.can_cast('m8[M]', 'm8[D]', casting='same_kind'))
        assert_(not np.can_cast('m8[Y]', 'm8[D]', casting='same_kind'))
        # Can cast datetime same_kind across the date/time boundary
        assert_(np.can_cast('M8[D]', 'M8[h]', casting='same_kind'))
        # Can cast timedelta same_kind across the date/time boundary
        assert_(np.can_cast('m8[D]', 'm8[h]', casting='same_kind'))
        assert_(np.can_cast('m8[h]', 'm8[D]', casting='same_kind'))

        # Cannot cast safely if the integer multiplier doesn't divide
        assert_(not np.can_cast('M8[7h]', 'M8[3h]', casting='safe'))
        assert_(not np.can_cast('M8[3h]', 'M8[6h]', casting='safe'))
        # But can cast same_kind
        assert_(np.can_cast('M8[7h]', 'M8[3h]', casting='same_kind'))
        # Can cast safely if the integer multiplier does divide
        assert_(np.can_cast('M8[6h]', 'M8[3h]', casting='safe'))

        # We can always cast types with generic units (corresponding to NaT) to
        # more specific types
        assert_(np.can_cast('m8', 'm8[h]', casting='same_kind'))
        assert_(np.can_cast('m8', 'm8[h]', casting='safe'))
        assert_(np.can_cast('M8', 'M8[h]', casting='same_kind'))
        assert_(np.can_cast('M8', 'M8[h]', casting='safe'))
        # but not the other way around
        assert_(not np.can_cast('m8[h]', 'm8', casting='same_kind'))
        assert_(not np.can_cast('m8[h]', 'm8', casting='safe'))
        assert_(not np.can_cast('M8[h]', 'M8', casting='same_kind'))
        assert_(not np.can_cast('M8[h]', 'M8', casting='safe'))

    def test_compare_generic_nat(self):
        # regression tests for gh-6452
        assert_(np.datetime64('NaT') !=
                np.datetime64('2000') + np.timedelta64('NaT'))
        assert_(np.datetime64('NaT') != np.datetime64('NaT', 'us'))
        assert_(np.datetime64('NaT', 'us') != np.datetime64('NaT'))

    @pytest.mark.parametrize("size", [
        3, 21, 217, 1000])
    def test_datetime_nat_argsort_stability(self, size):
        # NaT < NaT should be False internally for
        # sort stability
        expected = np.arange(size)
        arr = np.tile(np.datetime64('NaT'), size)
        assert_equal(np.argsort(arr, kind='mergesort'), expected)
    
    @pytest.mark.parametrize("size", [
        3, 21, 217, 1000])
    def test_timedelta_nat_argsort_stability(self, size):
        # NaT < NaT should be False internally for
        # sort stability
        expected = np.arange(size)
        arr = np.tile(np.timedelta64('NaT'), size)
        assert_equal(np.argsort(arr, kind='mergesort'), expected)

    @pytest.mark.parametrize("arr, expected", [
        # the example provided in gh-12629
        (['NaT', 1, 2, 3],
         [1, 2, 3, 'NaT']),
        # multiple NaTs
        (['NaT', 9, 'NaT', -707],
         [-707, 9, 'NaT', 'NaT']),
        # this sort explores another code path for NaT
        ([1, -2, 3, 'NaT'],
         [-2, 1, 3, 'NaT']),
        # 2-D array
        ([[51, -220, 'NaT'],
          [-17, 'NaT', -90]],
         [[-220, 51, 'NaT'],
          [-90, -17, 'NaT']]),
        ])
    @pytest.mark.parametrize("dtype", [
        'M8[ns]', 'M8[us]',
        'm8[ns]', 'm8[us]'])
    def test_datetime_timedelta_sort_nat(self, arr, expected, dtype):
        # fix for gh-12629 and gh-15063; NaT sorting to end of array
        arr = np.array(arr, dtype=dtype)
        expected = np.array(expected, dtype=dtype)
        arr.sort()
        assert_equal(arr, expected)

    def test_datetime_scalar_construction(self):
        # Construct with different units
        assert_equal(np.datetime64('1950-03-12', 'D'),
                     np.datetime64('1950-03-12'))
        assert_equal(np.datetime64('1950-03-12T13', 's'),
                     np.datetime64('1950-03-12T13', 'm'))

        # Default construction means NaT
        assert_equal(np.datetime64(), np.datetime64('NaT'))

        # Some basic strings and repr
        assert_equal(str(np.datetime64('NaT')), 'NaT')
        assert_equal(repr(np.datetime64('NaT')),
                     "numpy.datetime64('NaT')")
        assert_equal(str(np.datetime64('2011-02')), '2011-02')
        assert_equal(repr(np.datetime64('2011-02')),
                     "numpy.datetime64('2011-02')")

        # None gets constructed as NaT
        assert_equal(np.datetime64(None), np.datetime64('NaT'))

        # Default construction of NaT is in generic units
        assert_equal(np.datetime64().dtype, np.dtype('M8'))
        assert_equal(np.datetime64('NaT').dtype, np.dtype('M8'))

        # Construction from integers requires a specified unit
        assert_raises(ValueError, np.datetime64, 17)

        # When constructing from a scalar or zero-dimensional array,
        # it either keeps the units or you can override them.
        a = np.datetime64('2000-03-18T16', 'h')
        b = np.array('2000-03-18T16', dtype='M8[h]')

        assert_equal(a.dtype, np.dtype('M8[h]'))
        assert_equal(b.dtype, np.dtype('M8[h]'))

        assert_equal(np.datetime64(a), a)
        assert_equal(np.datetime64(a).dtype, np.dtype('M8[h]'))

        assert_equal(np.datetime64(b), a)
        assert_equal(np.datetime64(b).dtype, np.dtype('M8[h]'))

        assert_equal(np.datetime64(a, 's'), a)
        assert_equal(np.datetime64(a, 's').dtype, np.dtype('M8[s]'))

        assert_equal(np.datetime64(b, 's'), a)
        assert_equal(np.datetime64(b, 's').dtype, np.dtype('M8[s]'))

        # Construction from datetime.date
        assert_equal(np.datetime64('1945-03-25'),
                     np.datetime64(datetime.date(1945, 3, 25)))
        assert_equal(np.datetime64('2045-03-25', 'D'),
                     np.datetime64(datetime.date(2045, 3, 25), 'D'))
        # Construction from datetime.datetime
        assert_equal(np.datetime64('1980-01-25T14:36:22.5'),
                     np.datetime64(datetime.datetime(1980, 1, 25,
                                                14, 36, 22, 500000)))

        # Construction with time units from a date is okay
        assert_equal(np.datetime64('1920-03-13', 'h'),
                     np.datetime64('1920-03-13T00'))
        assert_equal(np.datetime64('1920-03', 'm'),
                     np.datetime64('1920-03-01T00:00'))
        assert_equal(np.datetime64('1920', 's'),
                     np.datetime64('1920-01-01T00:00:00'))
        assert_equal(np.datetime64(datetime.date(2045, 3, 25), 'ms'),
                     np.datetime64('2045-03-25T00:00:00.000'))

        # Construction with date units from a datetime is also okay
        assert_equal(np.datetime64('1920-03-13T18', 'D'),
                     np.datetime64('1920-03-13'))
        assert_equal(np.datetime64('1920-03-13T18:33:12', 'M'),
                     np.datetime64('1920-03'))
        assert_equal(np.datetime64('1920-03-13T18:33:12.5', 'Y'),
                     np.datetime64('1920'))

    def test_datetime_scalar_construction_timezone(self):
        # verify that supplying an explicit timezone works, but is deprecated
        with assert_warns(DeprecationWarning):
            assert_equal(np.datetime64('2000-01-01T00Z'),
                         np.datetime64('2000-01-01T00'))
        with assert_warns(DeprecationWarning):
            assert_equal(np.datetime64('2000-01-01T00-08'),
                         np.datetime64('2000-01-01T08'))

    def test_datetime_array_find_type(self):
        dt = np.datetime64('1970-01-01', 'M')
        arr = np.array([dt])
        assert_equal(arr.dtype, np.dtype('M8[M]'))

        # at the moment, we don't automatically convert these to datetime64

        dt = datetime.date(1970, 1, 1)
        arr = np.array([dt])
        assert_equal(arr.dtype, np.dtype('O'))

        dt = datetime.datetime(1970, 1, 1, 12, 30, 40)
        arr = np.array([dt])
        assert_equal(arr.dtype, np.dtype('O'))

        # find "supertype" for non-dates and dates

        b = np.bool_(True)
        dm = np.datetime64('1970-01-01', 'M')
        d = datetime.date(1970, 1, 1)
        dt = datetime.datetime(1970, 1, 1, 12, 30, 40)

        arr = np.array([b, dm])
        assert_equal(arr.dtype, np.dtype('O'))

        arr = np.array([b, d])
        assert_equal(arr.dtype, np.dtype('O'))

        arr = np.array([b, dt])
        assert_equal(arr.dtype, np.dtype('O'))

        arr = np.array([d, d]).astype('datetime64')
        assert_equal(arr.dtype, np.dtype('M8[D]'))

        arr = np.array([dt, dt]).astype('datetime64')
        assert_equal(arr.dtype, np.dtype('M8[us]'))

    @pytest.mark.parametrize("unit", [
    # test all date / time units and use
    # "generic" to select generic unit
    ("Y"), ("M"), ("W"), ("D"), ("h"), ("m"),
    ("s"), ("ms"), ("us"), ("ns"), ("ps"),
    ("fs"), ("as"), ("generic") ])
    def test_timedelta_np_int_construction(self, unit):
        # regression test for gh-7617
        if unit != "generic":
            assert_equal(np.timedelta64(np.int64(123), unit),
                         np.timedelta64(123, unit))
        else:
            assert_equal(np.timedelta64(np.int64(123)),
                         np.timedelta64(123))

    def test_timedelta_scalar_construction(self):
        # Construct with different units
        assert_equal(np.timedelta64(7, 'D'),
                     np.timedelta64(1, 'W'))
        assert_equal(np.timedelta64(120, 's'),
                     np.timedelta64(2, 'm'))

        # Default construction means 0
        assert_equal(np.timedelta64(), np.timedelta64(0))

        # None gets constructed as NaT
        assert_equal(np.timedelta64(None), np.timedelta64('NaT'))

        # Some basic strings and repr
        assert_equal(str(np.timedelta64('NaT')), 'NaT')
        assert_equal(repr(np.timedelta64('NaT')),
                     "numpy.timedelta64('NaT')")
        assert_equal(str(np.timedelta64(3, 's')), '3 seconds')
        assert_equal(repr(np.timedelta64(-3, 's')),