### This is example of the script that will be run in the test environment.

### You can change the rest of the code to define and test your solution.
### However, you should not change the signature of the provided function.
### The script saves "submission.parquet" file in the current directory.
### You can use any additional files and subdirectories to organize your code.


import gc

import pandas as pd
from tqdm import tqdm

from dataset import load_hoho_dataset
from process_sample import process_sample

# check the https://github.com/s23dr/hoho2025/blob/main/hoho2025/example_solutions.py for the example solution

if __name__ == "__main__":
    dataset = load_hoho_dataset(testing=True)

    solution = []

    num_sample = 0
    for subset_name in dataset.keys():
        print(f"Predicting {subset_name}")
        for i, sample in enumerate(tqdm(dataset[subset_name])):
            res = process_sample(sample)
            solution.append(res)

            num_sample += 1
            if num_sample % 10 == 0:
                gc.collect()

    print("------------ Saving results ---------------")
    sub = pd.DataFrame(solution, columns=["order_id", "wf_vertices", "wf_edges"])
    sub.to_parquet("submission.parquet")
    print("------------ Done ------------ ")