///|
/// A mention is a candidate occurrence that can be linked across documents.
pub(all) struct EntityMention {
  document_id : String
  finding_id : String
  kind : PhiKind
  value : String
  normalized : String
  start : Int
  end : Int
  confidence : Int
  source : String
} derive(Debug, Eq)

///|
pub(all) struct EntityCluster {
  cluster_id : String
  kind : PhiKind
  canonical : String
  mentions : Array[EntityMention]
  confidence : Int
} derive(Debug)

///|
pub(all) struct EntityLinker {
  mut clusters : Map[String, EntityCluster]
  mut mention_to_cluster : Map[String, String]
  salt : String
} derive(Debug)

///|
pub fn entity_mention(
  document_id : String,
  finding : Finding,
  source : String,
) -> EntityMention {
  {
    document_id,
    finding_id: finding.id,
    kind: finding.kind,
    value: finding.text,
    normalized: normalize_entity_value(finding.text, finding.kind),
    start: finding.start,
    end: finding.end,
    confidence: finding.confidence,
    source,
  }
}

///|
pub fn normalize_entity_value(value : String, kind : PhiKind) -> String {
  let normalized = normalize_for_matching(value)
  match kind {
    Email => normalized.to_lower()
    Phone | IdNumber | MedicalRecord | Insurance =>
      keep_ascii_letters_and_digits(normalized).to_lower()
    PersonName | Address | Organization => normalized.to_lower()
    Date => normalized
    Custom(_) => normalized
  }
}

///|
pub fn entity_mention_key(mention : EntityMention) -> String {
  phi_kind_name(mention.kind) + "\u{1f}" + mention.normalized
}

///|
pub fn EntityLinker::new(salt : String) -> EntityLinker {
  { clusters: Map([]), mention_to_cluster: Map([]), salt }
}

///|
pub fn EntityLinker::default() -> EntityLinker {
  EntityLinker::new("moonbit-entity-linker")
}

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

///|
pub fn EntityLinker::mention_count(self : EntityLinker) -> Int {
  self.clusters
  .values()
  .fold(init=0, (sum, cluster) => sum + cluster.mentions.length())
}

///|
fn cluster_id(linker : EntityLinker, mention : EntityMention) -> String {
  "cluster-" + stable_hash(linker.salt + "\u{1f}" + entity_mention_key(mention))
}

///|
fn cluster_confidence(mentions : Array[EntityMention]) -> Int {
  if mentions.is_empty() {
    0
  } else {
    mentions.fold(init=0, (sum, item) => sum + item.confidence) /
    mentions.length()
  }
}

///|
pub fn EntityLinker::link(
  self : EntityLinker,
  mention : EntityMention,
) -> String {
  let id = cluster_id(self, mention)
  match self.clusters.get(id) {
    Some(cluster) => {
      let mentions = cluster.mentions.copy()
      mentions.push(mention)
      self.clusters[id] = {
        ..cluster,
        mentions,
        confidence: cluster_confidence(mentions),
      }
    }
    None =>
      self.clusters[id] = {
        cluster_id: id,
        kind: mention.kind,
        canonical: mention.normalized,
        mentions: [mention],
        confidence: mention.confidence,
      }
  }
  self.mention_to_cluster[mention.finding_id] = id
  id
}

///|
pub fn EntityLinker::link_finding(
  self : EntityLinker,
  document_id : String,
  finding : Finding,
  source : String,
) -> String {
  self.link(entity_mention(document_id, finding, source))
}

///|
pub fn EntityLinker::link_findings(
  self : EntityLinker,
  document_id : String,
  findings : Array[Finding],
  source : String,
) -> Array[String] {
  findings.map(fn(finding) { self.link_finding(document_id, finding, source) })
}

///|
pub fn EntityLinker::cluster(
  self : EntityLinker,
  cluster_id : String,
) -> EntityCluster? {
  self.clusters.get(cluster_id)
}

///|
pub fn EntityLinker::cluster_for_mention(
  self : EntityLinker,
  finding_id : String,
) -> EntityCluster? {
  match self.mention_to_cluster.get(finding_id) {
    Some(id) => self.clusters.get(id)
    None => None
  }
}

///|
pub fn EntityLinker::all_clusters(self : EntityLinker) -> Array[EntityCluster] {
  self.clusters.values().to_array()
}

///|
pub fn EntityLinker::clusters_for_kind(
  self : EntityLinker,
  kind : PhiKind,
) -> Array[EntityCluster] {
  self.clusters.values().filter(fn(cluster) { cluster.kind == kind }).to_array()
}

///|
pub fn EntityLinker::clusters_for_document(
  self : EntityLinker,
  document_id : String,
) -> Array[EntityCluster] {
  self.clusters
  .values()
  .filter(fn(cluster) {
    cluster.mentions.any(fn(mention) { mention.document_id == document_id })
  })
  .to_array()
}

///|
pub fn EntityLinker::documents(self : EntityLinker) -> Array[String] {
  let result : Map[String, Unit] = Map([])
  for cluster in self.clusters.values() {
    for mention in cluster.mentions {
      result[mention.document_id] = ()
    }
  }
  result.keys().to_array()
}

///|
pub fn EntityLinker::mention_ids(self : EntityLinker) -> Array[String] {
  self.mention_to_cluster.keys().to_array()
}

///|
pub fn EntityLinker::merge(
  self : EntityLinker,
  other : EntityLinker,
) -> EntityLinker {
  for cluster in other.all_clusters() {
    for mention in cluster.mentions {
      ignore(self.link(mention))
    }
  }
  self
}

///|
pub fn EntityLinker::collision_count(self : EntityLinker) -> Int {
  let mut count = 0
  for cluster in self.clusters.values() {
    let values = cluster.mentions.map(fn(mention) { mention.value })
    let unique : Map[String, Unit] = Map([])
    for value in values {
      unique[value] = ()
    }
    if unique.length() > 1 {
      count += 1
    }
  }
  count
}

///|
pub fn EntityLinker::kind_counts(self : EntityLinker) -> Map[String, Int] {
  let counts : Map[String, Int] = Map([])
  for cluster in self.clusters.values() {
    let key = phi_kind_name(cluster.kind)
    counts[key] = counts.get_or_default(key, 0) + 1
  }
  counts
}

///|
pub fn EntityLinker::document_counts(self : EntityLinker) -> Map[String, Int] {
  let counts : Map[String, Int] = Map([])
  for cluster in self.clusters.values() {
    for mention in cluster.mentions {
      counts[mention.document_id] = counts.get_or_default(
          mention.document_id,
          0,
        ) +
        1
    }
  }
  counts
}

///|
pub fn EntityLinker::high_confidence_clusters(
  self : EntityLinker,
  threshold : Int,
) -> Array[EntityCluster] {
  self.clusters
  .values()
  .filter(fn(cluster) { cluster.confidence >= threshold })
  .to_array()
}

///|
pub fn EntityLinker::low_confidence_clusters(
  self : EntityLinker,
  threshold : Int,
) -> Array[EntityCluster] {
  self.clusters
  .values()
  .filter(fn(cluster) { cluster.confidence < threshold })
  .to_array()
}

///|
pub fn entity_cluster_summary(cluster : EntityCluster) -> String {
  [
    "id=\{cluster.cluster_id}",
    "kind=\{phi_kind_name(cluster.kind)}",
    "canonical=\{cluster.canonical}",
    "mentions=\{cluster.mentions.length()}",
    "confidence=\{cluster.confidence}",
  ].join("\n")
}

///|
pub fn entity_cluster_json(cluster : EntityCluster) -> String {
  let mentions = cluster.mentions
    .map(fn(item) {
      "{" +
      "\"document_id\":\{json_escape(item.document_id)}," +
      "\"finding_id\":\{json_escape(item.finding_id)}," +
      "\"kind\":\{json_escape(phi_kind_name(item.kind))}," +
      "\"start\":\{item.start},\"end\":\{item.end}," +
      "\"confidence\":\{item.confidence}," +
      "\"source\":\{json_escape(item.source)}" +
      "}"
    })
    .join(",")
  "{" +
  "\"cluster_id\":\{json_escape(cluster.cluster_id)}," +
  "\"kind\":\{json_escape(phi_kind_name(cluster.kind))}," +
  "\"canonical\":\{json_escape(cluster.canonical)}," +
  "\"confidence\":\{cluster.confidence}," +
  "\"mentions\":[" +
  mentions +
  "]}"
}

///|
pub fn EntityLinker::to_json(self : EntityLinker) -> String {
  "[" + self.all_clusters().map(entity_cluster_json).join(",") + "]"
}

///|
pub fn EntityLinker::checksum(self : EntityLinker) -> String {
  let values = self
    .all_clusters()
    .map(fn(cluster) {
      cluster.cluster_id +
      ":" +
      cluster.canonical +
      ":\{cluster.mentions.length()}"
    })
  values.sort()
  stable_hash(values.join("\n"))
}

///|
pub fn EntityLinker::summary(self : EntityLinker) -> String {
  [
    "clusters=\{self.size()}",
    "mentions=\{self.mention_count()}",
    "documents=\{self.documents().length()}",
    "collisions=\{self.collision_count()}",
    "checksum=\{self.checksum()}",
  ].join("\n")
}

///|
pub fn cluster_representative(cluster : EntityCluster) -> EntityMention? {
  if cluster.mentions.is_empty() {
    None
  } else {
    Some(
      cluster.mentions.fold(init=cluster.mentions[0], (best, item) => {
        if item.confidence > best.confidence {
          item
        } else {
          best
        }
      }),
    )
  }
}

///|
pub fn cluster_span_count(cluster : EntityCluster) -> Int {
  cluster.mentions.map(fn(item) { item.end - item.start }).length()
}

///|
pub fn cluster_is_cross_document(cluster : EntityCluster) -> Bool {
  let documents : Map[String, Unit] = Map([])
  for mention in cluster.mentions {
    documents[mention.document_id] = ()
  }
  documents.length() > 1
}

///|
pub fn cluster_source_count(cluster : EntityCluster) -> Int {
  let sources : Map[String, Unit] = Map([])
  for mention in cluster.mentions {
    sources[mention.source] = ()
  }
  sources.length()
}

///|
pub fn cluster_has_external_source(cluster : EntityCluster) -> Bool {
  cluster.mentions.any(fn(mention) { mention.source != "builtin" })
}

///|
pub fn cluster_to_pseudonym(
  cluster : EntityCluster,
  vault : PseudonymVault,
) -> String {
  vault.lookup_or_create(cluster.canonical, cluster.kind)
}

///|
pub fn linked_findings(
  linker : EntityLinker,
  findings : Array[Finding],
) -> Map[String, String] {
  let result : Map[String, String] = Map([])
  for finding in findings {
    match linker.mention_to_cluster.get(finding.id) {
      Some(cluster) => result[finding.id] = cluster
      None => ()
    }
  }
  result
}

///|
pub fn entity_linking_is_stable(
  linker : EntityLinker,
  mention : EntityMention,
) -> Bool {
  let first = cluster_id(linker, mention)
  let second = cluster_id(linker, mention)
  first == second
}

///|
pub fn entity_linking_report(linker : EntityLinker) -> String {
  [
    linker.summary(),
    "kinds=\{map_to_json(linker.kind_counts())}",
    "documents=\{map_to_json(linker.document_counts())}",
  ].join("\n")
}