///|
/// A deterministic schedule for changing a training hyperparameter by step.
pub struct LinearSchedule {
  start : Double
  end : Double
  duration : Int
}

///|
pub fn LinearSchedule::new(
  start : Double,
  end : Double,
  duration : Int,
) -> LinearSchedule {
  { start, end, duration: if duration < 0 { 0 } else { duration } }
}

///|
pub fn LinearSchedule::value(self : LinearSchedule, step : Int) -> Double {
  if self.duration == 0 {
    return self.end
  }
  let position = if step < 0 {
    0
  } else if step > self.duration {
    self.duration
  } else {
    step
  }
  let ratio = position.to_double() / self.duration.to_double()
  self.start + (self.end - self.start) * ratio
}

///|
pub fn LinearSchedule::start(self : LinearSchedule) -> Double {
  self.start
}

///|
pub fn LinearSchedule::end(self : LinearSchedule) -> Double {
  self.end
}

///|
pub fn LinearSchedule::duration(self : LinearSchedule) -> Int {
  self.duration
}

///|
pub struct ExponentialSchedule {
  initial : Double
  final_value : Double
  decay_steps : Int
}

///|
pub fn ExponentialSchedule::new(
  initial : Double,
  final_value : Double,
  decay_steps : Int,
) -> ExponentialSchedule {
  {
    initial,
    final_value,
    decay_steps: if decay_steps < 0 {
      0
    } else {
      decay_steps
    },
  }
}

///|
pub fn ExponentialSchedule::value(
  self : ExponentialSchedule,
  step : Int,
) -> Double {
  if self.decay_steps == 0 {
    return self.final_value
  }
  let bounded = if step < 0 {
    0
  } else if step > self.decay_steps {
    self.decay_steps
  } else {
    step
  }
  let ratio = bounded.to_double() / self.decay_steps.to_double()
  self.initial + (self.final_value - self.initial) * ratio
}

///|
pub fn ExponentialSchedule::initial(self : ExponentialSchedule) -> Double {
  self.initial
}

///|
pub fn ExponentialSchedule::final_value(self : ExponentialSchedule) -> Double {
  self.final_value
}

///|
pub fn ExponentialSchedule::decay_steps(self : ExponentialSchedule) -> Int {
  self.decay_steps
}

///|
pub fn schedule_values(schedule : LinearSchedule, steps : Int) -> Array[Double] {
  let result : Array[Double] = []
  if steps <= 0 {
    return result
  }
  let mut i = 0
  while i < steps {
    result.push(schedule.value(i))
    i = i + 1
  }
  result
}

///|
pub fn schedule_delta(schedule : LinearSchedule, step : Int) -> Double {
  schedule.value(step + 1) - schedule.value(step)
}

///|
pub fn schedule_is_finished(schedule : LinearSchedule, step : Int) -> Bool {
  step >= schedule.duration()
}

///|
pub fn schedule_progress(schedule : LinearSchedule, step : Int) -> Double {
  if schedule.duration() == 0 {
    1.0
  } else {
    let bounded = if step < 0 {
      0
    } else if step > schedule.duration() {
      schedule.duration()
    } else {
      step
    }
    bounded.to_double() / schedule.duration().to_double()
  }
}

///|
pub struct WarmupCosineSchedule {
  warmup_steps : Int
  total_steps : Int
  maximum : Double
  minimum : Double
}

///|
pub fn WarmupCosineSchedule::new(
  warmup_steps : Int,
  total_steps : Int,
  maximum : Double,
  minimum : Double,
) -> WarmupCosineSchedule {
  {
    warmup_steps: if warmup_steps < 0 {
      0
    } else {
      warmup_steps
    },
    total_steps: if total_steps < 0 {
      0
    } else {
      total_steps
    },
    maximum,
    minimum,
  }
}

///|
pub fn WarmupCosineSchedule::value(
  self : WarmupCosineSchedule,
  step : Int,
) -> Double {
  if self.total_steps == 0 {
    return self.minimum
  }
  if step < self.warmup_steps && self.warmup_steps > 0 {
    return self.maximum * step.to_double() / self.warmup_steps.to_double()
  }
  if step >= self.total_steps {
    return self.minimum
  }
  let remaining = self.total_steps - self.warmup_steps
  if remaining <= 0 {
    return self.minimum
  }
  let position = step - self.warmup_steps
  let ratio = position.to_double() / remaining.to_double()
  self.minimum + (self.maximum - self.minimum) * (1.0 - ratio)
}

///|
pub fn WarmupCosineSchedule::warmup_steps(self : WarmupCosineSchedule) -> Int {
  self.warmup_steps
}

///|
pub fn WarmupCosineSchedule::total_steps(self : WarmupCosineSchedule) -> Int {
  self.total_steps
}

///|
pub fn WarmupCosineSchedule::maximum(self : WarmupCosineSchedule) -> Double {
  self.maximum
}

///|
pub fn WarmupCosineSchedule::minimum(self : WarmupCosineSchedule) -> Double {
  self.minimum
}

///|
pub struct BatchCursor {
  total : Int
  batch_size : Int
  mut position : Int
}

///|
pub fn BatchCursor::new(total : Int, batch_size : Int) -> BatchCursor {
  {
    total: if total < 0 {
      0
    } else {
      total
    },
    batch_size: if batch_size < 1 {
      1
    } else {
      batch_size
    },
    position: 0,
  }
}

///|
pub fn BatchCursor::total(self : BatchCursor) -> Int {
  self.total
}

///|
pub fn BatchCursor::batch_size(self : BatchCursor) -> Int {
  self.batch_size
}

///|
pub fn BatchCursor::position(self : BatchCursor) -> Int {
  self.position
}

///|
pub fn BatchCursor::has_next(self : BatchCursor) -> Bool {
  self.position < self.total
}

///|
pub fn BatchCursor::remaining(self : BatchCursor) -> Int {
  if self.total > self.position {
    self.total - self.position
  } else {
    0
  }
}

///|
pub fn BatchCursor::next_indices(self : BatchCursor) -> Array[Int] {
  let result : Array[Int] = []
  if !self.has_next() {
    return result
  }
  let mut i = 0
  while i < self.batch_size && self.position + i < self.total {
    result.push(self.position + i)
    i = i + 1
  }
  self.position = self.position + result.length()
  result
}

///|
pub fn BatchCursor::reset(self : BatchCursor) -> Unit {
  self.position = 0
}

///|
pub fn BatchCursor::batch_count(self : BatchCursor) -> Int {
  if self.total == 0 {
    0
  } else {
    (self.total + self.batch_size - 1) / self.batch_size
  }
}