import os
import pickle
import pandas as pd

from src.graph.dataset_builder import build_graph_dataset


def main():
    print("Loading dataset...")
    df = pd.read_csv("data/processed/transactions.csv")
    users = pd.read_csv("data/processed/users.csv")

    print("Building graph dataset...")
    graph_data = build_graph_dataset(df, users)

    os.makedirs("data/graph", exist_ok=True)

    with open("data/graph/graph.pkl", "wb") as f:
        pickle.dump(graph_data, f)

    print("Graph dataset saved")


if __name__ == "__main__":
    main()