# Copyright 2017-2024 Lawrence Livermore National Security, LLC and other
# Hatchet Project Developers. See the top-level LICENSE file for details.
#
# SPDX-License-Identifier: MIT

from hatchet import GraphFrame


def test_graphframe(mock_graph_literal):
    """Sanity test a GraphFrame object with known data."""
    gf = GraphFrame.from_literal(mock_graph_literal)

    assert len(gf.dataframe) == 24


def test_graphframe_to_literal(mock_graph_literal):
    gf = GraphFrame.from_literal(mock_graph_literal)
    graph_literal = gf.to_literal()

    test_literal_output = gf.from_literal(graph_literal)

    assert len(test_literal_output.graph.roots) == len(gf.graph.roots)
    assert len(test_literal_output.graph) == len(gf.graph)


def test_with_duplicates(mock_graph_literal_duplicates):
    gf = GraphFrame.from_literal(mock_graph_literal_duplicates)

    assert len(gf.graph) == 6

    graph_literal = gf.to_literal()
    assert mock_graph_literal_duplicates.sort() == graph_literal.sort()


def test_with_duplicate_in_first_node(mock_graph_literal_duplicate_first):
    gf = GraphFrame.from_literal(mock_graph_literal_duplicate_first)

    assert len(gf.graph) == 6

    graph_literal = gf.to_literal()
    assert mock_graph_literal_duplicate_first.sort() == graph_literal.sort()


def test_inclusive_time_calculation_mock_graph(mock_graph_literal):
    """Validate calculate_inclusive_metrics() on known dataset."""
    gf = GraphFrame.from_literal(mock_graph_literal)

    # save original time (inc) column for correctness check
    gf.dataframe["orig_inc_time"] = gf.dataframe["time (inc)"]

    # remove original time (inc) column since it will be generated by calculate_inclusive_metrics()
    del gf.dataframe["time (inc)"]

    gf.calculate_inclusive_metrics()
    assert all(
        gf.dataframe["time (inc)"].values == gf.dataframe["orig_inc_time"].values
    )


def test_inclusive_time_calculation_small_mock_graphs(
    small_mock1, small_mock2, small_mock3
):
    gf1 = GraphFrame.from_literal(small_mock1)
    gf1.dataframe["orig_inc_time"] = gf1.dataframe["time (inc)"]
    del gf1.dataframe["time (inc)"]
    gf1.calculate_inclusive_metrics()
    assert all(
        gf1.dataframe["time (inc)"].values == gf1.dataframe["orig_inc_time"].values
    )

    gf2 = GraphFrame.from_literal(small_mock2)
    gf2.dataframe["orig_inc_time"] = gf2.dataframe["time (inc)"]
    del gf2.dataframe["time (inc)"]
    gf2.calculate_inclusive_metrics()
    assert all(
        gf2.dataframe["time (inc)"].values == gf2.dataframe["orig_inc_time"].values
    )

    gf3 = GraphFrame.from_literal(small_mock3)
    gf3.dataframe["orig_inc_time"] = gf3.dataframe["time (inc)"]
    del gf3.dataframe["time (inc)"]
    gf3.calculate_inclusive_metrics()
    assert all(
        gf3.dataframe["time (inc)"].values == gf3.dataframe["orig_inc_time"].values
    )


def test_inclusive_time_calculation_mock_dag_graphs(
    mock_dag_literal1,
    mock_dag_literal2,
):
    gf4 = GraphFrame.from_literal(mock_dag_literal1)
    gf4.dataframe["orig_inc_time"] = gf4.dataframe["time (inc)"]
    del gf4.dataframe["time (inc)"]
    gf4.calculate_inclusive_metrics()
    assert all(
        gf4.dataframe["time (inc)"].values == gf4.dataframe["orig_inc_time"].values
    )

    gf5 = GraphFrame.from_literal(mock_dag_literal2)
    gf5.dataframe["orig_inc_time"] = gf5.dataframe["time (inc)"]
    del gf5.dataframe["time (inc)"]
    gf5.calculate_inclusive_metrics()
    assert all(
        gf5.dataframe["time (inc)"].values == gf5.dataframe["orig_inc_time"].values
    )


def test_inclusive_time_calculation_mock_dag_modules(
    mock_dag_literal_module,
    mock_dag_literal_module_complex,
    mock_dag_literal_module_more_complex,
):
    gf6 = GraphFrame.from_literal(mock_dag_literal_module)
    gf6.dataframe["orig_inc_time"] = gf6.dataframe["time (inc)"]
    del gf6.dataframe["time (inc)"]
    gf6.calculate_inclusive_metrics()
    assert all(
        gf6.dataframe["time (inc)"].values == gf6.dataframe["orig_inc_time"].values
    )

    gf7 = GraphFrame.from_literal(mock_dag_literal_module_complex)
    gf7.dataframe["orig_inc_time"] = gf7.dataframe["time (inc)"]
    del gf7.dataframe["time (inc)"]
    gf7.calculate_inclusive_metrics()
    assert all(
        gf7.dataframe["time (inc)"].values == gf7.dataframe["orig_inc_time"].values
    )

    gf8 = GraphFrame.from_literal(mock_dag_literal_module_more_complex)
    gf8.dataframe["orig_inc_time"] = gf8.dataframe["time (inc)"]
    del gf8.dataframe["time (inc)"]
    gf8.calculate_inclusive_metrics()
    assert all(
        gf8.dataframe["time (inc)"].values == gf8.dataframe["orig_inc_time"].values
    )
