import numpy as np
import pytest

import pandas as pd


@pytest.fixture
def datetime_frame() -> pd.DataFrame:
    """
    Fixture for DataFrame of floats with DatetimeIndex

    Columns are ['A', 'B', 'C', 'D']
    """
    return pd.DataFrame(
        np.random.default_rng(2).standard_normal((10, 4)),
        columns=pd.Index(list("ABCD")),
        index=pd.date_range("2000-01-01", periods=10, freq="B"),
    )


@pytest.fixture
def float_string_frame():
    """
    Fixture for DataFrame of floats and strings with index of unique strings

    Columns are ['A', 'B', 'C', 'D', 'foo'].
    """
    df = pd.DataFrame(
        np.random.default_rng(2).standard_normal((30, 4)),
        index=pd.Index([f"foo_{i}" for i in range(30)], dtype=object),
        columns=pd.Index(list("ABCD")),
    )
    df["foo"] = "bar"
    return df


@pytest.fixture
def mixed_float_frame():
    """
    Fixture for DataFrame of different float types with index of unique strings

    Columns are ['A', 'B', 'C', 'D'].
    """
    df = pd.DataFrame(
        {
            col: np.random.default_rng(2).random(30, dtype=dtype)
            for col, dtype in zip(
                list("ABCD"), ["float32", "float32", "float32", "float64"], strict=True
            )
        },
        index=pd.Index([f"foo_{i}" for i in range(30)], dtype=object),
    )
    # not supported by numpy random
    df["C"] = df["C"].astype("float16")
    return df


@pytest.fixture
def mixed_int_frame():
    """
    Fixture for DataFrame of different int types with index of unique strings

    Columns are ['A', 'B', 'C', 'D'].
    """
    return pd.DataFrame(
        {
            col: np.ones(30, dtype=dtype)
            for col, dtype in zip(
                list("ABCD"), ["int32", "uint64", "uint8", "int64"], strict=True
            )
        },
        index=pd.Index([f"foo_{i}" for i in range(30)], dtype=object),
    )


@pytest.fixture
def timezone_frame():
    """
    Fixture for DataFrame of date_range Series with different time zones

    Columns are ['A', 'B', 'C']; some entries are missing

               A                         B                         C
    0 2013-01-01 2013-01-01 00:00:00-05:00 2013-01-01 00:00:00+01:00
    1 2013-01-02                       NaT                       NaT
    2 2013-01-03 2013-01-03 00:00:00-05:00 2013-01-03 00:00:00+01:00
    """
    df = pd.DataFrame(
        {
            "A": pd.date_range("20130101", periods=3, unit="ns"),
            "B": pd.date_range("20130101", periods=3, tz="US/Eastern", unit="ns"),
            "C": pd.date_range("20130101", periods=3, tz="CET", unit="ns"),
        }
    )
    df.iloc[1, 1] = pd.NaT
    df.iloc[1, 2] = pd.NaT
    return df