///|
pub struct AugmentationConfig {
  brightness : Int
  contrast_numerator : Int
  contrast_denominator : Int
  noise_seed : UInt
  noise_magnitude : Int
  blur_radius : Int
  rotate_quarter_turns : Int
  flip_horizontal : Bool
} derive(Debug, Eq)

///|
pub fn AugmentationConfig::identity() -> AugmentationConfig {
  {
    brightness: 0,
    contrast_numerator: 1,
    contrast_denominator: 1,
    noise_seed: 1,
    noise_magnitude: 0,
    blur_radius: 0,
    rotate_quarter_turns: 0,
    flip_horizontal: false,
  }
}

///|
pub fn AugmentationConfig::with_seed(
  seed : UInt,
  brightness? : Int = 0,
  noise_magnitude? : Int = 0,
  blur_radius? : Int = 0,
  rotate_quarter_turns? : Int = 0,
  flip_horizontal? : Bool = false,
) -> AugmentationConfig {
  {
    brightness,
    contrast_numerator: 1,
    contrast_denominator: 1,
    noise_seed: seed,
    noise_magnitude,
    blur_radius: max_i(0, blur_radius),
    rotate_quarter_turns,
    flip_horizontal,
  }
}

///|
pub fn AugmentationConfig::with_contrast(
  self : AugmentationConfig,
  numerator : Int,
  denominator : Int,
) -> AugmentationConfig {
  {
    ..self,
    contrast_numerator: numerator,
    contrast_denominator: if denominator == 0 {
      1
    } else {
      denominator
    },
  }
}

///|
pub fn augment_image(
  image : GrayImage,
  config : AugmentationConfig,
) -> GrayImage {
  let bright = image.add_brightness(config.brightness)
  let contrasted = bright.multiply_contrast(
    config.contrast_numerator,
    config.contrast_denominator,
  )
  let noisy = if config.noise_magnitude == 0 {
    contrasted
  } else {
    add_uniform_noise(
      contrasted,
      seed=config.noise_seed,
      magnitude=config.noise_magnitude,
    )
  }
  let blurred = if config.blur_radius == 0 {
    noisy
  } else {
    box_blur(noisy, config.blur_radius)
  }
  let flipped = if config.flip_horizontal {
    blurred.flip_horizontal()
  } else {
    blurred
  }
  flipped.rotate(config.rotate_quarter_turns)
}

///|
pub fn augment_rect(
  rect : Rect,
  source : Size,
  config : AugmentationConfig,
) -> Rect {
  let mut result = rect
  if config.flip_horizontal {
    result = Rect::new(
      x=source.width - rect.right(),
      y=rect.y,
      width=rect.width,
      height=rect.height,
    )
  }
  let turns = (config.rotate_quarter_turns % 4 + 4) % 4
  if turns == 0 {
    clamp_rect(result, source)
  } else if turns == 2 {
    clamp_rect(
      Rect::new(
        x=source.width - result.right(),
        y=source.height - result.bottom(),
        width=result.width,
        height=result.height,
      ),
      source,
    )
  } else {
    clamp_rect(
      Rect::new(
        x=result.y,
        y=source.width - result.right(),
        width=result.height,
        height=result.width,
      ),
      Size::new(width=source.height, height=source.width),
    )
  }
}

///|
pub fn augment_boxes(
  boxes : Array[Box2D],
  source : Size,
  config : AugmentationConfig,
) -> Array[Box2D] {
  let out : Array[Box2D] = []
  let target_size = if config.rotate_quarter_turns % 2 == 0 {
    source
  } else {
    Size::new(width=source.height, height=source.width)
  }
  for box in boxes {
    out.push(
      Box2D::new(
        label=box.label,
        rect=augment_rect(box.rect, source, config).clip_to(target_size),
        occluded=box.occluded,
      ),
    )
  }
  out
}

///|
fn Rect::clip_to(self : Rect, size : Size) -> Rect {
  clamp_rect(self, size)
}

///|
pub fn augment_scene(
  scene : SceneOutput,
  config : AugmentationConfig,
) -> SceneOutput {
  SceneOutput::new(
    plan=scene.plan,
    image=augment_image(scene.image, config),
    mask=scene.mask,
    boxes=augment_boxes(scene.boxes, scene.image.size(), config),
    flow=scene.flow,
  )
}

///|
pub fn augmentation_grid(
  seeds : Array[UInt],
  brightness_values : Array[Int],
) -> Array[AugmentationConfig] {
  let out : Array[AugmentationConfig] = []
  for seed in seeds {
    for brightness in brightness_values {
      out.push(AugmentationConfig::with_seed(seed, brightness~))
    }
  }
  out
}