"""
Implement test for targetviz report
"""

import tempfile
from time import sleep
from unittest.mock import MagicMock, Mock

import matplotlib
import numpy as np
import pandas as pd
import pytest
from sklearn import datasets

from targetviz.analyzers import BaseAnalyzer, ColumnAnalyzer
from targetviz.config import config
from targetviz.profile_report import targetviz_report

matplotlib.use("Agg")  # Use non-interactive backend for tests


def _has_pyarrow():
    try:
        import pyarrow  # noqa: F401

        return hasattr(pd, "ArrowDtype")
    except ImportError:
        return False


def sklearn_to_df(sklearn_dataset):
    """
    Create pandas DataFrame from sklearn dataset
    :param sklearn_dataset: result of loading datasets with sklearn, for
        example: sklearn.datasets.fetch_california_housing()
    :return: pandas DataFrame with variables and target
    """
    data = pd.DataFrame(sklearn_dataset.data, columns=sklearn_dataset.feature_names)
    data["target"] = pd.Series(sklearn_dataset.target)
    return data


def create_cat_dataset():
    """
    Creates a pandas DataFrame with categorical variables to use in some test
    :return: pandas DataFrame
    """
    np.random.seed(42)
    data = pd.DataFrame(
        {
            "cat_var": np.random.randint(9, size=100),
            "target": np.random.normal(10, size=100),
        }
    )
    data["cat_var"] = data.cat_var.map(
        {
            0: "t0",
            1: "t1",
            2: "t2",
            3: "t3",
            4: "t4",
            5: "t5",
            6: "t6",
            7: "t7",
            8: "t8",
        }
    )

    return data


def create_dt_dataset():
    """
    Create pandas DataFrame with some date values, and both numeric and categorical
        variables
    :return: pandas DataFrame
    """
    np.random.seed(42)
    starting_date = pd.Timestamp(2019, 7, 1)
    data = pd.DataFrame(
        {
            "date": (
                starting_date + pd.to_timedelta(np.random.randint(1, 100, size=100), unit="d")
            ),
            "target": np.random.normal(10, size=100),
        }
    )

    return data


def create_nan_dataset():
    """
    Create pandas DataFrame with some nan values, and both numeric and categorical
        variables
    :return: pandas DataFrame
    """
    np.random.seed(42)
    data = pd.DataFrame(
        {
            "cat_var": np.random.randint(9, size=100),
            "reg_var": np.random.normal(5, size=100),
            "target": np.random.normal(10, size=100),
        }
    )
    data["cat_var"] = data.cat_var.map(
        {
            0: "t0",
            1: "t1",
            2: "t2",
            3: "t3",
            4: "t4",
            5: "t5",
            6: "t6",
            7: "t7",
            8: "t8",
        }
    )
    data.loc[list(np.random.randint(100, size=10)), "cat_var"] = np.nan
    data.loc[list(np.random.randint(100, size=10)), "reg_var"] = np.nan

    return data


def create_df_unique():
    """
    Function gets dataset with only unique values, to test that pd.qcut is not a problem

    :return:
        (pd.DataFrame)
    """
    np.random.seed(42)
    data = pd.DataFrame(
        {
            "target": (np.random.uniform(size=10000) > 0.1) * 1,
            "pred": np.random.choice(
                np.array([-999, 0, 1.1, 3.2]),
                replace=True,
                p=np.array([0.95, 0.02, 0.02, 0.01]),
            ),
        }
    )
    return data


def test_targetviz():
    """
    Test cases for targetviz, asserting that no error arises during runtime
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        df_cal = sklearn_to_df(datasets.fetch_california_housing())
        target_col = "target"
        assert targetviz_report(df_cal, target_col, output_dir=temp_dir + "/") is not None

        data = create_cat_dataset()
        assert targetviz_report(data, "target", output_dir=temp_dir + "/") is not None

        df_nan = create_nan_dataset()
        assert (
            targetviz_report(df_nan, "target", pct_outliers=0.05, output_dir=temp_dir + "/")
            is not None
        )

        df_dt = create_dt_dataset()
        assert (
            targetviz_report(df_dt, "target", pct_outliers=0.05, output_dir=temp_dir + "/")
            is not None
        )

        df_unique = create_df_unique()
        assert targetviz_report(df_unique, "target", output_dir=temp_dir + "/") is not None


def test_calc_explained_var():
    """
    Test some cases for explained variance
    """
    df_cal = sklearn_to_df(datasets.fetch_california_housing())
    target_col = "target"
    n_breaks = 5

    # Test 1
    col = "AveRooms"
    new_col = pd.qcut(df_cal[col], n_breaks, duplicates="drop").cat.remove_unused_categories()

    config_ = config
    config_.target_type = "NUM"
    col_analysis = ColumnAnalyzer(Mock, target_col, config_, Mock)
    col_analysis.rate_non_nulls = 1
    np.testing.assert_almost_equal(
        col_analysis.calc_explained_variance(df_cal, new_col), 0.1240, decimal=3
    )

    # Test 2
    col_analysis = ColumnAnalyzer(Mock, target_col, config_, Mock)
    col_analysis.rate_non_nulls = 1
    col = "HouseAge"
    new_col = pd.qcut(df_cal[col], n_breaks, duplicates="drop").cat.remove_unused_categories()
    np.testing.assert_almost_equal(
        col_analysis.calc_explained_variance(df_cal, new_col), 0.0103, decimal=3
    )

    # Test 3
    col_analysis.rate_non_nulls = 0.5
    np.testing.assert_almost_equal(
        col_analysis.calc_explained_variance(df_cal, new_col), 0.00519, decimal=3
    )


def test_get_desc():
    """
    Test get_desc function, that is in charge of returning statistic description of
        variables
    """
    np.random.seed(42)
    quantiles = [0.01, 0.05, 0.25, 0.75, 0.95, 0.99]

    series = pd.Series(np.random.uniform(0, 1, size=20))
    expected_result = [
        20,
        20,
        0,
        "0.00%",
        "0.97",
        "0.02",
        "0.46",
        "0.40",
        "0.02",
        "0.31",
        "0.03",
        "0.06",
        "0.18",
        "0.71",
        "0.95",
        "0.97",
    ]

    base = BaseAnalyzer(Mock, Mock, MagicMock(), Mock)
    result = base.get_desc(series, full_samp=True, quantiles=quantiles)
    assert expected_result == result

    series[list(np.random.randint(20, size=3))] = np.nan
    expected_result2 = [
        20,
        17,
        3,
        "0.15%",
        "0.97",
        "0.02",
        "0.48",
        "0.43",
        "0.02",
        "0.31",
        "0.03",
        "0.05",
        "0.21",
        "0.73",
        "0.95",
        "0.97",
    ]
    assert expected_result2 == base.get_desc(series, full_samp=True, quantiles=quantiles)
    expected_result3 = [
        20,
        17,
        "-",
        "-",
        "0.97",
        "0.02",
        "0.48",
        "0.43",
        "0.02",
        "0.31",
        "0.03",
        "0.05",
        "0.21",
        "0.73",
        "0.95",
        "0.97",
    ]
    assert expected_result3 == base.get_desc(series, full_samp=False, quantiles=quantiles)


def test_numeric_name():
    with tempfile.TemporaryDirectory() as temp_dir:
        sleep(1)  # make sure folder names change since last test
        df_names_num = pd.DataFrame({"pred": [0, 1, 2, 3], "0": [2, 3, 4, 5], 0: [7, 8, 9, 0]})
        targetviz_report(df_names_num, "pred", n_breaks=2, output_dir=temp_dir + "/")


def test_nullable_types():
    with tempfile.TemporaryDirectory() as temp_dir:
        # Test with Int (nullable int)
        df_int = pd.DataFrame(
            {"target": [1, 2, 3, 4, 5], "int_col": pd.array([1, 2, 3, pd.NA, 5], dtype="Int64")}
        )
        assert targetviz_report(df_int, "target", output_dir=temp_dir + "/") is not None

        # Test with Float (nullable float)
        df_float = pd.DataFrame(
            {
                "target": [1.0, 2.0, 3.0, 4.0, 5.0],
                "float_col": pd.array([1.1, 2.2, 3.3, pd.NA, 5.5], dtype="Float64"),
            }
        )
        assert targetviz_report(df_float, "target", output_dir=temp_dir + "/") is not None

        # Test with Bool (nullable boolean)
        df_bool = pd.DataFrame(
            {
                "target": [True, False, True, False, True],
                "bool_col": pd.array([True, False, True, pd.NA, False], dtype="boolean"),
            }
        )
        assert targetviz_report(df_bool, "target", output_dir=temp_dir + "/") is not None

        # Test with mixed types
        df_mixed = pd.DataFrame(
            {
                "target": [1, 2, 3, 4, 5],
                "int_col": pd.array([1, 2, 3, pd.NA, 5], dtype="Int64"),
                "float_col": pd.array([1.1, 2.2, 3.3, pd.NA, 5.5], dtype="Float64"),
                "bool_col": pd.array([True, False, True, pd.NA, False], dtype="boolean"),
            }
        )
        assert targetviz_report(df_mixed, "target", output_dir=temp_dir + "/") is not None


def test_uint_types():
    """Test with unsigned integer types (numpy uint and nullable UInt)."""
    np.random.seed(42)
    n = 50

    with tempfile.TemporaryDirectory() as temp_dir:
        # numpy uint types
        df = pd.DataFrame(
            {
                "target": np.random.normal(10, size=n),
                "uint8_col": np.random.randint(0, 100, size=n).astype(np.uint8),
                "uint16_col": np.random.randint(0, 1000, size=n).astype(np.uint16),
                "uint32_col": np.random.randint(0, 10000, size=n).astype(np.uint32),
                "uint64_col": np.random.randint(0, 10000, size=n).astype(np.uint64),
            }
        )
        assert targetviz_report(df, "target", output_dir=temp_dir + "/") is not None

        # Nullable UInt types (pandas extension types)
        df_nullable = pd.DataFrame(
            {
                "target": list(range(1, 11)),
                "uint8_col": pd.array([1, 2, 3, pd.NA, 5, 6, 7, 8, 9, 10], dtype="UInt8"),
                "uint16_col": pd.array(
                    [100, 200, 300, pd.NA, 500, 600, 700, 800, 900, 1000], dtype="UInt16"
                ),
                "uint32_col": pd.array(
                    [1000, 2000, 3000, pd.NA, 5000, 6000, 7000, 8000, 9000, 10000], dtype="UInt32"
                ),
                "uint64_col": pd.array(
                    [10000, 20000, 30000, pd.NA, 50000, 60000, 70000, 80000, 90000, 100000],
                    dtype="UInt64",
                ),
            }
        )
        assert targetviz_report(df_nullable, "target", output_dir=temp_dir + "/") is not None


def test_nullable_int_float_all_sizes():
    """Test all sizes of nullable Int and Float types."""
    with tempfile.TemporaryDirectory() as temp_dir:
        df = pd.DataFrame(
            {
                "target": list(range(1, 11)),
                "int8_col": pd.array([1, 2, 3, pd.NA, 5, 6, 7, 8, 9, 10], dtype="Int8"),
                "int16_col": pd.array(
                    [100, 200, 300, pd.NA, 500, 600, 700, 800, 900, 1000], dtype="Int16"
                ),
                "int32_col": pd.array(
                    [1000, 2000, 3000, pd.NA, 5000, 6000, 7000, 8000, 9000, 10000], dtype="Int32"
                ),
                "float32_col": pd.array(
                    [1.1, 2.2, 3.3, pd.NA, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0], dtype="Float32"
                ),
            }
        )
        assert targetviz_report(df, "target", output_dir=temp_dir + "/") is not None


@pytest.mark.skipif(not _has_pyarrow(), reason="pyarrow or ArrowDtype not available")
def test_pyarrow_numeric_types():
    """Test with pyarrow-backed numeric types."""
    import pyarrow as pa

    with tempfile.TemporaryDirectory() as temp_dir:
        df = pd.DataFrame(
            {
                "target": list(range(1, 11)),
                "int_col": pd.array(
                    [1, 2, 3, None, 5, 6, 7, 8, 9, 10], dtype=pd.ArrowDtype(pa.int64())
                ),
                "float_col": pd.array(
                    [1.1, 2.2, 3.3, None, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0],
                    dtype=pd.ArrowDtype(pa.float64()),
                ),
                "uint_col": pd.array(
                    [1, 2, 3, None, 5, 6, 7, 8, 9, 10], dtype=pd.ArrowDtype(pa.uint32())
                ),
            }
        )
        assert targetviz_report(df, "target", output_dir=temp_dir + "/") is not None


@pytest.mark.skipif(not _has_pyarrow(), reason="pyarrow or ArrowDtype not available")
def test_pyarrow_string_type():
    """Test with pyarrow-backed string type."""
    import pyarrow as pa

    np.random.seed(42)
    n = 50
    with tempfile.TemporaryDirectory() as temp_dir:
        df = pd.DataFrame(
            {
                "target": np.random.normal(10, size=n),
                "str_col": pd.array(
                    ["cat_" + str(i % 5) for i in range(n)], dtype=pd.ArrowDtype(pa.string())
                ),
            }
        )
        assert targetviz_report(df, "target", output_dir=temp_dir + "/") is not None


def test_string_dtype():
    """Test with pandas StringDtype."""
    np.random.seed(42)
    n = 50
    with tempfile.TemporaryDirectory() as temp_dir:
        df = pd.DataFrame(
            {
                "target": np.random.normal(10, size=n),
                "str_col": pd.array(
                    ["cat_" + str(i % 5) for i in range(n)], dtype=pd.StringDtype()
                ),
            }
        )
        assert targetviz_report(df, "target", output_dir=temp_dir + "/") is not None


def test_pandas_future_string_inference():
    """Test pandas 3 string inference mode (future.infer_string)."""
    if not hasattr(pd, "options"):
        pytest.skip("pd.options not available")
    if not hasattr(pd.options, "future"):
        pytest.skip("pd.options.future not available")
    if not hasattr(pd.options.future, "infer_string"):
        pytest.skip("future.infer_string not available in this pandas version")

    with pd.option_context("future.infer_string", True):
        np.random.seed(42)
        n = 50
        df = pd.DataFrame(
            {
                "target": np.random.normal(10, size=n),
                "str_col": ["cat_" + str(i % 5) for i in range(n)],
            }
        )
        # Verify the string type was inferred (not object)
        assert df["str_col"].dtype != np.dtype("object")
        with tempfile.TemporaryDirectory() as temp_dir:
            assert targetviz_report(df, "target", output_dir=temp_dir + "/") is not None
