///|
pub enum VariantKind {
  Clean
  LowLight
  HighContrast
  Defocused
  SaltPepper
  Shifted
  Perspective
} derive(Debug, Eq)

///|
pub struct VariantPlan {
  name : String
  kind : VariantKind
  seed : UInt
  amount : Int
} derive(Debug, Eq)

///|
pub fn VariantPlan::new(
  name~ : String,
  kind~ : VariantKind,
  seed~ : UInt,
  amount? : Int = 0,
) -> VariantPlan {
  { name, kind, seed, amount: max_i(0, amount) }
}

///|
pub struct SyntheticPair {
  plan : VariantPlan
  reference : GrayImage
  variant : GrayImage
  flow : Array[FlowVector]
} derive(Debug, Eq)

///|
pub fn SyntheticPair::new(
  plan~ : VariantPlan,
  reference~ : GrayImage,
  variant~ : GrayImage,
  flow? : Array[FlowVector] = [],
) -> SyntheticPair {
  { plan, reference, variant, flow }
}

///|
pub fn render_variant(
  plan : VariantPlan,
  reference : GrayImage,
) -> SyntheticPair {
  match plan.kind {
    Clean => SyntheticPair::new(plan~, reference~, variant=reference.copy())
    LowLight =>
      SyntheticPair::new(
        plan~,
        reference~,
        variant=reference.multiply_contrast(6, 10).add_brightness(-plan.amount),
      )
    HighContrast =>
      SyntheticPair::new(
        plan~,
        reference~,
        variant=reference.multiply_contrast(max_i(11, plan.amount), 10),
      )
    Defocused =>
      SyntheticPair::new(
        plan~,
        reference~,
        variant=box_blur(reference, max_i(1, plan.amount)),
      )
    SaltPepper =>
      SyntheticPair::new(
        plan~,
        reference~,
        variant=add_salt_pepper(
          reference,
          seed=plan.seed,
          amount_per_1000=plan.amount,
        ),
      )
    Shifted => {
      let dx = plan.amount
      let dy = plan.amount / 2
      SyntheticPair::new(
        plan~,
        reference~,
        variant=reference.translate(dx, dy, Gray::black()),
        flow=constant_flow(reference.size(), 1, dx~, dy~),
      )
    }
    Perspective => {
      let transform = Projective2D::new(
        h11=1000,
        h12=0,
        h13=plan.amount,
        h21=0,
        h22=1000,
        h23=plan.amount / 2,
        h31=plan.amount / 100,
        h32=plan.amount / 150,
      )
      SyntheticPair::new(
        plan~,
        reference~,
        variant=reference.warp_projective(transform),
      )
    }
  }
}

///|
pub fn variant_kind_name(kind : VariantKind) -> String {
  match kind {
    Clean => "clean"
    LowLight => "low-light"
    HighContrast => "high-contrast"
    Defocused => "defocused"
    SaltPepper => "salt-pepper"
    Shifted => "shifted"
    Perspective => "perspective"
  }
}

///|
pub fn default_variant_plans(seed : UInt) -> Array[VariantPlan] {
  [
    VariantPlan::new(name="clean", kind=Clean, seed~, amount=0),
    VariantPlan::new(name="low-light", kind=LowLight, seed=seed + 1, amount=12),
    VariantPlan::new(
      name="contrast",
      kind=HighContrast,
      seed=seed + 2,
      amount=16,
    ),
    VariantPlan::new(name="defocus", kind=Defocused, seed=seed + 3, amount=2),
    VariantPlan::new(
      name="salt-pepper",
      kind=SaltPepper,
      seed=seed + 4,
      amount=40,
    ),
    VariantPlan::new(name="shifted", kind=Shifted, seed=seed + 5, amount=3),
    VariantPlan::new(
      name="perspective",
      kind=Perspective,
      seed=seed + 6,
      amount=4,
    ),
  ]
}

///|
pub fn render_variants(
  reference : GrayImage,
  plans : Array[VariantPlan],
) -> Array[SyntheticPair] {
  let out : Array[SyntheticPair] = []
  for plan in plans {
    out.push(render_variant(plan, reference))
  }
  out
}

///|
pub struct VariantQuality {
  name : String
  kind : VariantKind
  mae : Int
  mse : Int
  flow_error : Int
} derive(Debug, Eq)

///|
pub fn VariantQuality::from_pair(pair : SyntheticPair) -> VariantQuality {
  {
    name: pair.plan.name,
    kind: pair.plan.kind,
    mae: image_mae(pair.reference, pair.variant),
    mse: image_mse(pair.reference, pair.variant),
    flow_error: flow_endpoint_error_l1(pair.flow, pair.flow),
  }
}

///|
pub fn variant_quality_report(pairs : Array[SyntheticPair]) -> String {
  let sb = StringBuilder()
  sb.write_string("[")
  for i, pair in pairs {
    if i > 0 {
      sb.write_string(",")
    }
    let quality = VariantQuality::from_pair(pair)
    sb.write_string(
      "{\"name\":\"\{escape_json(quality.name)}\",\"kind\":\"\{variant_kind_name(quality.kind)}\",\"mae\":\{quality.mae},\"mse\":\{quality.mse},\"flow_error\":\{quality.flow_error}}",
    )
  }
  sb.write_string("]")
  sb.to_string()
}

///|
pub fn variant_manifest(
  reference : GrayImage,
  plans : Array[VariantPlan],
  prefix : String,
) -> DatasetManifest {
  let manifest = DatasetManifest::new(name="synthetic-variants")
  let reference_tag = if reference.size().is_valid() {
    "reference"
  } else {
    "empty-reference"
  }
  for plan in plans {
    manifest.add(
      Asset::new(
        name=plan.name,
        image_path=prefix + "/" + plan.name + ".pgm",
        annotation_path=prefix + "/" + plan.name + ".json",
        seed=plan.seed,
        tags=["variant", variant_kind_name(plan.kind), reference_tag],
      ),
    )
  }
  manifest
}

///|
pub fn validate_variant_pair(pair : SyntheticPair) -> ValidationReport {
  let report = ValidationReport::new()
  if pair.reference.size() != pair.variant.size() {
    report.error("variant must preserve reference image dimensions")
  }
  if pair.plan.name.length() == 0 {
    report.error("variant name must not be empty")
  }
  if pair.plan.kind == Shifted && pair.flow.length() == 0 {
    report.error("shifted variant must include ground-truth flow")
  }
  report
}

///|
pub fn pair_checksum(pair : SyntheticPair) -> UInt {
  let mut checksum = image_checksum(pair.reference)
  checksum = (checksum ^ image_checksum(pair.variant)) * 16777619
  for vector in pair.flow {
    checksum = (checksum ^ vector.dx.reinterpret_as_uint()) * 16777619
    checksum = (checksum ^ vector.dy.reinterpret_as_uint()) * 16777619
  }
  checksum
}