import albumentations as A

def get_augmentation_compose(augmentation_config):
    augmentations = []
    for augmentation_name, augmentation_params in augmentation_config.items():
        # Get default config and update with provided params
        config = AUGMENTATIONS_CONFIG.get(augmentation_name, {}).copy()
        if augmentation_params:
            config.update(augmentation_params)
            
        augmentations.append(
            AUGMENTATIONS[augmentation_name](**config)
        )
    return A.Compose(
        augmentations,
        bbox_params=A.BboxParams(format="coco", label_fields=["class_labels"]),
    )

# Dictionary mapping augmentation names to their Albumentations class instances
AUGMENTATIONS = {
    "additive_noise": A.AdditiveNoise,
    "advanced_blur": A.AdvancedBlur,
    "affine": A.Affine,
    "at_least_one_bbox_random_crop": A.AtLeastOneBBoxRandomCrop,
    "auto_contrast": A.AutoContrast,
    "bbox_safe_random_crop": A.BBoxSafeRandomCrop,
    "blur": A.Blur,
    "center_crop": A.CenterCrop,
    "channel_dropout": A.ChannelDropout,
    "channel_shuffle": A.ChannelShuffle,
    "chromatic_aberration": A.ChromaticAberration,
    "clahe": A.CLAHE,
    "coarse_dropout": A.CoarseDropout,
    "color_jitter": A.ColorJitter,
    "constrained_coarse_dropout": A.ConstrainedCoarseDropout,
    "crop": A.Crop,
    "crop_and_pad": A.CropAndPad,
    "crop_non_empty_mask_if_exists": A.CropNonEmptyMaskIfExists,
    "d4": A.D4,
    "defocus": A.Defocus,
    "downscale": A.Downscale,
    "elastic_transform": A.ElasticTransform,
    "emboss": A.Emboss,
    "equalize": A.Equalize,
    "erasing": A.Erasing,
    "fancy_pca": A.FancyPCA,
    "fda": A.FDA,
    "frequency_masking": A.FrequencyMasking,
    "from_float": A.FromFloat,
    "gauss_noise": A.GaussNoise,
    "gaussian_blur": A.GaussianBlur,
    "glass_blur": A.GlassBlur,
    "grid_distortion": A.GridDistortion,
    "grid_dropout": A.GridDropout,
    "grid_elastic_deform": A.GridElasticDeform,
    "he_stain": A.HEStain,
    "histogram_matching": A.HistogramMatching,
    "horizontal_flip": A.HorizontalFlip,
    "hue_saturation_value": A.HueSaturationValue,
    "illumination": A.Illumination,
    "image_compression": A.ImageCompression,
    "invert_img": A.InvertImg,
    "iso_noise": A.ISONoise,
    "lambda": A.Lambda,
    "longest_max_size": A.LongestMaxSize,
    "mask_dropout": A.MaskDropout,
    "median_blur": A.MedianBlur,
    "morphological": A.Morphological,
    "motion_blur": A.MotionBlur,
    "multiplicative_noise": A.MultiplicativeNoise,
    "no_op": A.NoOp,
    "normalize": A.Normalize,
    "optical_distortion": A.OpticalDistortion,
    "pad": A.Pad,
    "pad_if_needed": A.PadIfNeeded,
    "perspective": A.Perspective,
    "piecewise_affine": A.PiecewiseAffine,
    "pixel_distribution_adaptation": A.PixelDistributionAdaptation,
    "pixel_dropout": A.PixelDropout,
    "planckian_jitter": A.PlanckianJitter,
    "plasma_brightness_contrast": A.PlasmaBrightnessContrast,
    "plasma_shadow": A.PlasmaShadow,
    "posterize": A.Posterize,
    "random_brightness_contrast": A.RandomBrightnessContrast,
    "random_crop": A.RandomCrop,
    "random_crop_from_borders": A.RandomCropFromBorders,
    "random_crop_near_bbox": A.RandomCropNearBBox,
    "random_fog": A.RandomFog,
    "random_gamma": A.RandomGamma,
    "random_gravel": A.RandomGravel,
    "random_grid_shuffle": A.RandomGridShuffle,
    "random_rain": A.RandomRain,
    "random_resized_crop": A.RandomResizedCrop,
    "random_rotate_90": A.RandomRotate90,
    "random_scale": A.RandomScale,
    "random_shadow": A.RandomShadow,
    "random_sized_bbox_safe_crop": A.RandomSizedBBoxSafeCrop,
    "random_sized_crop": A.RandomSizedCrop,
    "random_snow": A.RandomSnow,
    "random_sun_flare": A.RandomSunFlare,
    "random_tone_curve": A.RandomToneCurve,
    "resize": A.Resize,
    "rgb_shift": A.RGBShift,
    "ringing_overshoot": A.RingingOvershoot,
    "rotate": A.Rotate,
    "safe_rotate": A.SafeRotate,
    "salt_and_pepper": A.SaltAndPepper,
    "sharpen": A.Sharpen,
    "shift_scale_rotate": A.ShiftScaleRotate,
    "shot_noise": A.ShotNoise,
    "smallest_max_size": A.SmallestMaxSize,
    "solarize": A.Solarize,
    "spatter": A.Spatter,
    "superpixels": A.Superpixels,
    "template_transform": A.TemplateTransform,
    "text_image": A.TextImage,
    "thin_plate_spline": A.ThinPlateSpline,
    "time_masking": A.TimeMasking,
    "time_reverse": A.TimeReverse,
    "to_float": A.ToFloat,
    "to_gray": A.ToGray,
    "to_rgb": A.ToRGB,
    "to_sepia": A.ToSepia,
    "transpose": A.Transpose,
    "unsharp_mask": A.UnsharpMask,
    "vertical_flip": A.VerticalFlip,
    "zoom_blur": A.ZoomBlur,
}

# Default configurations for augmentations
AUGMENTATIONS_CONFIG = {
    "planckian_jitter": {"factor": 0.1},
    "plasma_brightness_contrast": {"brightness": 0.2, "contrast": 0.2},
    "plasma_shadow": {"shadow_dimension": 4, "shadow_roi": (0, 0, 1, 1)},
    "posterize": {"num_bits": 4},
    "random_brightness_contrast": {"brightness_limit": 0.2, "contrast_limit": 0.2},
    "random_crop": {"height": 256, "width": 256},
    "random_crop_from_borders": {"crop_left": 0.1, "crop_right": 0.1, "crop_top": 0.1, "crop_bottom": 0.1},
    "random_crop_near_bbox": {"max_part_shift": 0.3},
    "random_fog": {"fog_coef_lower": 0.3, "fog_coef_upper": 1},
    "random_gamma": {"gamma_limit": (80, 120)},
    "random_gravel": {"gravel_roi": (0, 0, 1, 1)},
    "random_grid_shuffle": {"grid": (3, 3)},
    "random_rain": {"drop_length": 20, "drop_width": 1, "drop_color": (200, 200, 200)},
    "random_resized_crop": {"height": 256, "width": 256, "scale": (0.8, 1.0)},
    "random_rotate_90": {"p": 0.5},
    "random_scale": {"scale_limit": 0.1},
    "random_shadow": {"shadow_roi": (0, 0.5, 1, 1), "num_shadows_lower": 1, "num_shadows_upper": 2},
    "random_sized_bbox_safe_crop": {"height": 256, "width": 256, "erosion_rate": 0.0},
    "random_sized_crop": {"min_max_height": (256, 512), "height": 256, "width": 256},
    "random_snow": {"snow_point_lower": 0.1, "snow_point_upper": 0.3, "brightness_coeff": 2.5},
    "random_sun_flare": {"flare_roi": (0, 0, 1, 0.5), "angle_lower": 0, "angle_upper": 1},
    "random_tone_curve": {"scale": 0.1},
    "resize": {"height": 256, "width": 256},
    "rgb_shift": {"r_shift_limit": 20, "g_shift_limit": 20, "b_shift_limit": 20},
    "ringing_overshoot": {"blur_limit": 7},
    "rotate": {"limit": 90},
    "safe_rotate": {"limit": 90},
    "salt_and_pepper": {"p": 0.5},
    "sharpen": {"alpha": (0.2, 0.5), "lightness": (0.5, 1.0)},
    "shift_scale_rotate": {"shift_limit": 0.0625, "scale_limit": 0.1, "rotate_limit": 45},
    "shot_noise": {"gauss": 10.0, "poisson": 10.0},
    "smallest_max_size": {"max_size": 1024},
    "solarize": {"threshold": 128},
    "spatter": {"mean": 0.65, "std": 0.3},
    "superpixels": {"p_replace": 0.1, "n_segments": 100},
    "template_transform": {"template": None},
    "text_image": {"text": "Sample", "font_size": 32},
    "thin_plate_spline": {"points_per_dim": 4},
    "time_masking": {"num_masks_min": 1, "num_masks_max": 5},
    "time_reverse": {"p": 0.5},
    "to_float": {"max_value": 255.0},
    "to_gray": {"p": 0.5},
    "to_rgb": {"p": 0.5},
    "to_sepia": {"p": 0.5},
    "transpose": {"p": 0.5},
    "unsharp_mask": {"blur_limit": (3, 7)},
    "vertical_flip": {"p": 0.5},
    "zoom_blur": {"max_factor": 1.4}
}
