ABS_DATASET_DOMAIN = "https://dl.fbaipublicfiles.com"
ABS_DATASET_PATH = f"{ABS_DATASET_DOMAIN}/omnisealbench/"


def get_dataset_config(dataset_name):
    if dataset_name == "voxpopuli_1k/audio":
        return {
            "first_cols": [
                "snr",
                "sisnr",
                "stoi",
                "pesq",
            ],
            "attack_scores": [
                "bit_acc",
                "log10_p_value",
                "TPR",
                "FPR",
            ],
            "categories": {
                "speed": "Time",
                "updownresample": "Time",
                "echo": "Time",
                "random_noise": "Amplitude",
                "lowpass_filter": "Amplitude",
                "highpass_filter": "Amplitude",
                "bandpass_filter": "Amplitude",
                "smooth": "Amplitude",
                "boost_audio": "Amplitude",
                "duck_audio": "Amplitude",
                "shush": "Amplitude",
                "pink_noise": "Amplitude",
                "aac_compression": "Compression",
                "mp3_compression": "Compression",
            },
            "attacks_with_variations": [
                "random_noise",
                "lowpass_filter",
                "highpass_filter",
                "boost_audio",
                "duck_audio",
                "shush",
            ],
        }
    elif dataset_name == "ravdess_1k/audio":
        return {
            "first_cols": ["snr", "sisnr", "stoi", "pesq"],
            "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
            "categories": {
                "speed": "Time",
                "updownresample": "Time",
                "echo": "Time",
                "random_noise": "Amplitude",
                "lowpass_filter": "Amplitude",
                "highpass_filter": "Amplitude",
                "bandpass_filter": "Amplitude",
                "smooth": "Amplitude",
                "boost_audio": "Amplitude",
                "duck_audio": "Amplitude",
                "shush": "Amplitude",
                "pink_noise": "Amplitude",
                "aac_compression": "Compression",
                "mp3_compression": "Compression",
            },
            "attacks_with_variations": [
                "random_noise",
                "lowpass_filter",
                "highpass_filter",
                "boost_audio",
                "duck_audio",
                "shush",
            ],
        }
    elif dataset_name == "val2014_1k/image":
        return {
            "first_cols": ["psnr", "ssim", "lpips", "decoder_time"],
            "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
            "categories": {
                "proportion": "Geometric",
                "collage": "Inpainting",
                "center_crop": "Geometric",
                "rotate": "Geometric",
                "jpeg": "Compression",
                "brightness": "Visual",
                "contrast": "Visual",
                "saturation": "Visual",
                "sharpness": "Visual",
                "resize": "Geometric",
                "overlay_text": "Inpainting",
                "hflip": "Geometric",
                "perspective": "Geometric",
                "median_filter": "Visual",
                "hue": "Visual",
                "gaussian_blur": "Visual",
                "comb": "Mixed",
                "avg": "Averages",
                "none": "Baseline",
            },
            "attacks_with_variations": [
                "center_crop",
                "jpeg",
                "brightness",
                "contrast",
                "saturation",
                "sharpness",
                "resize",
                "perspective",
                "median_filter",
                "hue",
                "gaussian_blur",
            ],
        }
    elif dataset_name == "sav_val_full/video":
        return {
            "first_cols": ["psnr", "ssim", "msssim", "lpips", "vmaf", "decoder_time"],
            "attack_scores": ["bit_acc", "log10_p_value", "TPR", "FPR"],
            "categories": {
                "HorizontalFlip": "Geometric",
                "Rotate": "Geometric",
                "Resize": "Geometric",
                "Crop": "Geometric",
                "Perspective": "Geometric",
                "Brightness": "Visual",
                "Contrast": "Visual",
                "Saturation": "Visual",
                "Grayscale": "Visual",
                "Hue": "Visual",
                "JPEG": "Compression",
                "GaussianBlur": "Visual",
                "MedianFilter": "Visual",
                "H264": "Compression",
                "H264rgb": "Compression",
                "H265": "Compression",
                "VP9": "Compression",
                "H264_Crop_Brightness0": "Mixed",
                "H264_Crop_Brightness1": "Mixed",
                "H264_Crop_Brightness2": "Mixed",
                "H264_Crop_Brightness3": "Mixed",
            },
            "attacks_with_variations": [
                "Rotate",
                "Resize",
                "Crop",
                "Brightness",
                "Contrast",
                "Saturation",
                "H264",
                "H264rgb",
                "H265",
            ],
        }
    else:
        raise ValueError(f"Unknown dataset: {dataset_name}")