graph

////
//// Here's a handy index you can use to browse through the various graph
//// functions.
////
//// | operation kind | functions |
//// |---|---|
//// | creating graphs | [`new`](#new) |
//// | turning graphs into lists | [`nodes`](#nodes) |
//// | querying a graph | [`size`](#size), [`has_node`](#has_node), [`has_edge`](#has_edge), [`get_context`](#get_context), [`match`](#match) |
//// | adding/removing elements from a graph | [`insert_node`](#insert_node), [`insert_directed_edge`](#insert_directed_edge), [`insert_undirected_edge`](#insert_undirected_edge), [`remove_node`](#remove_node), [`remove_directed_edge`](#remove_directed_edge), [`remove_undirected_edge`](#remove_undirected_edge) |
//// | transforming graphs | [`fold`](#fold), [`reverse`](#reverse), [`map_contexts`](#map_contexts), [`map_values`](#map_values), [`map_labels`](#map_labels), [`reverse_edges`](#reverse_edges), [`to_directed`](#to_directed) |
////

import gleam/dict.{type Dict}
import gleam/result

// --- THE GRAPH TYPE ----------------------------------------------------------

/// The direction of a directed graph.
///
pub type Directed

/// The direction of an undirected graph.
///
pub type Undirected

/// A directed or undirected graph. A graph is made up of nodes and edges
/// connecting them: each node holds a `value` and each edge has a `label`.
///
/// The graph also carries along its `direction` (either `Directed` or
/// `Undirected`) in its type so that it's impossible to mix up `Directed` and
/// `Undirected` graphs inadvertently.
///
pub opaque type Graph(direction, value, label) {
  Graph(Dict(Int, Context(value, label)))
}

/// A node making up a graph. Every node is identified by a number and can hold
/// an arbitrary value.
///
pub type Node(value) {
  Node(id: Int, value: value)
}

/// The context associated with a node in a graph: it contains the node itself
/// and all the incoming and outgoing edges. Edges are stored in a dict going
/// from neighbour's id to the edge label.
///
pub type Context(value, label) {
  Context(
    incoming: Dict(Int, label),
    node: Node(value),
    outgoing: Dict(Int, label),
  )
}

// --- CREATING GRAPHS ---------------------------------------------------------

/// Creates a new empty graph.
///
/// ## Examples
///
/// ```gleam
/// nodes(new())
/// // -> []
/// ```
///
pub fn new() -> Graph(direction, value, label) {
  Graph(dict.new())
}

// --- TURNING GRAPHS INTO LISTS -----------------------------------------------

/// Returns a list of all the nodes contained in the graph.
///
/// ## Examples
///
/// ```gleam
/// new() |> nodes
/// // -> []
/// ```
///
/// ```gleam
/// new() |> insert_node(Node(1, "a node")) |> nodes
/// // -> [Node(1, "a node")]
/// ```
///
pub fn nodes(graph: Graph(direction, value, label)) -> List(Node(value)) {
  let Graph(graph) = graph
  use acc, _node_id, Context(node: node, ..) <- dict.fold(over: graph, from: [])
  [node, ..acc]
}

// --- QUERYING A GRAPH --------------------------------------------------------

/// Returns the number of nodes of the graph.
///
/// ## Examples
///
/// ```gleam
/// new() |> size
/// // -> 0
/// ```
///
/// ```gleam
/// new() |> insert_node(Node(1, "a node")) |> size
/// // -> 1
/// ```
///
pub fn size(graph: Graph(direction, value, label)) -> Int {
  let Graph(graph) = graph
  dict.size(graph)
}

/// Returns `True` if the graph contains a node with the given id.
///
/// ## Examples
///
/// ```gleam
/// let my_graph = new() |> insert_node(Node(1, "a node"))
///
/// my_graph |> has_node(1)
/// // -> True
///
/// my_graph |> has_node(2)
/// // -> False
/// ```
///
pub fn has_node(graph: Graph(direction, value, label), node_id: Int) -> Bool {
  let Graph(graph) = graph
  dict.has_key(graph, node_id)
}

/// Returns `True` if the graph has an edge connecting the two nodes with the
/// given ids.
///
/// ## Examples
///
/// ```gleam
/// let my_graph =
///   new()
///   |> insert_node(Node(1, "a node"))
///   |> insert_node(Node(2, "other node"))
///   |> insert_directed_edge("edge label", from: 1, to: 2)
///
/// my_graph |> has_edge(from: 1, to: 2)
/// // -> True
///
/// my_graph |> has_edge(from: 2, to: 1)
/// // -> False
/// ```
///
pub fn has_edge(
  graph: Graph(direction, value, label),
  from source: Int,
  to destination: Int,
) -> Bool {
  case get_context(graph, source) {
    Ok(Context(outgoing: outgoing, ..)) -> dict.has_key(outgoing, destination)
    Error(_) -> False
  }
}

/// Returns the context associated with the node with the given id, if present.
/// Otherwise returns `Error(Nil)`.
///
/// ## Examples
///
/// ```gleam
/// new() |> get(1)
/// // -> Error(Nil)
/// ```
///
/// ```gleam
/// new() |> insert_node(Node(1, "a node")) |> get_context(of: 1)
/// // -> Ok(Context(node: Node(1, "a node"), ..))
/// ```
///
pub fn get_context(
  graph: Graph(direction, value, label),
  of node: Int,
) -> Result(Context(value, label), Nil) {
  let Graph(graph) = graph
  dict.get(graph, node)
}

/// If the graph contains a node with the given id, returns a tuple containing
/// the context of that node (with all edges looping back to itself removed) and
/// the "remaining" graph: that is, the original graph where that node has been
/// removed.
///
pub fn match(
  graph: Graph(direction, value, label),
  node_id: Int,
) -> Result(#(Context(value, label), Graph(direction, value, label)), Nil) {
  use Context(incoming, node, outgoing) <- result.try(get_context(
    graph,
    node_id,
  ))
  let rest = remove_node(graph, node_id)
  let new_incoming = dict.delete(incoming, node_id)
  let new_outgoing = dict.delete(outgoing, node_id)
  Ok(#(Context(new_incoming, node, new_outgoing), rest))
}

// --- ADDING/REMOVING ELEMENTS FROM A GRAPH -----------------------------------

/// Adds a node to the given graph.
/// If the graph already contains a node with the same id, that will be replaced
/// by the new one.
/// The newly added node won't be connected to any existing node.
///
/// ## Examples
///
/// ```gleam
/// new() |> insert_node(Node(1, "a node")) |> nodes
/// // -> [Node(1, "a node")]
/// ```
///
pub fn insert_node(
  graph: Graph(direction, value, label),
  node: Node(value),
) -> Graph(direction, value, label) {
  let Graph(graph) = graph
  let empty_context = Context(dict.new(), node, dict.new())
  let new_graph = dict.insert(graph, node.id, empty_context)
  Graph(new_graph)
}

/// Adds an edge connecting two nodes in a directed graph.
///
/// ```gleam
/// let my_graph =
///   new()
///   |> insert_node(Node(1, "a node"))
///   |> insert_node(Node(2, "other node"))
///   |> insert_directed_edge("edge label", from: 1, to: 2)
///
/// my_graph |> has_edge(from: 1, to: 2)
/// // -> True
///
/// my_graph |> has_edge(from: 2, to: 1)
/// // -> False
/// ```
///
pub fn insert_directed_edge(
  graph: Graph(Directed, value, label),
  labelled label: label,
  from source: Int,
  to destination: Int,
) -> Graph(Directed, value, label) {
  graph
  |> update_context(of: source, with: add_outgoing_edge(_, destination, label))
  |> update_context(of: destination, with: add_incoming_edge(_, source, label))
}

/// Adds an edge connecting two nodes in an undirected graph.
///
/// ## Examples
///
/// ```gleam
/// let my_graph =
///   new()
///   |> insert_node(Node(1, "a node"))
///   |> insert_node(Node(2, "other node"))
///   |> insert_undirected_edge("edge label", between: 1, and: 2)
///
/// my_graph |> has_edge(from: 1, to: 2)
/// // -> True
///
/// my_graph |> has_edge(from: 2, to: 1)
/// // -> True
/// ```
pub fn insert_undirected_edge(
  graph: Graph(Undirected, value, label),
  labelled label: label,
  between one: Int,
  and other: Int,
) -> Graph(Undirected, value, label) {
  graph
  |> update_context(of: one, with: fn(context) {
    add_outgoing_edge(context, other, label)
    |> add_incoming_edge(other, label)
  })
  |> update_context(of: other, with: fn(context) {
    add_outgoing_edge(context, one, label)
    |> add_incoming_edge(one, label)
  })
}

fn update_context(
  in graph: Graph(direction, value, label),
  of node: Int,
  with fun: fn(Context(value, label)) -> Context(value, label),
) -> Graph(direction, value, label) {
  let Graph(graph) = graph
  case dict.get(graph, node) {
    Ok(context) -> Graph(dict.insert(graph, node, fun(context)))
    Error(_) -> Graph(graph)
  }
}

fn add_outgoing_edge(
  context: Context(value, label),
  to node: Int,
  labelled label: label,
) -> Context(value, label) {
  let Context(outgoing: outgoing, ..) = context
  Context(..context, outgoing: dict.insert(outgoing, node, label))
}

fn remove_outgoing_edge(
  context: Context(value, label),
  to node: Int,
) -> Context(value, label) {
  let Context(outgoing: outgoing, ..) = context
  Context(..context, outgoing: dict.delete(outgoing, node))
}

fn add_incoming_edge(
  context: Context(value, label),
  from node: Int,
  labelled label: label,
) -> Context(value, label) {
  let Context(incoming: incoming, ..) = context
  Context(..context, incoming: dict.insert(incoming, node, label))
}

fn remove_incoming_edge(
  context: Context(value, label),
  from node: Int,
) -> Context(value, label) {
  let Context(incoming: incoming, ..) = context
  Context(..context, incoming: dict.delete(incoming, node))
}

/// Removes a node with the given id from the graph. If there's no node with the
/// given id it does nothing.
///
pub fn remove_node(
  graph: Graph(direction, value, label),
  node_id: Int,
) -> Graph(direction, value, label) {
  case graph, get_context(graph, node_id) {
    _, Error(_) -> graph
    Graph(graph), Ok(Context(incoming, _, outgoing)) ->
      dict.delete(graph, node_id)
      |> remove_incoming_occurrences(of: node_id, from: outgoing)
      |> remove_outgoing_occurrences(of: node_id, from: incoming)
      |> Graph
  }
}

fn remove_incoming_occurrences(
  in graph: Dict(Int, Context(value, label)),
  of node: Int,
  from nodes: Dict(Int, a),
) -> Dict(Int, Context(value, label)) {
  use context, _ <- dict_map_shared_keys(graph, with: nodes)
  let Context(incoming: incoming, ..) = context
  Context(..context, incoming: dict.delete(incoming, node))
}

fn remove_outgoing_occurrences(
  in graph: Dict(Int, Context(value, label)),
  of node: Int,
  from nodes: Dict(Int, a),
) -> Dict(Int, Context(value, label)) {
  use context, _ <- dict_map_shared_keys(graph, with: nodes)
  let Context(outgoing: outgoing, ..) = context
  Context(..context, outgoing: dict.delete(outgoing, node))
}

/// Removes a directed edge connecting two nodes from a graph.
///
pub fn remove_directed_edge(
  graph: Graph(Directed, value, label),
  from source: Int,
  to destination: Int,
) -> Graph(Directed, value, label) {
  graph
  |> update_context(of: source, with: remove_outgoing_edge(_, to: destination))
  |> update_context(of: destination, with: remove_incoming_edge(_, from: source))
}

/// Removes an undirected edge connecting two nodes from a graph.
///
pub fn remove_undirected_edge(
  graph: Graph(Undirected, value, label),
  between one: Int,
  and other: Int,
) -> Graph(Undirected, value, label) {
  graph
  |> update_context(of: one, with: fn(context) {
    remove_outgoing_edge(context, to: other)
    |> remove_incoming_edge(from: other)
  })
  |> update_context(of: other, with: fn(context) {
    remove_outgoing_edge(context, to: one)
    |> remove_incoming_edge(from: one)
  })
}

// --- TRANSFORMING GRAPHS -----------------------------------------------------

/// Reduces the given graph into a single value by applying function to all its
/// contexts, one after the other.
///
/// > 🚨 Graph's contexts are not sorted in any way so your folding function
/// > should never rely on any accidental order the contexts might have.
///
/// ## Examples
///
/// ```gleam
/// // The size function could be implemented using a fold.
/// // The real implementation is more efficient because it doesn't have to
/// // traverse all contexts!
/// pub fn size(graph) {
///   fold(
///     over: graph,
///     from: 0,
///     with: fn(size, _context) { size + 1 },
///   )
/// }
/// ```
///
pub fn fold(
  over graph: Graph(direction, value, label),
  from initial: b,
  with fun: fn(b, Context(value, label)) -> b,
) -> b {
  let Graph(graph) = graph
  use acc, _node_id, context <- dict.fold(over: graph, from: initial)
  fun(acc, context)
}

/// Transform the contexts associated with each node.
///
/// > This function can add and remove arbitrary edges from the graph by
/// > updating the `incoming` and `outgoing` edges of a context.
/// > So we can't assume the final graph will still be `Undirected`, that's why
/// > it is always treated as a `Directed` one.
///
/// ## Examples
///
/// ```gleam
/// // The reverse function can be implemented with `map_contexts`
/// pub fn reverse(graph) {
///   map_contexts(in: graph, with: fn(context) {
///     Context(
///       ..context,
///       incoming: context.outgoing,
///       outgoing: context.incoming,
///     )
///   })
/// }
/// ```
///
pub fn map_contexts(
  in graph: Graph(direction, value, label),
  with fun: fn(Context(value, label)) -> Context(value, label),
) -> Graph(Directed, value, label) {
  use acc, context <- fold(over: graph, from: new())
  insert_context(acc, fun(context))
}

fn insert_context(
  graph: Graph(direction, value, label),
  context: Context(value, label),
) -> Graph(Directed, value, label) {
  let Graph(graph) = graph
  let new_graph = dict.insert(graph, context.node.id, context)
  Graph(new_graph)
}

/// Transforms the values of all the graph's nodes using the given function.
///
/// ## Examples
///
/// ```gleam
/// new()
/// |> insert_node(Node(1, "a node"))
/// |> map_nodes(fn(value) { value <> "!" })
/// |> nodes
/// // -> [Node(1, "my node!")]
/// ```
///
pub fn map_values(
  in graph: Graph(direction, value, label),
  with fun: fn(value) -> new_value,
) -> Graph(direction, new_value, label) {
  let Graph(graph) = graph
  // Since this function doesn't change the graph's topology I'm not
  // implementing it with a `graph.fold` or a `graph.map_contexts`, it would
  // increase code reuse but would rebuild a new graph each time by adding
  // each context one by one.
  Graph({
    use _node_id, context <- dict.map_values(graph)
    let Context(incoming, Node(id, value), outgoing) = context
    Context(incoming, Node(id, fun(value)), outgoing)
  })
}

/// Transforms the labels of all the graph's edges using the given function.
///
/// ## Examples
///
/// ```
/// new()
/// |> insert_node(Node(1, "a node"))
/// |> insert_undirected_edge(UndirectedEdge(1, 1, "label"))
/// |> map_labels(fn(label) { label <> "!" })
/// |> labels
/// // -> ["label!"]
/// ```
///
pub fn map_labels(
  in graph: Graph(direction, value, label),
  with fun: fn(label) -> new_label,
) -> Graph(direction, value, new_label) {
  // Since this function doesn't change the graph's topology I'm not
  // implementing it with a `graph.fold` or a `graph.map_contexts`, it would
  // increase code reuse but would rebuild a new graph each time by adding
  // each context one by one.
  let Graph(graph) = graph
  Graph({
    use _node_id, context <- dict.map_values(graph)
    let Context(incoming, node, outgoing) = context
    let new_incoming = dict.map_values(incoming, fn(_id, label) { fun(label) })
    let new_outgoing = dict.map_values(outgoing, fn(_id, label) { fun(label) })
    Context(new_incoming, node, new_outgoing)
  })
}

/// Flips the direction of every edge in the graph. All incoming edges will
/// become outgoing and vice-versa.
///
pub fn reverse_edges(
  graph: Graph(Directed, value, label),
) -> Graph(Directed, value, label) {
  // Since this function doesn't change the graph's structure I'm not
  // implementing it with a `graph.fold` or a `graph.map_contexts`, it would
  // increase code reuse but would rebuild a new graph each time by adding
  // each context one by one
  let Graph(graph) = graph
  Graph({
    use _node_id, context <- dict.map_values(graph)
    let Context(incoming, node, outgoing) = context
    Context(outgoing, node, incoming)
  })
}

/// Turns an undirected graph into a directed one. Every edge connecting two
/// nodes in the original graph will be considered as a pair of edges connecting
/// the nodes going in both directions.
///
pub fn to_directed(
  graph: Graph(Undirected, value, label),
) -> Graph(Directed, value, label) {
  let Graph(graph) = graph
  Graph(graph)
}

// --- DICT UTILITY FUNCTIONS --------------------------------------------------

fn dict_map_shared_keys(
  in one: Dict(k, a),
  with other: Dict(k, b),
  using fun: fn(a, b) -> a,
) -> Dict(k, a) {
  use one, key, other_value <- dict.fold(over: other, from: one)
  case dict.get(one, key) {
    Ok(one_value) -> dict.insert(one, key, fun(one_value, other_value))
    Error(_) -> one
  }
}