///|
pub(all) struct BudgetPoint {
  budget : Int
  assignments : Int
  target_assignments : Int
  yes : Int
  no : Int
  unknown : Int
  cheapest_target_cost : Int
} derive(Debug, Eq)

///|
pub(all) struct BudgetReport {
  policy_name : String
  baseline : Truth
  target : Truth
  max_budget : Int
  domain_points : Int
  evaluated : Int
  outside_budget : Int
  truncated : Bool
  points : Array[BudgetPoint]
} derive(Debug, Eq)

///|
priv struct BudgetState {
  model : Policy
  target : Truth
  max_budget : Int
  max_evaluations : Int
  points : Array[BudgetPoint]
  mut evaluated : Int
  mut outside_budget : Int
  mut truncated : Bool
}

///|
pub fn Policy::budget_report(
  self : Policy,
  max_budget? : Int = 16,
  max_evaluations? : Int = 10000,
) -> Result[BudgetReport, Diagnostic] {
  if max_budget < 0 || max_evaluations <= 0 {
    return Err({
      code: "E_BUDGET_LIMIT",
      message: "max_budget must be non-negative and max_evaluations must be positive",
      span: { offset: 0, line: 1, column: 1 },
    })
  }
  let validation = self.validate()
  if !validation.is_empty() {
    return Err(validation[0])
  }
  let baseline = self.evaluate().truth
  let target = target_for_budget(baseline)
  let state = BudgetState::{
    model: self,
    target,
    max_budget,
    max_evaluations,
    points: initial_budget_points(max_budget),
    evaluated: 0,
    outside_budget: 0,
    truncated: false,
  }
  enumerate_budget(state, 0, self.environment(), 0)
  Ok({
    policy_name: self.name,
    baseline,
    target,
    max_budget,
    domain_points: self.search_space_size(),
    evaluated: state.evaluated,
    outside_budget: state.outside_budget,
    truncated: state.truncated,
    points: state.points,
  })
}

///|
fn initial_budget_points(max_budget : Int) -> Array[BudgetPoint] {
  let points : Array[BudgetPoint] = []
  for budget in 0..<(max_budget + 1) {
    points.push({
      budget,
      assignments: 0,
      target_assignments: 0,
      yes: 0,
      no: 0,
      unknown: 0,
      cheapest_target_cost: -1,
    })
  }
  points
}

///|
fn target_for_budget(value : Truth) -> Truth {
  match value {
    Yes => No
    No | Unknown => Yes
  }
}

///|
fn enumerate_budget(
  state : BudgetState,
  index : Int,
  environment : Environment,
  cost : Int,
) -> Unit {
  if state.evaluated >= state.max_evaluations {
    state.truncated = true
    return
  }
  if index >= state.model.features.length() {
    state.evaluated = state.evaluated + 1
    if cost > state.max_budget {
      state.outside_budget = state.outside_budget + 1
      return
    }
    let outcome = evaluate(state.model.goal, environment).truth
    record_budget_outcome(state, cost, outcome)
    return
  }
  let item = state.model.features[index]
  if item.locked {
    enumerate_budget(state, index + 1, environment, cost)
    return
  }
  for candidate in item.domain {
    if state.evaluated >= state.max_evaluations {
      state.truncated = true
      return
    }
    if candidate == item.current {
      enumerate_budget(state, index + 1, environment, cost)
    } else {
      let next_environment = environment_set(environment, item.name, candidate)
      enumerate_budget(state, index + 1, next_environment, cost + item.cost)
    }
  }
}

///|
fn record_budget_outcome(
  state : BudgetState,
  cost : Int,
  outcome : Truth,
) -> Unit {
  for budget in cost..<(state.max_budget + 1) {
    let point = state.points[budget]
    let yes_delta = if outcome == Yes { 1 } else { 0 }
    let no_delta = if outcome == No { 1 } else { 0 }
    let unknown_delta = if outcome == Unknown { 1 } else { 0 }
    let target_delta = if outcome == state.target { 1 } else { 0 }
    let yes = point.yes + yes_delta
    let no = point.no + no_delta
    let unknown = point.unknown + unknown_delta
    let target_assignments = point.target_assignments + target_delta
    let cheapest_target_cost = if outcome == state.target {
      if point.cheapest_target_cost < 0 || cost < point.cheapest_target_cost {
        cost
      } else {
        point.cheapest_target_cost
      }
    } else {
      point.cheapest_target_cost
    }
    state.points[budget] = {
      budget: point.budget,
      assignments: point.assignments + 1,
      target_assignments,
      yes,
      no,
      unknown,
      cheapest_target_cost,
    }
  }
}

///|
pub fn BudgetReport::minimum_budget(self : BudgetReport) -> Int {
  for point in self.points {
    if point.target_assignments > 0 {
      return point.budget
    }
  }
  -1
}

///|
pub fn BudgetReport::point_for(
  self : BudgetReport,
  budget : Int,
) -> BudgetPoint? {
  for point in self.points {
    if point.budget == budget {
      return Some(point)
    }
  }
  None
}

///|
pub fn BudgetPoint::target_per_mille(self : BudgetPoint) -> Int {
  if self.assignments == 0 {
    0
  } else {
    self.target_assignments * 1000 / self.assignments
  }
}

///|
pub fn BudgetPoint::target_ratio(self : BudgetPoint) -> String {
  "\{self.target_assignments}/\{self.assignments}"
}

///|
pub fn BudgetPoint::has_target(self : BudgetPoint) -> Bool {
  self.target_assignments > 0
}

///|
pub fn BudgetReport::complete(self : BudgetReport) -> Bool {
  !self.truncated && self.evaluated >= self.domain_points
}

///|
pub fn BudgetReport::to_markdown(self : BudgetReport) -> String {
  let out = StringBuilder::new()
  out.write_string("# Budget curve: \{self.policy_name}\n\n")
  out.write_string("- Current outcome: **\{self.baseline.render()}**\n")
  out.write_string("- Target outcome: **\{self.target.render()}**\n")
  out.write_string("- Domain assignments: \{self.domain_points}\n")
  out.write_string("- Assignments evaluated: \{self.evaluated}\n")
  out.write_string("- Outside reported budget: \{self.outside_budget}\n")
  out.write_string("- Search truncated: \{self.truncated}\n\n")
  out.write_string(
    "| budget | assignments | target | true | false | unknown | target per mille | cheapest target |\n",
  )
  out.write_string("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |\n")
  for point in self.points {
    let cheapest = if point.cheapest_target_cost >= 0 {
      point.cheapest_target_cost.to_string()
    } else {
      "none"
    }
    out.write_string(
      "| \{point.budget} | \{point.assignments} | \{point.target_assignments} | \{point.yes} | \{point.no} | \{point.unknown} | \{point.target_per_mille()} | \{cheapest} |\n",
    )
  }
  out.to_string()
}

///|
pub fn BudgetReport::to_json(self : BudgetReport) -> String {
  let out = StringBuilder::new()
  out.write_string("{")
  out.write_string("\"policy\":\"\{budget_json_escape(self.policy_name)}\",")
  out.write_string("\"baseline\":\"\{self.baseline.render()}\",")
  out.write_string("\"target\":\"\{self.target.render()}\",")
  out.write_string("\"max_budget\":\{self.max_budget},")
  out.write_string("\"domain_points\":\{self.domain_points},")
  out.write_string("\"evaluated\":\{self.evaluated},")
  out.write_string("\"outside_budget\":\{self.outside_budget},")
  out.write_string("\"truncated\":\{self.truncated},")
  out.write_string("\"minimum_budget\":\{self.minimum_budget()},")
  out.write_string("\"points\":[")
  for index, point in self.points {
    if index > 0 {
      out.write_char(',')
    }
    out.write_string("{")
    out.write_string("\"budget\":\{point.budget},")
    out.write_string("\"assignments\":\{point.assignments},")
    out.write_string("\"target_assignments\":\{point.target_assignments},")
    out.write_string("\"yes\":\{point.yes},")
    out.write_string("\"no\":\{point.no},")
    out.write_string("\"unknown\":\{point.unknown},")
    out.write_string("\"target_per_mille\":\{point.target_per_mille()},")
    out.write_string("\"cheapest_target_cost\":\{point.cheapest_target_cost}")
    out.write_string("}")
  }
  out.write_string("]}")
  out.to_string()
}

///|
fn budget_json_escape(text : String) -> String {
  let out = StringBuilder::new()
  for ch in text {
    match ch {
      '"' => out.write_string("\\\"")
      '\\' => out.write_string("\\\\")
      '\n' => out.write_string("\\n")
      '\r' => out.write_string("\\r")
      '\t' => out.write_string("\\t")
      _ => out.write_char(ch)
    }
  }
  out.to_string()
}