///|
/// Online decision policies for recommendation, routing, and experimentation.
///
/// The policies in this file are stateful by design. Each policy exposes its
/// counters and reset operation so it can be checkpointed between service
/// instances and tested with deterministic event streams.
pub(all) enum BanditAction {
  Explore
  Exploit
} derive(Debug, Eq)

///|
pub fn bandit_action_catalog() -> Array[BanditAction] {
  [Explore, Exploit]
}

///|
pub struct ArmStatistics {
  arm : Int
  mut pulls : Int
  mut reward_sum : Double
  mut reward_square_sum : Double
  mut last_reward : Double
}

///|
pub fn ArmStatistics::new(arm : Int) -> ArmStatistics {
  { arm, pulls: 0, reward_sum: 0.0, reward_square_sum: 0.0, last_reward: 0.0 }
}

///|
pub fn ArmStatistics::arm(self : ArmStatistics) -> Int {
  self.arm
}

///|
pub fn ArmStatistics::pulls(self : ArmStatistics) -> Int {
  self.pulls
}

///|
pub fn ArmStatistics::reward_sum(self : ArmStatistics) -> Double {
  self.reward_sum
}

///|
pub fn ArmStatistics::last_reward(self : ArmStatistics) -> Double {
  self.last_reward
}

///|
pub fn ArmStatistics::mean(self : ArmStatistics) -> Double {
  if self.pulls == 0 {
    0.0
  } else {
    self.reward_sum / self.pulls.to_double()
  }
}

///|
pub fn ArmStatistics::variance(self : ArmStatistics) -> Double {
  if self.pulls <= 1 {
    0.0
  } else {
    let count = self.pulls.to_double()
    let value = (
        self.reward_square_sum - self.reward_sum * self.reward_sum / count
      ) /
      (count - 1.0)
    if value < 0.0 {
      0.0
    } else {
      value
    }
  }
}

///|
pub fn ArmStatistics::standard_error(self : ArmStatistics) -> Double {
  if self.pulls == 0 {
    0.0
  } else {
    (self.variance() / self.pulls.to_double()).sqrt()
  }
}

///|
pub fn ArmStatistics::update(self : ArmStatistics, reward : Double) -> Unit {
  let value = if reward.is_nan() || reward.is_inf() { 0.0 } else { reward }
  self.pulls += 1
  self.reward_sum += value
  self.reward_square_sum += value * value
  self.last_reward = value
}

///|
pub fn ArmStatistics::reset(self : ArmStatistics) -> Unit {
  self.pulls = 0
  self.reward_sum = 0.0
  self.reward_square_sum = 0.0
  self.last_reward = 0.0
}

///|
pub struct BanditHistory {
  capacity : Int
  arms : Array[Int]
  rewards : Array[Double]
  actions : Array[BanditAction]
}

///|
pub fn BanditHistory::new(capacity? : Int = 1024) -> BanditHistory {
  {
    capacity: if capacity < 1 {
      1
    } else {
      capacity
    },
    arms: [],
    rewards: [],
    actions: [],
  }
}

///|
pub fn BanditHistory::append(
  self : BanditHistory,
  arm : Int,
  reward : Double,
  action : BanditAction,
) -> Unit {
  self.arms.push(arm)
  self.rewards.push(reward)
  self.actions.push(action)
  if self.arms.length() > self.capacity {
    let _ = self.arms.remove(0)
    let _ = self.rewards.remove(0)
    let _ = self.actions.remove(0)
  }
}

///|
pub fn BanditHistory::size(self : BanditHistory) -> Int {
  self.arms.length()
}

///|
pub fn BanditHistory::arms(self : BanditHistory) -> Array[Int] {
  self.arms.map(value => value)
}

///|
pub fn BanditHistory::rewards(self : BanditHistory) -> Array[Double] {
  copy_vector(self.rewards)
}

///|
pub fn BanditHistory::actions(self : BanditHistory) -> Array[BanditAction] {
  self.actions.map(value => value)
}

///|
pub fn BanditHistory::mean_reward(self : BanditHistory) -> Double {
  if self.rewards.is_empty() {
    0.0
  } else {
    sum_values(self.rewards) / self.rewards.length().to_double()
  }
}

///|
pub fn BanditHistory::exploration_rate(self : BanditHistory) -> Double {
  if self.actions.is_empty() {
    0.0
  } else {
    self.actions
    .count_if(action => {
      match action {
        Explore => true
        _ => false
      }
    })
    .to_double() /
    self.actions.length().to_double()
  }
}

///|
pub fn BanditHistory::reset(self : BanditHistory) -> Unit {
  self.arms.clear()
  self.rewards.clear()
  self.actions.clear()
}

///|
pub struct EpsilonGreedyBandit {
  arms : Array[ArmStatistics]
  epsilon : Double
  rng : DeterministicRng
  history : BanditHistory
  mut steps : Int
  mut explored : Int
}

///|
pub fn EpsilonGreedyBandit::new(
  arm_count : Int,
  epsilon? : Double = 0.1,
  seed? : UInt64 = 1,
  history_capacity? : Int = 1024,
) -> EpsilonGreedyBandit {
  let size = if arm_count < 0 { 0 } else { arm_count }
  {
    arms: Array::makei(size, i => ArmStatistics::new(i)),
    epsilon: clamp(epsilon, 0.0, 1.0),
    rng: DeterministicRng::new(seed),
    history: BanditHistory::new(capacity=history_capacity),
    steps: 0,
    explored: 0,
  }
}

///|
pub fn EpsilonGreedyBandit::arm_count(self : EpsilonGreedyBandit) -> Int {
  self.arms.length()
}

///|
pub fn EpsilonGreedyBandit::epsilon(self : EpsilonGreedyBandit) -> Double {
  self.epsilon
}

///|
pub fn EpsilonGreedyBandit::statistics(
  self : EpsilonGreedyBandit,
  arm : Int,
) -> ArmStatistics? {
  self.arms.get(arm)
}

///|
fn EpsilonGreedyBandit::best_arm(self : EpsilonGreedyBandit) -> Int? {
  if self.arms.is_empty() {
    None
  } else {
    let mut best = 0
    let mut best_value = self.arms[0].mean()
    for i in 1.. best_value {
        best = i
        best_value = value
      }
    }
    Some(best)
  }
}

///|
pub fn EpsilonGreedyBandit::choose(
  self : EpsilonGreedyBandit,
) -> (Int, BanditAction)? {
  if self.arms.is_empty() {
    None
  } else {
    let mut selected = -1
    for i in 0..= 0 || self.rng.next_double() < self.epsilon
    let action = if should_explore { Explore } else { Exploit }
    let arm = if selected >= 0 {
      selected
    } else {
      self.best_arm().unwrap_or(0)
    }
    if should_explore {
      self.explored += 1
    }
    Some((arm, action))
  }
}

///|
pub fn EpsilonGreedyBandit::update(
  self : EpsilonGreedyBandit,
  arm : Int,
  reward : Double,
  action? : BanditAction = Exploit,
) -> Bool {
  match self.arms.get(arm) {
    None => false
    Some(statistics) => {
      statistics.update(reward)
      self.steps += 1
      self.history.append(arm, reward, action)
      true
    }
  }
}

///|
pub fn EpsilonGreedyBandit::steps(self : EpsilonGreedyBandit) -> Int {
  self.steps
}

///|
pub fn EpsilonGreedyBandit::explored(self : EpsilonGreedyBandit) -> Int {
  self.explored
}

///|
pub fn EpsilonGreedyBandit::history(
  self : EpsilonGreedyBandit,
) -> BanditHistory {
  self.history
}

///|
pub fn EpsilonGreedyBandit::reset(self : EpsilonGreedyBandit) -> Unit {
  for arm in self.arms {
    arm.reset()
  }
  self.history.reset()
  self.steps = 0
  self.explored = 0
}

///|
pub struct UcbBandit {
  arms : Array[ArmStatistics]
  exploration : Double
  history : BanditHistory
  mut steps : Int
}

///|
pub fn UcbBandit::new(
  arm_count : Int,
  exploration? : Double = 2.0,
  history_capacity? : Int = 1024,
) -> UcbBandit {
  let size = if arm_count < 0 { 0 } else { arm_count }
  {
    arms: Array::makei(size, i => ArmStatistics::new(i)),
    exploration: if exploration < 0.0 {
      0.0
    } else {
      exploration
    },
    history: BanditHistory::new(capacity=history_capacity),
    steps: 0,
  }
}

///|
pub fn UcbBandit::arm_count(self : UcbBandit) -> Int {
  self.arms.length()
}

///|
pub fn UcbBandit::statistics(self : UcbBandit, arm : Int) -> ArmStatistics? {
  self.arms.get(arm)
}

///|
pub fn UcbBandit::upper_bound(self : UcbBandit, arm : Int) -> Double {
  match self.arms.get(arm) {
    None => 0.0
    Some(statistics) =>
      if statistics.pulls() == 0 {
        1.0e30
      } else {
        statistics.mean() +
        self.exploration *
        (2.0 *
        @math.ln(self.steps.to_double() + 1.0) /
        statistics.pulls().to_double()).sqrt()
      }
  }
}

///|
pub fn UcbBandit::choose(self : UcbBandit) -> Int? {
  if self.arms.is_empty() {
    None
  } else {
    let mut best = 0
    let mut best_value = self.upper_bound(0)
    for i in 1.. best_value {
        best = i
        best_value = value
      }
    }
    Some(best)
  }
}

///|
pub fn UcbBandit::update(self : UcbBandit, arm : Int, reward : Double) -> Bool {
  match self.arms.get(arm) {
    None => false
    Some(statistics) => {
      statistics.update(reward)
      self.steps += 1
      self.history.append(arm, reward, Explore)
      true
    }
  }
}

///|
pub fn UcbBandit::steps(self : UcbBandit) -> Int {
  self.steps
}

///|
pub fn UcbBandit::history(self : UcbBandit) -> BanditHistory {
  self.history
}

///|
pub fn UcbBandit::reset(self : UcbBandit) -> Unit {
  for arm in self.arms {
    arm.reset()
  }
  self.history.reset()
  self.steps = 0
}

///|
pub struct SoftmaxBandit {
  arms : Array[ArmStatistics]
  temperature : Double
  rng : DeterministicRng
  history : BanditHistory
  mut steps : Int
}

///|
pub fn SoftmaxBandit::new(
  arm_count : Int,
  temperature? : Double = 0.5,
  seed? : UInt64 = 1,
  history_capacity? : Int = 1024,
) -> SoftmaxBandit {
  let size = if arm_count < 0 { 0 } else { arm_count }
  {
    arms: Array::makei(size, i => ArmStatistics::new(i)),
    temperature: if temperature <= 0.0 {
      0.1
    } else {
      temperature
    },
    rng: DeterministicRng::new(seed),
    history: BanditHistory::new(capacity=history_capacity),
    steps: 0,
  }
}

///|
pub fn SoftmaxBandit::arm_count(self : SoftmaxBandit) -> Int {
  self.arms.length()
}

///|
pub fn SoftmaxBandit::probabilities(self : SoftmaxBandit) -> Array[Double] {
  if self.arms.is_empty() {
    []
  } else {
    let maximum = self.arms
      .map(arm => arm.mean())
      .fold(init=-1.0e30, (left, right) => {
        if left > right {
          left
        } else {
          right
        }
      })
    let values = self.arms.map(arm => {
      @math.exp((arm.mean() - maximum) / self.temperature)
    })
    let total = sum_values(values)
    if total <= 0.0 {
      Array::make(values.length(), 1.0 / values.length().to_double())
    } else {
      values.map(value => value / total)
    }
  }
}

///|
pub fn SoftmaxBandit::choose(self : SoftmaxBandit) -> Int? {
  let probabilities = self.probabilities()
  if probabilities.is_empty() {
    None
  } else {
    let draw = self.rng.next_double()
    let mut cumulative = 0.0
    let mut selected = probabilities.length() - 1
    for i in 0.. Bool {
  match self.arms.get(arm) {
    None => false
    Some(statistics) => {
      statistics.update(reward)
      self.steps += 1
      self.history.append(arm, reward, Exploit)
      true
    }
  }
}

///|
pub fn SoftmaxBandit::steps(self : SoftmaxBandit) -> Int {
  self.steps
}

///|
pub fn SoftmaxBandit::history(self : SoftmaxBandit) -> BanditHistory {
  self.history
}

///|
pub fn SoftmaxBandit::reset(self : SoftmaxBandit) -> Unit {
  for arm in self.arms {
    arm.reset()
  }
  self.history.reset()
  self.steps = 0
}

///|
pub struct PosteriorBandit {
  alpha : Array[Double]
  beta : Array[Double]
  rng : DeterministicRng
  mut steps : Int
}

///|
pub fn PosteriorBandit::new(
  arm_count : Int,
  prior? : Double = 1.0,
  seed? : UInt64 = 1,
) -> PosteriorBandit {
  let size = if arm_count < 0 { 0 } else { arm_count }
  let safe_prior = if prior <= 0.0 { 1.0 } else { prior }
  {
    alpha: Array::make(size, safe_prior),
    beta: Array::make(size, safe_prior),
    rng: DeterministicRng::new(seed),
    steps: 0,
  }
}

///|
pub fn PosteriorBandit::arm_count(self : PosteriorBandit) -> Int {
  self.alpha.length()
}

///|
pub fn PosteriorBandit::mean(self : PosteriorBandit, arm : Int) -> Double {
  match (self.alpha.get(arm), self.beta.get(arm)) {
    (Some(a), Some(b)) => a / (a + b)
    _ => 0.0
  }
}

///|
pub fn PosteriorBandit::uncertainty(
  self : PosteriorBandit,
  arm : Int,
) -> Double {
  match (self.alpha.get(arm), self.beta.get(arm)) {
    (Some(a), Some(b)) => {
      let total = a + b
      if total <= 0.0 {
        0.0
      } else {
        (a * b / (total * total * (total + 1.0))).sqrt()
      }
    }
    _ => 0.0
  }
}

///|
pub fn PosteriorBandit::sample(self : PosteriorBandit, arm : Int) -> Double {
  let value = self.mean(arm) + self.rng.normal() * self.uncertainty(arm)
  clamp(value, 0.0, 1.0)
}

///|
pub fn PosteriorBandit::choose(self : PosteriorBandit) -> Int? {
  if self.alpha.is_empty() {
    None
  } else {
    let mut selected = 0
    let mut best = self.sample(0)
    for i in 1.. best {
        selected = i
        best = value
      }
    }
    Some(selected)
  }
}

///|
pub fn PosteriorBandit::update(
  self : PosteriorBandit,
  arm : Int,
  reward : Double,
) -> Bool {
  if arm < 0 || arm >= self.alpha.length() {
    false
  } else {
    let value = clamp(reward, 0.0, 1.0)
    self.alpha[arm] += value
    self.beta[arm] += 1.0 - value
    self.steps += 1
    true
  }
}

///|
pub fn PosteriorBandit::steps(self : PosteriorBandit) -> Int {
  self.steps
}

///|
pub fn PosteriorBandit::reset(self : PosteriorBandit) -> Unit {
  self.alpha.fill(1.0)
  self.beta.fill(1.0)
  self.steps = 0
}

///|
pub struct ContextualBandit {
  models : Array[OnlineRidgeRegression]
  exploration : Double
  mut steps : Int
}

///|
pub fn ContextualBandit::new(
  arm_count : Int,
  dimension : Int,
  exploration? : Double = 0.1,
  learning_rate? : Double = 0.02,
) -> ContextualBandit {
  let size = if arm_count < 0 { 0 } else { arm_count }
  {
    models: Array::makei(size, _ => {
      OnlineRidgeRegression::new(dimension, learning_rate~)
    }),
    exploration: if exploration < 0.0 {
      0.0
    } else {
      exploration
    },
    steps: 0,
  }
}

///|
pub fn ContextualBandit::arm_count(self : ContextualBandit) -> Int {
  self.models.length()
}

///|
pub fn ContextualBandit::predict(
  self : ContextualBandit,
  arm : Int,
  features : Array[Double],
) -> Double {
  match self.models.get(arm) {
    None => 0.0
    Some(model) => model.predict(features)
  }
}

///|
pub fn ContextualBandit::choose(
  self : ContextualBandit,
  features : Array[Double],
) -> Int? {
  if self.models.is_empty() {
    None
  } else {
    let mut best_arm = 0
    let mut best_score = self.predict(0, features)
    for i in 1.. best_score {
        best_arm = i
        best_score = score
      }
    }
    Some(best_arm)
  }
}

///|
pub fn ContextualBandit::update(
  self : ContextualBandit,
  arm : Int,
  features : Array[Double],
  reward : Double,
) -> Bool {
  match self.models.get(arm) {
    None => false
    Some(model) => {
      model.update(features, reward)
      self.steps += 1
      true
    }
  }
}

///|
pub fn ContextualBandit::steps(self : ContextualBandit) -> Int {
  self.steps
}

///|
pub fn ContextualBandit::weights(
  self : ContextualBandit,
  arm : Int,
) -> Array[Double] {
  self.models.get(arm).map(model => model.weights()).unwrap_or([])
}

///|
pub fn ContextualBandit::reset(self : ContextualBandit) -> Unit {
  for model in self.models {
    model.reset()
  }
  self.steps = 0
}

///|
pub struct BanditEvaluator {
  mut observations : Int
  mut reward_sum : Double
  mut optimal_sum : Double
  mut regret_sum : Double
  mut successes : Int
}

///|
pub fn BanditEvaluator::new() -> BanditEvaluator {
  {
    observations: 0,
    reward_sum: 0.0,
    optimal_sum: 0.0,
    regret_sum: 0.0,
    successes: 0,
  }
}

///|
pub fn BanditEvaluator::observe(
  self : BanditEvaluator,
  reward : Double,
  optimal_reward : Double,
  success? : Bool = false,
) -> Unit {
  self.observations += 1
  self.reward_sum += reward
  self.optimal_sum += optimal_reward
  self.regret_sum += optimal_reward - reward
  if success {
    self.successes += 1
  }
}

///|
pub fn BanditEvaluator::observations(self : BanditEvaluator) -> Int {
  self.observations
}

///|
pub fn BanditEvaluator::mean_reward(self : BanditEvaluator) -> Double {
  if self.observations == 0 {
    0.0
  } else {
    self.reward_sum / self.observations.to_double()
  }
}

///|
pub fn BanditEvaluator::mean_optimal_reward(self : BanditEvaluator) -> Double {
  if self.observations == 0 {
    0.0
  } else {
    self.optimal_sum / self.observations.to_double()
  }
}

///|
pub fn BanditEvaluator::cumulative_regret(self : BanditEvaluator) -> Double {
  self.regret_sum
}

///|
pub fn BanditEvaluator::regret_rate(self : BanditEvaluator) -> Double {
  if self.observations == 0 {
    0.0
  } else {
    self.regret_sum / self.observations.to_double()
  }
}

///|
pub fn BanditEvaluator::success_rate(self : BanditEvaluator) -> Double {
  if self.observations == 0 {
    0.0
  } else {
    self.successes.to_double() / self.observations.to_double()
  }
}

///|
pub fn BanditEvaluator::reset(self : BanditEvaluator) -> Unit {
  self.observations = 0
  self.reward_sum = 0.0
  self.optimal_sum = 0.0
  self.regret_sum = 0.0
  self.successes = 0
}