///|
/// A small parameter that both backends need to reach.
///
/// The CPU kernels read `values` directly — going through the arena would mean
/// an FFI call per element — while the Metal kernels address the same numbers
/// in the arena at `offset`. Norm gains, biases and the type embedding total a
/// few hundred kilobytes, so holding both is cheaper than either alternative.
priv struct Param {
  values : FixedArray[Float]
  offset : Int64
}

///|
/// Plans the float32 arena, then fills it.
///
/// Every parameter is reserved up front so the arena can be allocated as one
/// block; `finish` then decodes each tensor from float16 into its slot.
priv struct Loader {
  file : @safetensors.File
  planned : Array[@safetensors.TensorInfo]
  targets : Array[Int64]
  mut total : Int64
}

///|
fn Loader::new(file : @safetensors.File) -> Loader {
  { file, planned: [], targets: [], total: 0L, }
}

///|
/// Reserve arena space for a tensor, returning the element offset it will live
/// at. The shape is checked here so a mismatched checkpoint fails before any
/// allocation happens.
fn Loader::reserve(
  self : Loader,
  name : String,
  shape : Array[Int],
) -> Int64 raise @error.LayaError {
  let info = self.file.info_shaped(name, shape)
  let offset = self.total
  self.planned.push(info)
  self.targets.push(offset)
  self.total += info.elements()
  offset
}

///|
/// Reserve a vector and also read it into a MoonBit array for the CPU path.
fn Loader::param(
  self : Loader,
  name : String,
  length : Int,
) -> Param raise @error.LayaError {
  let offset = self.reserve(name, [length])
  let values = @tensor.zeros(length)
  self.file.read_into_array(self.file.info(name), values)
  { values, offset, }
}

///|
/// Same, for a matrix the CPU path also reads elementwise.
fn Loader::param_matrix(
  self : Loader,
  name : String,
  rows : Int,
  cols : Int,
) -> Param raise @error.LayaError {
  let offset = self.reserve(name, [rows, cols])
  let values = @tensor.zeros(rows * cols)
  self.file.read_into_array(self.file.info(name), values)
  { values, offset, }
}

///|
fn Loader::finish(self : Loader) -> @tensor.Arena raise @error.LayaError {
  let arena = @tensor.Arena::new(self.total)
  for index, info in self.planned {
    self.file.read_into_arena(info, arena, self.targets[index])
  }
  arena
}

///|
/// The parameter blocks of one ModernBERT encoder layer, plus the attention
/// pattern it was configured with.
priv struct EncoderLayer {
  /// `None` on layer 0, where ModernBERT uses an identity in place of the
  /// pre-attention norm because the embedding norm already ran.
  attn_norm : Param?
  wqkv : Int64
  attn_wo : Int64
  mlp_norm : Param
  mlp_wi : Int64
  mlp_wo : Int64
  attention : Attention
  rope_base : Double
}

///|
/// One layer of the decision head — a `torch.nn.TransformerEncoderLayer` in the
/// original, so norms carry biases and the feed-forward uses ReLU.
priv struct HeadLayer {
  norm1_weight : Param
  norm1_bias : Param
  norm2_weight : Param
  norm2_bias : Param
  in_proj : Int64
  in_proj_bias : Int64
  out_proj : Int64
  out_proj_bias : Int64
  linear1 : Int64
  linear1_bias : Int64
  linear2 : Int64
  linear2_bias : Int64
}

///|
/// A loaded Laya checkpoint, ready to run.
struct Model {
  encoder : EncoderConfig
  agent : AgentConfig
  arena : @tensor.Arena
  file : @safetensors.File
  tok_embeddings : @safetensors.TensorInfo
  embed_norm : Param
  layers : Array[EncoderLayer]
  final_norm : Param
  /// `[3, hidden]`: one learned offset per question type, added to every
  /// position before the decision head runs.
  type_emb : Param
  head : Array[HeadLayer]
  head_heads : Int
  head_dim : Int
  scorer_norm_weight : Param
  scorer_norm_bias : Param
  scorer_dense : Int64
  scorer_dense_bias : Int64
  scorer_out : Int64
  scorer_out_bias : Int64
  act_hidden : Int
  act_in : Int64
  act_in_bias : Int64
  act_out : Int64
  act_out_bias : Int64
  act_outputs : Int
  /// Set by `Model::use_metal`; `None` means every pass runs on the CPU.
  mut metal : MetalState?
}

///|
/// The encoder configuration this model was built from.
pub fn Model::encoder_config(self : Model) -> EncoderConfig {
  self.encoder
}

///|
/// The calibration and prompt-budget configuration.
pub fn Model::agent_config(self : Model) -> AgentConfig {
  self.agent
}

///|
/// Load a checkpoint from an already-parsed pair of configs plus its weights.
pub fn Model::load(
  weights_path : String,
  encoder : EncoderConfig,
  agent : AgentConfig,
) -> Model raise @error.LayaError {
  let file = @safetensors.File::open(weights_path)
  let hidden = encoder.hidden_size
  let inner = encoder.intermediate_size
  let loader = Loader::new(file)

  // Encoder body. Wqkv packs query, key and value into one [3H, H] matrix and
  // Wi packs the GeGLU value and gate halves into one [2I, H] matrix.
  let layers = Array(capacity=encoder.num_hidden_layers)
  for index in 0..