thorhead_math.ex

defmodule Thorhead.Analysis do
  @moduledoc """
  Computing overall rarity distribution using generative functions.
  """
  alias Thorhead.Nfts.Metadata

  # Multiply two polynomials (PMFs)
  # Each polynomial is represented as a map: %{rarity => probability}
  def poly_mult(poly1, poly2) do
    Enum.reduce(poly1, %{}, fn {exp1, coeff1}, acc ->
      Enum.reduce(poly2, acc, fn {exp2, coeff2}, inner_acc ->
        exp = exp1 + exp2
        Map.update(inner_acc, exp, coeff1 * coeff2, &(&1 + coeff1 * coeff2))
      end)
    end)
  end

  # Sum probabilities for totals above given threshold
  def total_probability(poly, threshold) do
    poly
    |> Enum.filter(fn {exp, _} -> exp >= threshold end)
    |> Enum.reduce(0.0, fn {_, prob}, acc -> acc + prob end)
  end

  # Create a PMF for attributes with a flat weight map and corresponding rarity list.
  # E.g. for Background.
  def background_pmf do
    total = Metadata.total_background_weight()
    # Build a map of option => rarity from the rarity keyword list.
    rarity_map = Map.new(Metadata.background_rarity_list())
    Metadata.background_weight_map()
    |> Enum.reduce(%{}, fn {option, weight}, acc ->
      rarity = Map.get(rarity_map, option)
      prob = weight / total
      Map.update(acc, rarity, prob, &(&1 + prob))
    end)
  end

  # For attributes defined with bucketed weight maps, e.g. Head Style.
  def head_style_pmf do
    total = Metadata.total_head_style_weight()
    Metadata.head_style_weights()
    |> Enum.map(fn {_bucket, {weight, _options, rarity}} ->
      {rarity, weight / total}
    end)
    |> Enum.reduce(%{}, fn {rarity, prob}, acc ->
      Map.update(acc, rarity, prob, &(&1 + prob))
    end)
  end

  # For Head Height (flat weight map).
  def head_height_pmf do
    total = Metadata.total_head_height_weight()
    rarity_map = Map.new(Metadata.head_height_rarity_list())
    Metadata.head_height_weight_map()
    |> Enum.reduce(%{}, fn {option, weight}, acc ->
      rarity = Map.get(rarity_map, option)
      prob = weight / total
      Map.update(acc, rarity, prob, &(&1 + prob))
    end)
  end

  # For Skin Color (bucketed weight map)
  def skin_color_pmf do
    total = Metadata.total_skin_color_weight()
    Metadata.skin_color_weights()
    |> Enum.map(fn {_bucket, {weight, _colors, rarity}} ->
      {rarity, weight / total}
    end)
    |> Enum.reduce(%{}, fn {rarity, prob}, acc ->
      Map.update(acc, rarity, prob, &(&1 + prob))
    end)
  end

  # For Primary Color (bucketed weight map)
  def primary_color_pmf do
    total = Metadata.total_primary_color_weight()
    Metadata.primary_color_weights()
    |> Enum.map(fn {_bucket, {weight, _colors, rarity}} ->
      {rarity, weight / total}
    end)
    |> Enum.reduce(%{}, fn {rarity, prob}, acc ->
      Map.update(acc, rarity, prob, &(&1 + prob))
    end)
  end

  # Similar to Primary Color (bucketed weight map) but with the added condition that not every head style has an accent color. Ref: @head_style_config module attribute.
  # We assume that:
  # 1. If the chosen head style does NOT support accent color then accent_color is forced to :none (rarity 0),
  # 2. Otherwise, accent_color is chosen according to the Primary Color PMF.
  def accent_color_pmf do
    allowed_prob = head_style_allowed_probability()
    primary = primary_color_pmf()

    # For head styles with accent_color allowed, scale primary PMF by allowed_prob.
    primary_scaled =
      Enum.into(primary, %{}, fn {rarity, prob} ->
        {rarity, prob * allowed_prob}
      end)

    # For head styles that don't allow an accent color, accent_color becomes :none (mapped to rarity 0).
    # We add that probability to rarity 0.
    Map.update(primary_scaled, 0, 1 - allowed_prob, &(&1 + (1 - allowed_prob)))
  end

  # Helper: Compute the probability that a randomly generated head style supports accent color.
  # We use two pieces of data:
  #   - Metadata.head_style_weights(): bucketed structure {weight, options, _rarity}
  #   - Metadata.head_style_config(): map of head style options to a config where :accent_color is true/false.
  defp head_style_allowed_probability do
    total_weight = Metadata.total_head_style_weight()
    Metadata.head_style_weights()
    |> Enum.reduce(0.0, fn {_bucket, {weight, options, _}}, acc ->
      # For each bucket, determine the fraction of options that support accent_color.
      allowed_count = Enum.count(options, fn opt ->
        case Map.get(Metadata.head_style_config(), opt) do
          %{accent_color: true} -> true
          _ -> false
        end
      end)

      fraction_allowed =
        if length(options) > 0 do
          allowed_count / length(options)
        else
          0
        end

      acc + (weight / total_weight) * fraction_allowed
    end)
  end

  # For Holographic (flat weight map) 4096 total, 4095 No, 1 Yes. with an addition of 100 rarity for the Yes option.
  def holographic_pmf do
    total = 4096
    rarity_map = %{false => 0, true => 100}
    %{false => 4095, true => 1}
    |> Enum.reduce(%{}, fn {option, weight}, acc ->
      rarity = Map.get(rarity_map, option)
      prob = weight / total
      Map.update(acc, rarity, prob, &(&1 + prob))
    end)
  end

  # (Additional attribute PMF definitions similar to above would be placed here)
  # For this example, we'll assume PMFs for Background and Head Style are enough.

  # Calculate the number of rolls (n) needed to have at least target probability (t)
  # of obtaining an NFT with total rarity >= threshold.
  def rolls_needed(p, target_probability) when p > 0 and p < 1 do
    # (1 - p)^n <= (1 - target_probability)
    # n >= log(1 - target_probability) / log(1 - p)
    n = :math.log(1 - target_probability) / :math.log(1 - p)
    Float.ceil(n) |> trunc()
  end

  # Example of running the full calculation:
  def run_example(threshold) do
    # Compute PMFs for each attribute.
    attr_background = background_pmf()
    attr_head_style = head_style_pmf()
    attr_head_height = head_height_pmf()
    attr_skin_color = skin_color_pmf()
    attr_primary_color = primary_color_pmf()
    attr_accent_color = accent_color_pmf()
    attr_holographic = holographic_pmf()

    IO.inspect(attr_background, label: "Background PMF")
    IO.inspect(attr_head_style, label: "Head Style PMF")
    IO.inspect(attr_head_height, label: "Head Height PMF")
    IO.inspect(attr_skin_color, label: "Skin Color PMF")
    IO.inspect(attr_primary_color, label: "Primary Color PMF")
    IO.inspect(attr_accent_color, label: "Accent Color PMF")
    IO.inspect(attr_holographic, label: "Holographic PMF")

    # Convolve the PMFs
    total_poly =
      attr_background
      |> poly_mult(attr_head_style)
      |> poly_mult(attr_head_height)
      |> poly_mult(attr_skin_color)
      |> poly_mult(attr_primary_color)
      |> poly_mult(attr_accent_color)
      |> poly_mult(attr_holographic)

    IO.inspect(Enum.sort(total_poly), label: "Combined PMF", pretty: true, limit: :infinity)

    # Calculate probability of total rarity >= 20
    p_ge_20 = total_probability(total_poly, threshold)
    IO.puts("Probability that total rarity >= #{threshold} is: #{p_ge_20}")

    # Define a target guarantee probability (e.g. 99% chance at least one succeeds)
    target_probability = 0.99

    # Calculate number of rolls needed.
    rolls = rolls_needed(p_ge_20, target_probability)
    IO.puts("Number of rolls needed to have at least a #{target_probability * 100}% chance: #{rolls}")
  end
end