///|
pub struct FoldAssignment {
  asset_name : String
  fold : Int
  role : String
} derive(Debug, Eq)

///|
pub fn FoldAssignment::new(
  asset_name~ : String,
  fold~ : Int,
  role~ : String,
) -> FoldAssignment {
  { asset_name, fold: max_i(0, fold), role }
}

///|
pub struct CrossValidationPlan {
  folds : Int
  assignments : Array[FoldAssignment]
} derive(Debug, Eq)

///|
pub fn CrossValidationPlan::new(folds~ : Int) -> CrossValidationPlan {
  { folds: max_i(1, folds), assignments: [] }
}

///|
pub fn CrossValidationPlan::add(
  self : CrossValidationPlan,
  assignment : FoldAssignment,
) -> Unit {
  self.assignments.push(assignment)
}

///|
pub fn CrossValidationPlan::fold(
  self : CrossValidationPlan,
  fold : Int,
) -> Array[String] {
  let out : Array[String] = []
  let target = clamp(fold, low=0, high=self.folds - 1)
  for assignment in self.assignments {
    if assignment.fold == target && assignment.role == "validation" {
      out.push(assignment.asset_name)
    }
  }
  out
}

///|
pub fn CrossValidationPlan::training_fold(
  self : CrossValidationPlan,
  fold : Int,
) -> Array[String] {
  let out : Array[String] = []
  let target = clamp(fold, low=0, high=self.folds - 1)
  for assignment in self.assignments {
    if assignment.fold != target && assignment.role == "training" {
      out.push(assignment.asset_name)
    }
  }
  out
}

///|
pub fn CrossValidationPlan::valid(self : CrossValidationPlan) -> Bool {
  if self.assignments.length() == 0 {
    false
  } else {
    for assignment in self.assignments {
      if assignment.fold >= self.folds || assignment.asset_name.length() == 0 {
        return false
      }
    }
    true
  }
}

///|
pub fn CrossValidationPlan::to_json(self : CrossValidationPlan) -> String {
  let sb = StringBuilder()
  sb.write_string("{\"folds\":\{self.folds},\"assignments\":[")
  for i, assignment in self.assignments {
    if i > 0 {
      sb.write_string(",")
    }
    sb.write_string(
      "{\"asset\":\"\{escape_json(assignment.asset_name)}\",\"fold\":\{assignment.fold},\"role\":\"\{assignment.role}\"}",
    )
  }
  sb.write_string("]}")
  sb.to_string()
}

///|
pub fn cross_validation(
  manifest : DatasetManifest,
  folds : Int,
  seed : UInt,
) -> CrossValidationPlan {
  let plan = CrossValidationPlan::new(folds=max_i(1, folds))
  let count = plan.folds
  for i, asset in manifest.assets {
    let fold = ((i * 1103515245 + seed.reinterpret_as_int()) & 0x7fffffff) %
      count
    for candidate in 0.. SplitSummary {
  let names : Array[String] = []
  for asset in split.train {
    names.push(asset.name)
  }
  for asset in split.validation {
    names.push(asset.name)
  }
  for asset in split.holdout {
    names.push(asset.name)
  }
  let mut unique = 0
  for i, name in names {
    let mut repeated = false
    for j in 0.. Int {
  self.training + self.validation + self.holdout
}

///|
pub fn SplitSummary::is_disjoint(self : SplitSummary) -> Bool {
  self.total() == self.unique_assets
}

///|
pub fn balanced_split(
  manifest : DatasetManifest,
  train_per_1000 : Int,
  validation_per_1000 : Int,
) -> DatasetSplit {
  let train : Array[Asset] = []
  let validation : Array[Asset] = []
  let holdout : Array[Asset] = []
  let train_limit = clamp(train_per_1000, low=0, high=1000)
  let validation_limit = clamp(
    train_limit + validation_per_1000,
    low=train_limit,
    high=1000,
  )
  let per_tag : Array[String] = []
  for asset in manifest.assets {
    for tag in asset.tags {
      if !contains_string(per_tag, tag) {
        per_tag.push(tag)
      }
    }
  }
  for i, asset in manifest.assets {
    let tag_count = max_i(1, asset.tags.length())
    let bucket = (i * 997 + tag_count * 31) % 1000
    if bucket < train_limit {
      train.push(asset)
    } else if bucket < validation_limit {
      validation.push(asset)
    } else {
      holdout.push(asset)
    }
  }
  { train, validation, holdout }
}

///|
pub fn temporal_split(
  manifest : DatasetManifest,
  holdout_count : Int,
) -> DatasetSplit {
  let train : Array[Asset] = []
  let validation : Array[Asset] = []
  let holdout : Array[Asset] = []
  let holdout_start = max_i(
    0,
    manifest.assets.length() - max_i(0, holdout_count),
  )
  let validation_start = max_i(
    0,
    holdout_start - max_i(1, manifest.assets.length() / 5),
  )
  for i, asset in manifest.assets {
    if i >= holdout_start {
      holdout.push(asset)
    } else if i >= validation_start {
      validation.push(asset)
    } else {
      train.push(asset)
    }
  }
  { train, validation, holdout }
}

///|
pub fn split_summary_json(split : DatasetSplit) -> String {
  let summary = SplitSummary::from_manifest(split)
  "{\"training\":\{summary.training},\"validation\":\{summary.validation},\"holdout\":\{summary.holdout},\"unique_assets\":\{summary.unique_assets},\"disjoint\":\{summary.is_disjoint()}}"
}

///|
pub fn validate_split(split : DatasetSplit) -> ValidationReport {
  let report = ValidationReport::new()
  let summary = SplitSummary::from_manifest(split)
  if !summary.is_disjoint() {
    report.error("dataset split contains duplicate asset names")
  }
  if summary.total() == 0 {
    report.warning("dataset split is empty")
  }
  report
}

///|
pub fn split_checksum(split : DatasetSplit) -> UInt {
  let mut checksum : UInt = 2166136261
  for asset in split.train {
    checksum = (checksum ^ asset.seed) * 16777619
  }
  for asset in split.validation {
    checksum = (checksum ^ asset.seed) * 16777619
  }
  for asset in split.holdout {
    checksum = (checksum ^ asset.seed) * 16777619
  }
  checksum
}

///|
pub fn dataset_difference(
  left : DatasetManifest,
  right : DatasetManifest,
) -> Array[Asset] {
  let out : Array[Asset] = []
  for asset in left.assets {
    let mut found = false
    for candidate in right.assets {
      if candidate.name == asset.name {
        found = true
      }
    }
    if !found {
      out.push(asset)
    }
  }
  out
}

///|
pub fn dataset_intersection(
  left : DatasetManifest,
  right : DatasetManifest,
) -> Array[Asset] {
  let out : Array[Asset] = []
  for asset in left.assets {
    for candidate in right.assets {
      if candidate.name == asset.name {
        out.push(asset)
      }
    }
  }
  out
}