defmodule Statwise.Distributions.Normal do
  @moduledoc """
  Normal distribution helpers.
  """

  @sqrt2 :math.sqrt(2.0)

  @doc "Returns the standard normal cumulative distribution function."
  @spec cdf(number()) :: float()
  def cdf(x) when is_number(x) do
    0.5 * erfc(-x / @sqrt2)
  end

  @doc "Returns the standard normal survival function."
  @spec sf(number()) :: float()
  def sf(x) when is_number(x), do: 0.5 * erfc(x / @sqrt2)

  @doc "Returns the standard normal quantile for probability `p`."
  @spec ppf(float()) :: float()
  def ppf(p) when is_number(p) and p > 0.0 and p < 1.0 do
    # Peter John Acklam's rational approximation, refined enough for plotting
    # positions and statistical-test critical values.
    a = [
      -3.969_683_028_665_376e1,
      2.209_460_984_245_205e2,
      -2.759_285_104_469_687e2,
      1.383_577_518_672_69e2,
      -3.066_479_806_614_716e1,
      2.506_628_277_459_239
    ]

    b = [
      -5.447_609_879_822_406e1,
      1.615_858_368_580_409e2,
      -1.556_989_798_598_866e2,
      6.680_131_188_771_972e1,
      -1.328_068_155_288_572e1
    ]

    c = [
      -7.784_894_002_430_293e-3,
      -3.223_964_580_411_365e-1,
      -2.400_758_277_161_838,
      -2.549_732_539_343_734,
      4.374_664_141_464_968,
      2.938_163_982_698_783
    ]

    d = [
      7.784_695_709_041_462e-3,
      3.224_671_290_700_398e-1,
      2.445_134_137_142_996,
      3.754_408_661_907_416
    ]

    low = 0.02425
    high = 1.0 - low

    cond do
      p < low ->
        q = :math.sqrt(-2.0 * :math.log(p))
        rational(c, q) / rational(d ++ [1.0], q)

      p <= high ->
        q = p - 0.5
        r = q * q
        q * rational(a, r) / rational(b ++ [1.0], r)

      true ->
        q = :math.sqrt(-2.0 * :math.log(1.0 - p))
        -rational(c, q) / rational(d ++ [1.0], q)
    end
  end

  @doc "Returns a two-sided p-value from a z statistic."
  @spec two_sided_p_from_z(number()) :: float()
  def two_sided_p_from_z(z) when is_number(z), do: min(1.0, 2.0 * sf(abs(z)))

  @doc "Returns the Gauss error function."
  @spec erf(number()) :: float()
  def erf(x) when is_number(x), do: 1.0 - erfc(x)

  defp erfc(x) when x == 0, do: 1.0

  defp erfc(x) do
    z = abs(x)
    t = 1.0 / (1.0 + 0.5 * z)

    polynomial =
      1.00002368 +
        t *
          (0.37409196 +
             t *
               (0.09678418 +
                  t *
                    (-0.18628806 +
                       t *
                         (0.27886807 +
                            t *
                              (-1.13520398 +
                                 t * (1.48851587 + t * (-0.82215223 + t * 0.17087277)))))))

    tau = t * :math.exp(-z * z - 1.26551223 + t * polynomial)

    if x >= 0.0, do: tau, else: 2.0 - tau
  end

  defp rational(coefficients, x) do
    Enum.reduce(coefficients, 0.0, fn coefficient, acc -> acc * x + coefficient end)
  end
end