import time
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
import numpy as np
from datasets import load_dataset
import sklearn_submission
from sklearn_submission import predict_wireframe_sklearn

sklearn_submission.USE_BUNDLE_ADJUST = True
sklearn_submission.ADD_ISOLATED_TRACK_VERTICES = True

print("Loading dataset...")
dataset = load_dataset('usm3d/hoho22k_2026_trainval', split='train', streaming=True, trust_remote_code=True)

# Process 5 samples and time them
times = []
for idx, sample in enumerate(dataset):
    if idx >= 5:
        break
    start = time.time()
    try:
        predict_wireframe_sklearn(sample)
    except Exception as e:
        print(f"Error on sample {idx}: {e}")
    elapsed = time.time() - start
    times.append(elapsed)
    print(f"Sample {idx}: {elapsed:.2f} seconds")

print(f"Average time per sample: {np.mean(times):.2f} seconds")