import numpy as np

from template_code_location.dataframe_level_anonymisation.utils import (
    parse_value_list,
    normalize_hierarchy_levels,
)


# ------------------------------------
# Tests for parse_value_list
# ------------------------------------
def test_parse_value_list_all_strings_digits():
    values = ["1", "2", "3"]
    assert parse_value_list(values) == [1, 2, 3]


def test_parse_value_list_mixed_values():
    values = ["1", 2, "abc", "5"]
    assert parse_value_list(values) == [1, 2, "abc", 5]


def test_parse_value_list_no_digits():
    values = ["a", "b", "c"]
    assert parse_value_list(values) == ["a", "b", "c"]


# ------------------------------------
# Tests for normalize_hierarchy_levels
# ------------------------------------
def test_normalize_hierarchy_levels_level_0_converted_to_numpy_array():
    hierarchy = {"age": {"0": ["1", "2", "3"], "1": ["0-10", "11-20"]}}

    normalized = normalize_hierarchy_levels(hierarchy)

    assert "age" in normalized
    assert 0 in normalized["age"]
    assert isinstance(normalized["age"][0], np.ndarray)
    assert normalized["age"][0].tolist() == [1, 2, 3]  # converted via parse_value_list
    assert normalized["age"][1] == ["0-10", "11-20"]  # untouched


def test_normalize_hierarchy_levels_multiple_columns():
    hierarchy = {"age": {"0": ["10", "20"]}, "gender": {"0": ["M", "F"], "1": ["*"]}}

    normalized = normalize_hierarchy_levels(hierarchy)

    # First column
    assert isinstance(normalized["age"][0], np.ndarray)
    assert normalized["age"][0].tolist() == [10, 20]

    # Second column
    assert isinstance(normalized["gender"][0], np.ndarray)
    assert normalized["gender"][0].tolist() == ["M", "F"]
    assert normalized["gender"][1] == ["*"]


def test_normalize_hierarchy_levels_mixed_digit_non_digit_at_level_0():
    hierarchy = {"test": {"0": ["1", "x", "3"]}}

    normalized = normalize_hierarchy_levels(hierarchy)

    assert isinstance(normalized["test"][0], np.ndarray)
    assert normalized["test"][0].tolist() == ["1", "x", "3"]


def test_normalize_hierarchy_levels_empty_mapping():
    hierarchy = {"col": {}}
    normalized = normalize_hierarchy_levels(hierarchy)

    assert normalized == {"col": {}}