arcana

if Code.ensure_loaded?(Leidenfold) do
  defmodule Arcana.Graph.CommunityDetector.Leiden do
    @moduledoc """
    Leiden algorithm implementation for community detection.

    Uses the Leidenfold library (Rust NIF) to detect communities in entity graphs.
    The Leiden algorithm is a refinement of the Louvain algorithm that
    guarantees well-connected communities.

    ## Installation

    Add `leidenfold` to your dependencies in `mix.exs`:

        defp deps do
          [
            {:arcana, "~> 1.2"},
            {:leidenfold, "~> 0.2"}
          ]
        end

    Precompiled binaries are available for macOS (Apple Silicon) and Linux (x86_64, ARM64).

    ## Usage

        detector = {Arcana.Graph.CommunityDetector.Leiden, resolution: 1.0}
        {:ok, communities} = CommunityDetector.detect(detector, entities, relationships)

    ## Options

      - `:resolution` - Controls community granularity (default: 1.0)
        Higher values produce smaller communities
      - `:objective` - Quality function to optimize (default: :cpm)
        Options: :cpm, :modularity, :rber, :rbc, :significance, :surprise
      - `:iterations` - Number of optimization iterations (default: 2)
      - `:seed` - Random seed for reproducibility (default: 0 = random)
      - `:min_size` - Minimum community size to include (default: 1)
        Set to 2 to exclude singleton communities
      - `:max_level` - Maximum hierarchy levels to generate (default: 1)
        Higher levels contain coarser communities built by aggregating lower levels

    """

    @behaviour Arcana.Graph.CommunityDetector

    require Logger

    @impl true
    def detect([], _relationships, _opts), do: {:ok, []}

    def detect(entities, relationships, opts) do
      resolution = Keyword.get(opts, :resolution, 1.0)
      objective = Keyword.get(opts, :objective, :cpm)
      iterations = Keyword.get(opts, :iterations, 2)
      seed = Keyword.get(opts, :seed, 0)
      min_size = Keyword.get(opts, :min_size, 1)
      max_level = Keyword.get(opts, :max_level, 1)

      # Build index mappings
      entity_ids = Enum.map(entities, & &1.id)
      id_to_index = entity_ids |> Enum.with_index() |> Map.new()
      index_to_id = entity_ids |> Enum.with_index() |> Map.new(fn {id, idx} -> {idx, id} end)

      # Convert to weighted edge tuples with integer indices
      edges = to_weighted_edges(id_to_index, relationships)

      Logger.info(
        "[Leiden] Starting: #{length(entity_ids)} entities, #{length(edges)} edges, " <>
          "resolution=#{resolution}, objective=#{objective}, min_size=#{min_size}, max_level=#{max_level}"
      )

      :telemetry.span(
        [:arcana, :graph, :community_detection],
        %{entity_count: length(entities), detector: :leiden},
        fn ->
          start_time = System.monotonic_time(:millisecond)

          leiden_opts = [
            n_nodes: length(entity_ids),
            objective: objective,
            resolution: resolution,
            iterations: iterations,
            seed: seed,
            max_levels: max_level,
            min_size: min_size
          ]

          result =
            case Leidenfold.detect_hierarchical_from_weighted_edges(edges, leiden_opts) do
              {:ok, levels} ->
                elapsed = System.monotonic_time(:millisecond) - start_time

                # Convert each level's membership to community maps
                communities = format_hierarchical_levels(levels, index_to_id, min_size)

                by_level = Enum.group_by(communities, & &1.level)

                level_summary =
                  by_level
                  |> Enum.sort_by(&elem(&1, 0))
                  |> Enum.map_join(", ", fn {lvl, comms} -> "L#{lvl}=#{length(comms)}" end)

                Logger.info("[Leiden] Completed in #{elapsed}ms: #{level_summary}")
                {:ok, communities}

              {:error, reason} ->
                elapsed = System.monotonic_time(:millisecond) - start_time
                Logger.error("[Leiden] Failed after #{elapsed}ms: #{inspect(reason)}")
                {:error, reason}
            end

          metadata =
            case result do
              {:ok, communities} -> %{community_count: length(communities)}
              {:error, _} -> %{community_count: 0}
            end

          {result, metadata}
        end
      )
    end

    # Convert relationships to weighted edge tuples {source_idx, target_idx, weight}
    defp to_weighted_edges(id_to_index, relationships) do
      relationships
      |> Enum.filter(fn rel ->
        Map.has_key?(id_to_index, rel.source_id) and
          Map.has_key?(id_to_index, rel.target_id)
      end)
      |> Enum.map(fn rel ->
        source_idx = Map.fetch!(id_to_index, rel.source_id)
        target_idx = Map.fetch!(id_to_index, rel.target_id)
        weight = (Map.get(rel, :strength, 1) || 1) / 1
        {source_idx, target_idx, weight}
      end)
    end

    # Convert hierarchical levels from leidenfold to community maps
    defp format_hierarchical_levels(levels, index_to_id, min_size) do
      Enum.flat_map(levels, fn %{level: level, membership: membership} ->
        format_communities(membership, index_to_id, level, min_size)
      end)
    end

    # Convert membership list to community maps
    # membership[node_idx] = community_id
    defp format_communities(membership, index_to_id, level, min_size) do
      membership
      |> Enum.with_index()
      |> Enum.group_by(fn {community_id, _node_idx} -> community_id end, fn {_community_id,
                                                                             node_idx} ->
        Map.fetch!(index_to_id, node_idx)
      end)
      |> Enum.map(fn {_community_id, entity_ids} ->
        %{level: level, entity_ids: entity_ids}
      end)
      |> Enum.filter(fn %{entity_ids: ids} -> length(ids) >= min_size end)
    end
  end
end