From 30e2dc7adb7ed3e6102ffea92eca6a0c9890b856 Mon Sep 17 00:00:00 2001
From: victor felder <victor@draft.li>
Date: Wed, 6 Mar 2024 21:40:03 +0100
Subject: [PATCH] Support custom reference lines

---
 lib/sparkline_svg.ex                 | 12 ++++++--
 lib/sparkline_svg/core.ex            |  7 ++++-
 lib/sparkline_svg/reference_line.ex  |  6 +++-
 test/sparkline_svg_ref_line_test.exs | 43 ++++++++++++++++++++++++++++
 4 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/lib/sparkline_svg.ex b/lib/sparkline_svg.ex
index bc13031..a6de719 100644
--- a/lib/sparkline_svg.ex
+++ b/lib/sparkline_svg.ex
@@ -122,18 +122,26 @@ defmodule SparklineSvg do
   reference lines as you want. Reference lines are displayed as horizontal lines that span the
   entire width of the chart.
 
-  There are four types of currently supported reference lines:
+  There are currently four types of supported reference lines:
   - `:max` - show the maximum value of the chart.
   - `:min` - show the minimum value of the chart.
   - `:avg` - show the average value of the chart.
   - `:median` - show the median value of the chart.
 
+  You can implement custom reference lines by passing a function that receives `Core.points()`
+  and returns the `y` value at which to display the line. See examples in tests.
+
   ``` elixir
   svg =
     datapoints
     |> SparklineSvg.new()
     |> SparklineSvg.show_line()
     |> SparklineSvg.show_ref_line(:max, color: "red")
+    |> SparklineSvg.show_ref_line(fn points ->
+      Enum.reduce(points, 0, fn {_x, y}, acc ->
+        y + acc
+      end) / 3
+    end, color: "blue")
     |> SparklineSvg.to_svg!()
   ```
 
@@ -304,7 +312,7 @@ defmodule SparklineSvg do
           | list({NaiveDateTime.t(), NaiveDateTime.t()})
 
   @typedoc "The type of reference line."
-  @type ref_line :: :max | :min | :avg | :median
+  @type ref_line :: :max | :min | :avg | :median | (Core.points() -> Core.y())
 
   @typedoc "Padding options for the chart."
   @type padding ::
diff --git a/lib/sparkline_svg/core.ex b/lib/sparkline_svg/core.ex
index 1f5e2b8..e25e071 100644
--- a/lib/sparkline_svg/core.ex
+++ b/lib/sparkline_svg/core.ex
@@ -89,7 +89,12 @@ defmodule SparklineSvg.Core do
   defp calc_resize_ref_lines(ref_lines, datapoints, min_max_y, options) do
     ref_lines
     |> Enum.map(fn {type, ref_line} ->
-      value = calc_ref_line(type, datapoints)
+      value =
+        cond do
+          is_atom(type) -> calc_ref_line(type, datapoints)
+          is_function(type, 1) -> type.(datapoints)
+        end
+
       position = resize_y(value, min_max_y, options)
 
       {type, %ReferenceLine{ref_line | position: position, value: value}}
diff --git a/lib/sparkline_svg/reference_line.ex b/lib/sparkline_svg/reference_line.ex
index 627abea..1ec391d 100644
--- a/lib/sparkline_svg/reference_line.ex
+++ b/lib/sparkline_svg/reference_line.ex
@@ -39,8 +39,12 @@ defmodule SparklineSvg.ReferenceLine do
 
     cond do
       keys == [] -> {:ok, ref_lines}
-      Enum.all?(keys, &Enum.member?(@valid_types, &1)) -> {:ok, ref_lines}
+      Enum.all?(keys, &valid_type?/1) -> {:ok, ref_lines}
       true -> {:error, :invalid_ref_line_type}
     end
   end
+
+  defp valid_type?(type) when type in @valid_types, do: true
+  defp valid_type?(fun) when is_function(fun, 1), do: true
+  defp valid_type?(_), do: false
 end
diff --git a/test/sparkline_svg_ref_line_test.exs b/test/sparkline_svg_ref_line_test.exs
index 40a3ee3..b5d844e 100644
--- a/test/sparkline_svg_ref_line_test.exs
+++ b/test/sparkline_svg_ref_line_test.exs
@@ -85,4 +85,47 @@ defmodule SparklineSvgMRefLineTest do
     assert sparkline.ref_lines.avg.value == 1.75
     assert sparkline.ref_lines.median.value == 1.5
   end
+
+  test "valid custom ref lines" do
+    percentile =
+      fn percentile ->
+        fn datapoints ->
+          values_count = length(datapoints)
+
+          case percentile / 100 * (values_count + 1) do
+            n when n < 1 ->
+              0
+
+            n ->
+              sorted_values = Enum.map(datapoints, &elem(&1, 1)) |> Enum.sort()
+
+              case Enum.drop(sorted_values, max(0, trunc(n) - 1)) do
+                [a, b | _] -> a + (n - trunc(n)) * (b - a)
+                [a] -> a
+              end
+          end
+        end
+      end
+
+    percentile_99 = percentile.(99)
+    fixed = fn _ -> 0.33 end
+
+    {:ok, sparkline} =
+      SparklineSvg.new(1..200)
+      |> SparklineSvg.show_ref_line(percentile_99)
+      |> SparklineSvg.show_ref_line(fixed)
+      |> SparklineSvg.dry_run()
+
+    assert sparkline.ref_lines[percentile_99].value == 198.99
+    assert sparkline.ref_lines[fixed].value == 0.33
+  end
+
+  test "custom ref line with empty chart" do
+    fun = fn _ -> 7 end
+
+    {:ok, sparkline} =
+      SparklineSvg.new([]) |> SparklineSvg.show_ref_line(fun) |> SparklineSvg.dry_run()
+
+    assert sparkline.ref_lines[fun].value == nil
+  end
 end