from huggingface_hub import snapshot_download
import argparse

parser = argparse.ArgumentParser()
parser.add_argument(
    "--output-dir", "-o", required=True, help="Output directory for the dataset."
)
parser.add_argument(
    "--scenes", "-s", required=True, nargs="+", help="List of scenes to download.",
    choices=["010_0050", "020_0020", "040_0040", "050_0130", "050_0160", "060_0100", "060_0130", "070_0123"]
)
parser.add_argument(
    "--modalities", "-m", required=True, nargs="*", help="List of scenes to download.",
    choices=["rgb", "depth", "mask", "segmentation", "normal", "flow_fw", "flow_bw"]
)
parser.add_argument(
    "--tasks", "-t", required=True, nargs="+", help="Tasks to download.",
    choices=["Dense", "Sparse", "Mono"]
)

def main(args):
    # Collect scenes from multiple repos
    repo_type = "dataset"

    # Set modalities patterns
    all_modalities = ["rgb", "depth", "mask", "segmentation", "normal", "flow_fw", "flow_bw"]
    modalities_exclude = set(all_modalities) - set(args.modalities)
    exclude_patterns = []
    for mod in modalities_exclude:
        if mod == "rgb":
            exclude_patterns.append("*frame_????.png")
        elif mod == "depth":
            exclude_patterns.append("*_depth.npy")
        elif mod == "mask":
            exclude_patterns.append("*_dyn_mask.png")
        elif mod == "flow_fw":
            exclude_patterns.append("*_flow_fw.npy")
        elif mod == "flow_bw":
            exclude_patterns.append("*_flow_bw.npy")
        elif mod == "segmentation":
            exclude_patterns.append("*_segmentation.png")
        elif mod == "normal":
            exclude_patterns.append("*_normal.png")

    # Set task patterns
    allow_patterns = []
    for task in args.tasks:
        allow_patterns.append(f"*/{task}/*")
    if "segmentation" not in modalities_exclude:
        allow_patterns.append("*segmentation.json")

    # Download from respective repos
    for scene in args.scenes:
        if scene == "050_0130":
            repo_id = "charge-benchmark/Charge-050_0130"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
		elif scene == "010_0050":
            repo_id = "charge-benchmark/Charge-010_0050-Dense"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
			repo_id = "charge-benchmark/Charge-010_0050-Sparse-Mono"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
		elif scene == "020_0020":
            repo_id = "charge-benchmark/Charge-020_0020-Dense"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
			repo_id = "charge-benchmark/Charge-020_0020-Sparse-Mono"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
		elif scene == "040_0040":
            repo_id = "charge-benchmark/Charge-040_0040-DenseTrain"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
			repo_id = "charge-benchmark/Charge-040_0040-DenseTest"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
			repo_id = "charge-benchmark/Charge-040_0040-Sparse-Mono"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
		elif scene == "050_0160":
            repo_id = "charge-benchmark/Charge-050_0160"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
		elif scene == "060_0100":
            repo_id = "charge-benchmark/Charge-060_0100"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
		elif scene == "060_0130":
            repo_id = "charge-benchmark/Charge-060_0130"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
		elif scene == "070_0123":
            repo_id = "charge-benchmark/Charge-070_0123"
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                allow_patterns=allow_patterns,
                ignore_patterns=exclude_patterns,
                local_dir=args.output_dir
            )
        else:
            print(f"Incorrect scene {scene}.")
        

if __name__ == "__main__":
    args = parser.parse_args()
    main(args)