# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os

import numpy as np


def fix_lyft(root_folder='./data/lyft', version='v1.01'):
    # refer to https://www.kaggle.com/c/3d-object-detection-for-autonomous-vehicles/discussion/110000  # noqa
    lidar_path = 'lidar/host-a011_lidar1_1233090652702363606.bin'
    root_folder = os.path.join(root_folder, f'{version}-train')
    lidar_path = os.path.join(root_folder, lidar_path)
    assert os.path.isfile(lidar_path), f'Please download the complete Lyft ' \
        f'dataset and make sure {lidar_path} is present.'
    points = np.fromfile(lidar_path, dtype=np.float32, count=-1)
    try:
        points.reshape([-1, 5])
        print(f'This fix is not required for version {version}.')
    except ValueError:
        new_points = np.array(list(points) + [100.0, 1.0], dtype='float32')
        new_points.tofile(lidar_path)
        print(f'Appended 100.0 and 1.0 to the end of {lidar_path}.')


parser = argparse.ArgumentParser(description='Lyft dataset fixer arg parser')
parser.add_argument(
    '--root-folder',
    type=str,
    default='./data/lyft',
    help='specify the root path of Lyft dataset')
parser.add_argument(
    '--version',
    type=str,
    default='v1.01',
    help='specify Lyft dataset version')
args = parser.parse_args()

if __name__ == '__main__':
    fix_lyft(root_folder=args.root_folder, version=args.version)