///|
fn is_spatial_type(name : String) -> Bool {
  match name.to_upper() {
    "IFCPROJECT"
    | "IFCSITE"
    | "IFCBUILDING"
    | "IFCBUILDINGSTOREY"
    | "IFCSPACE" => true
    _ => false
  }
}

///|
fn spatial_relationships(
  index : ModelIndex,
) -> (Map[Int, Array[Int]], Map[Int, Array[Int]]) {
  let children : Map[Int, Array[Int]] = Map([])
  let products : Map[Int, Array[Int]] = Map([])
  for relation in index.all() {
    let name = relation.name.to_upper()
    if name == "IFCRELAGGREGATES" {
      match optional_reference_argument(relation, 5) {
        Some(parent) =>
          for child in relation_objects(relation, 4) {
            children.update_or_default(parent, [child], values => {
              values.push(child)
              values
            })
          }
        None => ()
      }
    } else if name == "IFCRELCONTAINEDINSPATIALSTRUCTURE" {
      match optional_reference_argument(relation, 5) {
        Some(parent) =>
          for product in relation_objects(relation, 4) {
            products.update_or_default(parent, [product], values => {
              values.push(product)
              values
            })
          }
        None => ()
      }
    }
  }
  (children, products)
}

///|
fn spatial_name(index : ModelIndex, id : Int) -> String? {
  match index.entity(id) {
    Some(entity) => optional_string_argument(entity, 2)
    None => None
  }
}

///|
fn make_spatial_node(
  index : ModelIndex,
  id : Int,
  children_map : Map[Int, Array[Int]],
  products_map : Map[Int, Array[Int]],
  path : Array[Int],
) -> SpatialNode? {
  if path.contains(id) {
    return None
  }
  match index.entity(id) {
    Some(entity) if is_spatial_type(entity.name) => {
      let next_path = path.copy()
      next_path.push(id)
      let children = children_map
        .get(id)
        .unwrap_or([])
        .fold(init=Array::new(), (result, child_id) => {
          match
            make_spatial_node(
              index, child_id, children_map, products_map, next_path,
            ) {
            Some(child) => {
              result.push(child)
              result
            }
            None => result
          }
        })
      Some(SpatialNode::{
        id,
        ifc_type: entity.name,
        name: spatial_name(index, id),
        children,
        products: products_map.get(id).unwrap_or([]),
      })
    }
    _ => None
  }
}

///|
fn parent_ids(children_map : Map[Int, Array[Int]]) -> Map[Int, Bool] {
  let parents : Map[Int, Bool] = Map([])
  for _, child_ids in children_map {
    for id in child_ids {
      parents[id] = true
    }
  }
  parents
}

///|
pub fn spatial_roots(index : ModelIndex) -> Array[SpatialNode] {
  let (children_map, products_map) = spatial_relationships(index)
  let parents = parent_ids(children_map)
  index
  .all()
  .fold(init=Array::new(), (result, entity) => {
    if is_spatial_type(entity.name) && !parents.contains(entity.id) {
      match
        make_spatial_node(index, entity.id, children_map, products_map, []) {
        Some(root) => result.push(root)
        None => ()
      }
    }
    result
  })
}

///|
pub fn spatial_root(index : ModelIndex) -> SpatialNode? {
  spatial_roots(index).get(0)
}

///|
pub fn flatten_spatial(root : SpatialNode) -> Array[SpatialNode] {
  let result : Array[SpatialNode] = [root]
  for child in root.children {
    for nested in flatten_spatial(child) {
      result.push(nested)
    }
  }
  result
}

///|
pub fn spatial_product_ids(root : SpatialNode) -> Array[Int] {
  flatten_spatial(root).fold(init=Array::new(), (result, node) => {
    for id in node.products {
      result.push(id)
    }
    result
  })
}