# https://github.com/Mohamedelrefaie/DrivAerNet/blob/main/RegDGCNN_SurfaceFields/run_pipeline.py
import argparse
import os

def preprocess_data(args):
    """
    Preprocess the dataset to create cached point cloud data.

    Args:
        args: Command line arguments

    Returns:
        True if preprocessing was successful, False otherwise
    """

    # Create cache directory if it doesn't exist
    cache_dir = args.cache_dir
    os.makedirs(cache_dir, exist_ok=True)

    try:
        # Import required modules for preprocessing
        from data_loader import SurfacePressureDataset

        # Create the dataset with preprocessing enabled
        dataset = SurfacePressureDataset(
            root_dir=args.data_dir,
            num_points=args.num_points,
            preprocess=True,
            cache_dir=cache_dir
        )
        
        # Process all files
        for i, vtk_file in enumerate(dataset.vtk_files):
            _ = dataset[i]  # This will trigger preprocessing and caching
        return True
    except Exception as e:
        return False


parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', type=str, default='../../../Extra_results/driveaer_raw')
parser.add_argument('--cache_dir', type=str, default='../../../driveaer_raw/processed_data')
# parser.add_argument('--cache_dir', type=str, default='../../../Extra_results/processed_data')
parser.add_argument('--num_points', type=int, default=10000)
args = parser.parse_args()
result = preprocess_data(args)
if result:
    print('Preprocessing complete')
else:
    print('Error in preprocessing')