from suntdataset.loader import SUNTLoader
from suntdataset.visualizer import SUNTVisualizer
import pandas as pd
import pytest


def test_load_batch_od_parquet():
    loader = SUNTLoader()
    df = loader.load_batch(dataset_type="od", periods=1)

    assert df is not None
    if not df.empty:
        assert "date_ref" in df.columns


def test_load_batch_gtfs_parquet():
    loader = SUNTLoader()
    dataset_types = loader.get_datasets().keys()

    for dataset_type in dataset_types:
        df = loader.load_batch(dataset_type=dataset_type, periods=1)
        assert df is not None
        if not df.empty:            
            print(f"{dataset_type} - len: {len(df)}")
            assert "date_ref" in df.columns


def test_build_od_graph_from_dataframe():
    pytest.importorskip("networkx")

    df = pd.DataFrame(
        {
            "trip_id": ["t1", "t1", "t1", "t2", "t2"],
            "stop_id": ["A", "B", "C", "A", "C"],
            "pt_sequence": [1, 2, 3, 1, 2],
            "stop_time": [
                "2024-03-01 08:00:00",
                "2024-03-01 08:10:00",
                "2024-03-01 08:20:00",
                "2024-03-01 09:00:00",
                "2024-03-01 09:12:00",
            ],
            "n-boardings": [5, 3, 1, 4, 2],
        }
    )

    viz = SUNTVisualizer(df)
    graph = viz.build_od_graph()

    assert graph.number_of_nodes() == 3
    assert graph.number_of_edges() == 3
    assert graph.has_edge("A", "B")
    assert graph.has_edge("B", "C")
    assert graph.has_edge("A", "C")
