from numpy import array
from pandas import DataFrame, MultiIndex, Series, concat

from acteval._aggregation import (
    descriptions_to_domain_level,
    descriptions_to_group_level,
    distances_to_domain_level,
    distances_to_group_level,
)
from acteval._pipeline import describe
from acteval.features._utils import equals
from acteval.features.structural import (
    feasibility_eval,
    time_consistency,
)
from acteval.population import Population


def test_time_consistency():
    population = DataFrame(
        [
            {"pid": 0, "start": 0, "end": 10, "duration": 10},
            {"pid": 0, "start": 10, "end": 20, "duration": 10},
            {"pid": 0, "start": 20, "end": 30, "duration": 10},
            {"pid": 1, "start": 0, "end": 10, "duration": 10},
            {"pid": 1, "start": 10, "end": 20, "duration": 10},
        ]
    )
    result = time_consistency(Population(population), target=30).aggregate()
    assert set(result.keys()) == {"starts at 0", "ends at 30", "duration is 30"}
    # Both persons start at 0 → only value 1, count 2
    assert equals(
        result,
        {
            "starts at 0": (array([1.0]), array([2])),
            "ends at 30": (array([0.0, 1.0]), array([1, 1])),
            "duration is 30": (array([0.0, 1.0]), array([1, 1])),
        },
    )


def test_feasibility_eval():
    schedule = DataFrame(
        [
            {"pid": 0, "act": "home"},
            {"pid": 0, "act": "work"},
            {"pid": 0, "act": "home"},
            {"pid": 1, "act": "home"},
            {"pid": 1, "act": "home"},
            {"pid": 2, "act": "home"},
            {"pid": 2, "act": "work"},
            {"pid": 2, "act": "shop"},
        ]
    )
    weights, metrics = feasibility_eval(Population(schedule), "observed")
    assert (
        weights.reset_index(drop=True)
        .astype("int32")
        .equals(Series([3, 3, 3, 3, 3, 3, 3, 3], dtype="int32"))
    )
    assert metrics.reset_index(drop=True).equals(
        Series([2 / 3, 1 / 3, 0, 1 / 3, 1 / 3, 1 / 3, 0, 0])
    )


def test_describe_structural():
    index = MultiIndex.from_tuples(
        [
            ("feasibility", "invalid", "all"),
            ("feasibility", "not home based", "all"),
            ("feasibility", "not home based", "starts"),
            ("feasibility", "not home based", "ends"),
            ("feasibility", "consecutive", "all"),
            ("feasibility", "consecutive", "home"),
            ("feasibility", "consecutive", "work"),
            ("feasibility", "consecutive", "education"),
        ],
        names=["domain", "feature", "segment"],
    )

    observed_weights = Series(
        [3, 3, 3, 3, 3, 3, 3, 3], index=index, name="observed__weight"
    )
    observed_metrics = Series(
        [2 / 3, 1 / 3, 0, 1 / 3, 1 / 3, 1 / 3, 0, 0],
        index=index,
        name="observed",
    )
    weights = Series([3, 3, 3, 3, 3, 3, 3, 3], index=index, name="y__weight")
    metrics = Series(
        [2 / 3, 1 / 3, 0, 1 / 3, 1 / 3, 1 / 3, 0, 0], index=index, name="y"
    )
    metrics = concat([observed_weights, observed_metrics, weights, metrics], axis=1)
    metrics["unit"] = "prob. invalid"
    frames = describe(metrics, metrics)
    assert len(frames["descriptions"]) == 8
    assert len(frames["group_descriptions"]) == 3
    assert len(frames["domain_descriptions"]) == 1

    assert len(frames["distances"]) == 8
    assert len(frames["group_distances"]) == 3
    assert len(frames["domain_distances"]) == 1


def test_describe_splits_structural():
    index = MultiIndex.from_tuples(
        [
            ("feasibility", "invalid", "all", "a"),
            ("feasibility", "not home based", "all", "a"),
            ("feasibility", "not home based", "starts", "a"),
            ("feasibility", "not home based", "ends", "a"),
            ("feasibility", "consecutive", "all", "a"),
            ("feasibility", "consecutive", "home", "a"),
            ("feasibility", "consecutive", "work", "a"),
            ("feasibility", "consecutive", "education", "a"),
            ("feasibility", "invalid", "all", "b"),
            ("feasibility", "not home based", "all", "b"),
            ("feasibility", "not home based", "starts", "b"),
            ("feasibility", "not home based", "ends", "b"),
            ("feasibility", "consecutive", "all", "b"),
            ("feasibility", "consecutive", "home", "b"),
            ("feasibility", "consecutive", "work", "b"),
            ("feasibility", "consecutive", "education", "b"),
        ],
        names=["domain", "feature", "segment", "label"],
    )

    observed_weights = Series([3] * 16, index=index, name="observed__weight")
    observed_metrics = Series([1 / 3] * 16, index=index, name="observed")
    weights = Series([3] * 16, index=index, name="y__weight")
    metrics = Series([1 / 3, 0] * 8, index=index, name="y")
    metrics = concat([observed_weights, observed_metrics, weights, metrics], axis=1)
    metrics["unit"] = "prob. invalid"
    frames = describe(metrics, metrics)
    assert len(frames["descriptions"]) == 8
    assert len(frames["group_descriptions"]) == 3
    assert len(frames["domain_descriptions"]) == 1

    assert len(frames["distances"]) == 8
    assert len(frames["group_distances"]) == 3
    assert len(frames["domain_distances"]) == 1

    label_group_desc = descriptions_to_group_level(metrics, extra=["label"])
    label_group_dist = distances_to_group_level(metrics, extra=["label"])
    label_domain_desc = descriptions_to_domain_level(label_group_desc, extra=["label"])
    label_domain_dist = distances_to_domain_level(label_group_dist, extra=["label"])
    assert len(metrics) == 16
    assert len(label_group_desc) == 6
    assert len(label_domain_desc) == 2

    assert len(metrics) == 16
    assert len(label_group_dist) == 6
    assert len(label_domain_dist) == 2
