///|
pub(all) struct EntityInput {
  label : String
  kind : PhiKind
  start : Int
  end : Int
  confidence : Int
  source : String
} derive(Debug, Eq)

///|
pub(all) struct EntityBatch {
  source : String
  entities : Array[EntityInput]
  invalid_count : Int
} derive(Debug, Eq)

///|
pub fn entity_input(
  label~ : String,
  kind~ : PhiKind,
  start~ : Int,
  end~ : Int,
  confidence~ : Int,
  source~ : String,
) -> EntityInput {
  { label, kind, start, end, confidence, source }
}

///|
pub fn external_entity(item : EntityInput) -> ExternalEntity {
  {
    kind: item.kind,
    label: item.label,
    start: item.start,
    end: item.end,
    confidence: item.confidence,
    source: item.source,
  }
}

///|
pub fn entity_batch(
  source : String,
  entities : Array[EntityInput],
) -> EntityBatch {
  { source, entities, invalid_count: 0 }
}

///|
pub fn valid_entity_input(input : String, item : EntityInput) -> Bool {
  item.start >= 0 &&
  item.end > item.start &&
  item.end <= input.length() &&
  item.confidence >= 0 &&
  item.confidence <= 100
}

///|
pub fn normalize_entity_batch(
  input : String,
  batch : EntityBatch,
) -> EntityBatch {
  let valid = []
  let mut invalid = 0
  for item in batch.entities {
    if valid_entity_input(input, item) {
      valid.push({
        ..item,
        label: trim_ascii_space(item.label),
        source: trim_ascii_space(item.source),
      })
    } else {
      invalid += 1
    }
  }
  { ..batch, entities: valid, invalid_count: invalid }
}

///|
pub fn entity_batch_to_external(batch : EntityBatch) -> Array[ExternalEntity] {
  batch.entities.map(external_entity)
}

///|
pub fn merge_entity_batch(
  input : String,
  findings : Array[Finding],
  batch : EntityBatch,
  mode : ReplacementMode,
) -> Array[Finding] {
  let normalized = normalize_entity_batch(input, batch)
  merge_external_entities(
    input,
    findings,
    entity_batch_to_external(normalized),
    mode~,
  )
}

///|
pub fn ner_merge_and_redact(
  input : String,
  batch : EntityBatch,
  mode : ReplacementMode,
) -> DeidResult {
  let builtin = scan(input) catch { _ => [] }
  let findings = merge_entity_batch(input, builtin, batch, mode)
  let (text, offsets) = apply_findings(input, findings)
  {
    text,
    findings,
    offsets,
    audit: build_audit(input, text, findings, offsets),
  }
}

///|
pub fn entities_for_source(
  batch : EntityBatch,
  source : String,
) -> Array[EntityInput] {
  batch.entities.filter(fn(item) { item.source == source })
}

///|
pub fn entities_for_kind(
  batch : EntityBatch,
  kind : PhiKind,
) -> Array[EntityInput] {
  batch.entities.filter(fn(item) { item.kind == kind })
}

///|
pub fn entity_source_counts(batch : EntityBatch) -> Map[String, Int] {
  let counts : Map[String, Int] = Map([])
  for item in batch.entities {
    counts[item.source] = counts.get_or_default(item.source, 0) + 1
  }
  counts
}

///|
pub fn entity_kind_counts(batch : EntityBatch) -> Map[String, Int] {
  let counts : Map[String, Int] = Map([])
  for item in batch.entities {
    let key = phi_kind_name(item.kind)
    counts[key] = counts.get_or_default(key, 0) + 1
  }
  counts
}

///|
pub fn entity_confidence_average(batch : EntityBatch) -> Float {
  if batch.entities.is_empty() {
    0.0
  } else {
    Float::from_int(
      batch.entities.fold(init=0, (sum, item) => sum + item.confidence),
    ) /
    Float::from_int(batch.entities.length())
  }
}

///|
pub fn entity_batch_summary(batch : EntityBatch) -> String {
  [
    "source=\{batch.source}",
    "entities=\{batch.entities.length()}",
    "invalid=\{batch.invalid_count}",
    "sources=\{entity_source_counts(batch).length()}",
    "kinds=\{entity_kind_counts(batch).length()}",
    "average_confidence=\{entity_confidence_average(batch)}",
  ].join("\n")
}

///|
pub fn external_entities_csv(entities : Array[ExternalEntity]) -> String {
  let lines = ["source,label,kind,start,end,confidence"]
  for item in entities {
    lines.push(
      [
        csv_cell(item.source),
        csv_cell(item.label),
        csv_cell(phi_kind_name(item.kind)),
        "\{item.start}",
        "\{item.end}",
        "\{item.confidence}",
      ].join(","),
    )
  }
  lines.join("\n")
}

///|
pub fn entity_batch_from_tsv(text : String, source : String) -> EntityBatch {
  let entities : Array[EntityInput] = []
  let mut invalid = 0
  for line in split_lines_with_offsets(text) {
    let parts = line.text.split("\t").to_array()
    if parts.length() >= 5 {
      let start = decimal_value(parts[2].to_owned())
      let end = decimal_value(parts[3].to_owned())
      let confidence = decimal_value(parts[4].to_owned())
      let kind = match parts[1].to_owned() {
        "name" => PersonName
        "id" => IdNumber
        "phone" => Phone
        "email" => Email
        "date" => Date
        "address" => Address
        "medical_record" => MedicalRecord
        "insurance" => Insurance
        "organization" => Organization
        value => Custom(value)
      }
      let item : EntityInput = {
        label: parts[0].to_owned(),
        kind,
        start,
        end,
        confidence,
        source,
      }
      entities.push(item)
    } else if !line.text.trim().is_empty() {
      invalid += 1
    }
  }
  { source, entities, invalid_count: invalid }
}

///|
pub fn entity_batch_to_tsv(batch : EntityBatch) -> String {
  batch.entities
  .map(fn(item) {
    [
      item.label,
      phi_kind_name(item.kind),
      "\{item.start}",
      "\{item.end}",
      "\{item.confidence}",
      item.source,
    ].join("\t")
  })
  .join("\n")
}

///|
pub fn entity_spans(batch : EntityBatch) -> Array[Span] {
  batch.entities.map(fn(item) { { start: item.start, end: item.end } })
}

///|
pub fn entity_batch_overlap_count(batch : EntityBatch) -> Int {
  let spans = entity_spans(batch)
  let mut count = 0
  for i in 0.. Bool {
  normalize_entity_batch(input, batch).invalid_count == 0 &&
  entity_batch_overlap_count(batch) == 0
}