from datetime import (
    date,
    timedelta,
)
import re

import numpy as np
import pytest

from pandas._libs.tslibs.timezones import maybe_get_tz
import pandas.util._test_decorators as td

import pandas as pd
import pandas._testing as tm


def _compare_with_tz(a, b):
    tm.assert_frame_equal(a, b)

    # compare the zones on each element
    for c in a.columns:
        for i in a.index:
            a_e = a.loc[i, c]
            b_e = b.loc[i, c]
            if not (a_e == b_e and a_e.tz == b_e.tz):
                raise AssertionError(f"invalid tz comparison [{a_e}] [{b_e}]")


# use maybe_get_tz instead of dateutil.tz.gettz to handle the windows
# filename issues.
gettz_dateutil = lambda x: maybe_get_tz("dateutil/" + x)
gettz_pytz = lambda x: x


@pytest.mark.parametrize("gettz", [gettz_dateutil, gettz_pytz])
def test_append_with_timezones(temp_hdfstore, gettz):
    # as columns

    # Single-tzinfo, no DST transition
    df_est = pd.DataFrame(
        {
            "A": [
                pd.Timestamp("20130102 2:00:00", tz=gettz("US/Eastern")).as_unit("ns")
                + timedelta(hours=1) * i
                for i in range(5)
            ]
        }
    )

    # frame with all columns having same tzinfo, but different sides
    #  of DST transition
    df_crosses_dst = pd.DataFrame(
        {
            "A": pd.Timestamp("20130102", tz=gettz("US/Eastern")).as_unit("ns"),
            "B": pd.Timestamp("20130603", tz=gettz("US/Eastern")).as_unit("ns"),
        },
        index=range(5),
    )

    df_mixed_tz = pd.DataFrame(
        {
            "A": pd.Timestamp("20130102", tz=gettz("US/Eastern")).as_unit("ns"),
            "B": pd.Timestamp("20130102", tz=gettz("EET")).as_unit("ns"),
        },
        index=range(5),
    )

    df_different_tz = pd.DataFrame(
        {
            "A": pd.Timestamp("20130102", tz=gettz("US/Eastern")).as_unit("ns"),
            "B": pd.Timestamp("20130102", tz=gettz("CET")).as_unit("ns"),
        },
        index=range(5),
    )

    temp_hdfstore.append("df_tz", df_est, data_columns=["A"])
    result = temp_hdfstore["df_tz"]
    _compare_with_tz(result, df_est)
    tm.assert_frame_equal(result, df_est)

    # select with tz aware
    expected = df_est[df_est.A >= df_est.A[3]]
    result = temp_hdfstore.select("df_tz", where="A>=df_est.A[3]")
    _compare_with_tz(result, expected)

    # ensure we include dates in DST and STD time here.
    temp_hdfstore.remove("df_tz")
    temp_hdfstore.append("df_tz", df_crosses_dst)
    result = temp_hdfstore["df_tz"]
    _compare_with_tz(result, df_crosses_dst)
    tm.assert_frame_equal(result, df_crosses_dst)

    msg = (
        r"invalid info for \[values_block_1\] for \[tz\], "
        r"existing_value \[(dateutil/.*)?(US/Eastern|America/New_York)\] "
        r"conflicts with new value \[(dateutil/.*)?EET\]"
    )
    with pytest.raises(ValueError, match=msg):
        temp_hdfstore.append("df_tz", df_mixed_tz)

    # this is ok
    temp_hdfstore.remove("df_tz")
    temp_hdfstore.append("df_tz", df_mixed_tz, data_columns=["A", "B"])
    result = temp_hdfstore["df_tz"]
    _compare_with_tz(result, df_mixed_tz)
    tm.assert_frame_equal(result, df_mixed_tz)

    # can't append with diff timezone
    msg = (
        r"invalid info for \[B\] for \[tz\], "
        r"existing_value \[(dateutil/.*)?EET\] "
        r"conflicts with new value \[(dateutil/.*)?CET\]"
    )
    with pytest.raises(ValueError, match=msg):
        temp_hdfstore.append("df_tz", df_different_tz)


@pytest.mark.parametrize("gettz", [gettz_dateutil, gettz_pytz])
def test_append_with_timezones_as_index(temp_hdfstore, gettz):
    # GH#4098 example

    dti = pd.date_range("2000-1-1", periods=3, freq="h", tz=gettz("US/Eastern"))
    dti = dti._with_freq(None)  # freq doesn't round-trip

    df = pd.DataFrame({"A": pd.Series(range(3), index=dti)})

    temp_hdfstore.put("df", df, track_times=False)
    result = temp_hdfstore.select("df")
    tm.assert_frame_equal(result, df)

    temp_hdfstore.remove("df")
    temp_hdfstore.append("df", df)
    result = temp_hdfstore.select("df")
    tm.assert_frame_equal(result, df)


@pytest.mark.parametrize("as_index", [True, False])
@pytest.mark.parametrize(
    "stored_tz, appended_tz", [("US/Pacific", None), (None, "US/Pacific")]
)
def test_append_tz_naive_aware_mismatch_raises(
    temp_hdfstore, as_index, stored_tz, appended_tz
):
    # GH#68583 - the append used to be accepted, and the rows already on disk
    # read back as different instants: one stored tz governs the whole column.
    def frame(start, tz):
        dti = pd.date_range(start, periods=2, tz=tz)._with_freq(None)
        if as_index:
            return pd.DataFrame({"v": [1.0, 2.0]}, index=dti)
        return pd.DataFrame({"v": dti})

    stored = frame("2000-01-01", stored_tz)
    temp_hdfstore.append("df", stored)

    name = "index" if as_index else "values_block_0"
    if appended_tz is None:
        msg = f"cannot append tz-naive data to tz-aware col [{name}]"
    else:
        msg = f"cannot append tz-aware data to tz-naive col [{name}]"
    with pytest.raises(TypeError, match=re.escape(msg)):
        temp_hdfstore.append("df", frame("2000-01-03", appended_tz))

    tm.assert_frame_equal(temp_hdfstore["df"], stored)


def test_roundtrip_tz_aware_index(temp_hdfstore, unit):
    # GH 17618
    ts = pd.Timestamp("2000-01-01 01:00:00", tz="US/Eastern")
    dti = pd.DatetimeIndex([ts]).as_unit(unit)
    df = pd.DataFrame(data=[0], index=dti)

    temp_hdfstore.put("frame", df, format="fixed", track_times=False)
    recons = temp_hdfstore["frame"]
    tm.assert_frame_equal(recons, df)

    value = recons.index[0]._value
    denom = {"ns": 1, "us": 1000, "ms": 10**6, "s": 10**9}[unit]
    assert value == 946706400000000000 // denom


def test_store_index_name_with_tz(temp_hdfstore):
    # GH 13884
    df = pd.DataFrame({"A": [1, 2]})
    df.index = pd.DatetimeIndex([1234567890123456787, 1234567890123456788])
    df.index = df.index.tz_localize("UTC")
    df.index.name = "foo"

    temp_hdfstore.put("frame", df, format="table", track_times=False)
    recons = temp_hdfstore["frame"]
    tm.assert_frame_equal(recons, df)


def test_tseries_select_index_column(temp_hdfstore):
    # GH7777
    # selecting a UTC datetimeindex column did
    # not preserve UTC tzinfo set before storing

    # check that no tz still works
    rng = pd.date_range("1/1/2000", "1/30/2000")
    frame = pd.DataFrame(
        np.random.default_rng(2).standard_normal((len(rng), 4)), index=rng
    )

    temp_hdfstore.append("frame", frame)
    result = temp_hdfstore.select_column("frame", "index")
    assert rng.tz == pd.DatetimeIndex(result.values).tz

    # check utc
    rng = pd.date_range("1/1/2000", "1/30/2000", tz="UTC")
    frame = pd.DataFrame(
        np.random.default_rng(2).standard_normal((len(rng), 4)), index=rng
    )

    temp_hdfstore.remove("frame")
    temp_hdfstore.append("frame", frame)
    result = temp_hdfstore.select_column("frame", "index")
    assert rng.tz == result.dt.tz

    # double check non-utc
    rng = pd.date_range("1/1/2000", "1/30/2000", tz="US/Eastern")
    frame = pd.DataFrame(
        np.random.default_rng(2).standard_normal((len(rng), 4)), index=rng
    )

    temp_hdfstore.remove("frame")
    temp_hdfstore.append("frame", frame)
    result = temp_hdfstore.select_column("frame", "index")
    assert rng.tz == result.dt.tz


def test_timezones_fixed_format_frame_non_empty(temp_hdfstore):
    # index
    rng = pd.date_range("1/1/2000", "1/30/2000", tz="US/Eastern")
    rng = rng._with_freq(None)  # freq doesn't round-trip
    df = pd.DataFrame(
        np.random.default_rng(2).standard_normal((len(rng), 4)), index=rng
    )
    temp_hdfstore["df"] = df
    result = temp_hdfstore["df"]
    tm.assert_frame_equal(result, df)


def test_timezones_fixed_format_frame_non_empty_as_data(temp_hdfstore):
    # GH11411
    rng = pd.date_range("1/1/2000", "1/30/2000", tz="US/Eastern")
    rng = rng._with_freq(None)  # freq doesn't round-trip
    df = pd.DataFrame(
        {
            "A": rng,
            "B": rng.tz_convert("UTC").tz_localize(None),
            "C": rng.tz_convert("CET"),
            "D": range(len(rng)),
        },
        index=rng,
    )
    temp_hdfstore["df"] = df
    result = temp_hdfstore["df"]
    tm.assert_frame_equal(result, df)


def test_timezones_fixed_format_empty(temp_hdfstore, tz_aware_fixture, frame_or_series):
    # GH 20594

    dtype = pd.DatetimeTZDtype(tz=tz_aware_fixture)

    obj = pd.Series(dtype=dtype, name="A", index=pd.Index([]))
    if frame_or_series is pd.DataFrame:
        obj = obj.to_frame()

    temp_hdfstore["obj"] = obj
    result = temp_hdfstore["obj"]
    tm.assert_equal(result, obj)


def test_timezones_fixed_format_series_nonempty(temp_hdfstore, tz_aware_fixture):
    # GH 20594

    dtype = pd.DatetimeTZDtype(tz=tz_aware_fixture)

    s = pd.Series([0], dtype=dtype)
    temp_hdfstore["s"] = s
    result = temp_hdfstore["s"]
    tm.assert_series_equal(result, s)


def test_fixed_offset_tz(temp_hdfstore):
    rng = pd.date_range("1/1/2000 00:00:00-07:00", "1/30/2000 00:00:00-07:00")
    frame = pd.DataFrame(
        np.random.default_rng(2).standard_normal((len(rng), 4)), index=rng
    )

    temp_hdfstore["frame"] = frame
    recons = temp_hdfstore["frame"]
    tm.assert_index_equal(recons.index, rng, check_freq=False)
    assert rng.tz == recons.index.tz


@td.skip_if_windows
def test_store_timezone(temp_hdfstore):
    # GH2852
    # issue storing datetime.date with a timezone as it resets when read
    # back in a new timezone

    # original method
    today = date(2013, 9, 10)
    df = pd.DataFrame([1, 2, 3], index=[today, today, today])
    temp_hdfstore["obj1"] = df
    result = temp_hdfstore["obj1"]
    tm.assert_frame_equal(result, df)

    # with tz setting
    with tm.set_timezone("EST5EDT"):
        today = date(2013, 9, 10)
        df = pd.DataFrame([1, 2, 3], index=[today, today, today])
        temp_hdfstore["obj2"] = df

    with tm.set_timezone("CST6CDT"):
        result = temp_hdfstore["obj2"]

    tm.assert_frame_equal(result, df)


def test_dst_transitions(temp_hdfstore):
    # make sure we are not failing on transitions
    times = pd.date_range(
        "2013-10-26 23:00",
        "2013-10-27 01:00",
        tz="Europe/London",
        freq="h",
        ambiguous="infer",
    )
    times = times._with_freq(None)  # freq doesn't round-trip

    for i in [times, times + pd.Timedelta("10min")]:
        df = pd.DataFrame({"A": range(len(i)), "B": i}, index=i)
        temp_hdfstore.append("df", df)
        result = temp_hdfstore.select("df")
        tm.assert_frame_equal(result, df)
        temp_hdfstore.remove("df")


def test_read_with_where_tz_aware_index(temp_hdfstore):
    # GH 11926
    periods = 10
    dts = pd.date_range("20151201", periods=periods, freq="D", tz="UTC", unit="ns")
    mi = pd.MultiIndex.from_arrays([dts, range(periods)], names=["DATE", "NO"])
    expected = pd.DataFrame({"MYCOL": 0}, index=mi)

    key = "mykey"
    with pd.HDFStore(temp_hdfstore) as store:
        store.append(key, expected, format="table", append=True)
    result = pd.read_hdf(temp_hdfstore, key, where="DATE > 20151130")
    tm.assert_frame_equal(result, expected, check_freq=False)