"""
Order Statistics

In statistics, an <url>:order statistic:
https://en.wikipedia.org/wiki/Order_statistic</url> gives \
the $k$-th smallest value.

Together with <url>:rank statistics:
https://en.wikipedia.org/wiki/Ranking</url> these are \
fundamental tools in non-parametric statistics and inference.

Important special cases of order statistics are finding \
minimum and maximum value of a sample and sample quantiles.
"""

from mpmath import ceil as mpceil, floor as mpfloor

from mathics.algorithm.introselect import introselect
from mathics.builtin.list.math import _RankedTakeLargest, _RankedTakeSmallest
from mathics.core.atoms import Integer, Integer1
from mathics.core.attributes import A_PROTECTED, A_READ_PROTECTED
from mathics.core.builtin import Builtin
from mathics.core.expression import Evaluation, Expression
from mathics.core.list import ListExpression
from mathics.core.symbols import Atom, SymbolFloor, SymbolPlus, SymbolTimes, SymbolTrue
from mathics.core.systemsymbols import (
    SymbolRankedMax,
    SymbolRankedMin,
    SymbolSort,
    SymbolSubtract,
)
from mathics.eval.numerify import numerify


class Quantile(Builtin):
    """

    <url>
    :Quantile:
    https://en.wikipedia.org/wiki/Quantile</url> (<url>
    :WMA link:
    https://reference.wolfram.com/language/ref/Quantile.html</url>)

    In statistics and probability, quantiles are cut points dividing the \
    range of a probability distribution into continuous intervals with \
    equal probabilities, or dividing the observations in a sample in the same way.

    Quantile is also known as value at risk (VaR) or fractile.
    <dl>
      <dt>'Quantile'[$list$, $q$]
      <dd>returns the $q$th quantile of $list$.

      <dt>'Quantile'[$list$, $q$, {{$a,b$}, {$c,d$}}]
      <dd>uses the quantile definition specified by parameters $a$, $b$, $c$, $d$.

      For a list of length $n$, 'Quantile'[$list$, $q$, {{$a ,b$}, {$c, d$}}] depends \
      on $x=a+(n+b)q$.

      If $x$ is an integer, the result is $s[[x]]$, where $s$='Sort[list,Less]'.

      Otherwise, the result is:
      's[[Floor[x]]] + (s[[Ceiling[x]]] - s[[Floor[x]]])(c + d FractionalPart[x])',
      with the indices taken to be 1 or n if they are out of range.

      The default choice of parameters is '{{0,0},{1,0}}'.
    </dl>

    Common choices of parameters include:
    <ul>
      <li>'{{0, 0}, {1, 0}}' inverse empirical CDF (default)
      <li>'{{0, 0}, {0, 1}}' linear interpolation (California method)
     </ul>

    'Quantile[list,q]' always gives a result equal to an element of list.

    >> Quantile[Range[11], 1/3]
     = 4

    >> Quantile[Range[16], 1/4]
     = 4

    >> Quantile[{1, 2, 3, 4, 5, 6, 7}, {1/4, 3/4}]
     = {2, 6}
    """

    messages = {
        "nquan": "The quantile `1` has to be between 0 and 1.",
    }

    attributes = A_PROTECTED | A_READ_PROTECTED
    rules = {
        "Quantile[list_List, q_, abcd_]": "Quantile[list, {q}, abcd]",
        "Quantile[list_List, q_]": "Quantile[list, q, {{0, 0}, {1, 0}}]",
    }
    summary_text = "cut points dividing the range of a probability distribution into continuous intervals"

    def eval(
        self,
        data: ListExpression,
        qs: ListExpression,
        a,
        b,
        c,
        d,
        evaluation: Evaluation,
    ):
        """Quantile[data_List, qs_List, {{a_, b_}, {c_, d_}}]"""

        n = len(data.elements)
        partially_sorted = data.get_mutable_elements()

        def ranked(i):
            return introselect(partially_sorted, min(max(0, i - 1), n - 1))

        numeric_qs = qs.evaluate(evaluation)
        if numeric_qs is not None:
            numeric_qs = numerify(numeric_qs, evaluation)
        results = []

        for q in numeric_qs.elements:
            py_q = q.to_mpmath()

            if py_q is None or not 0.0 <= py_q <= 1.0:
                evaluation.message("Quantile", "nquan", q)
                return

            x = (Integer(n) + b) * q + a

            numeric_x = numerify(x.evaluate(evaluation), evaluation)

            if isinstance(numeric_x, Integer):
                results.append(ranked(numeric_x.value))
            else:
                py_x = numeric_x.to_mpmath()

                if py_x is None:
                    return

                if c.get_int_value() == 1 and d.get_int_value() == 0:  # k == 1?
                    results.append(ranked(int(mpceil(py_x))))
                else:
                    py_floor_x = mpfloor(py_x)
                    s0 = ranked(int(py_floor_x))
                    s1 = ranked(int(mpceil(py_x)))

                    k = Expression(
                        SymbolPlus,
                        c,
                        Expression(
                            SymbolTimes,
                            d,
                            Expression(SymbolSubtract, x, Expression(SymbolFloor, x)),
                        ),
                    )

                    results.append(
                        Expression(
                            SymbolPlus,
                            s0,
                            Expression(
                                SymbolTimes, k, Expression(SymbolSubtract, s1, s0)
                            ),
                        )
                    )

        if len(results) == 1:
            return results[0]
        else:
            return ListExpression(*results)


class Quartiles(Builtin):
    """
    <url>:Quartile:
    https://en.wikipedia.org/wiki/Quartile</url> (<url>
    :WMA link:
    https://reference.wolfram.com/language/ref/Quartiles.html</url>)
    <dl>
      <dt>'Quartiles'[$list$]
      <dd>returns the 1/4, 1/2, and 3/4 quantiles of $list$.
    </dl>

    >> Quartiles[Range[25]]
     = {27 / 4, 13, 77 / 4}
    """

    attributes = A_PROTECTED | A_READ_PROTECTED
    rules = {
        "Quartiles[list_List]": "Quantile[list, {1/4, 1/2, 3/4}, {{1/2, 0}, {0, 1}}]",
    }
    summary_text = "list of quartiles"


class RankedMax(Builtin):
    """
    <url>:WMA link:https://reference.wolfram.com/language/ref/RankedMax.html</url>

    <dl>
      <dt>'RankedMax'[$list$, $n$]
      <dd>returns the $n$th largest element of $list$ (with $n$ = 1 yielding the largest element,
      $n$ = 2 yielding the second largest element, and so on).
    </dl>

    >> RankedMax[{482, 17, 181, -12}, 2]
     = 181
    """

    messages = {
        "intpm": "Expected positive integer at position 2 in ``.",
        "rank": "The specified rank `1` is not between 1 and `2`.",
    }
    attributes = A_PROTECTED | A_READ_PROTECTED
    summary_text = "the n-th largest item"

    def eval(self, element, n: Integer, evaluation: Evaluation):
        "RankedMax[element_List, n_Integer]"
        py_n = n.value
        if py_n < 1:
            evaluation.message(
                "RankedMax", "intpm", Expression(SymbolRankedMax, element, n)
            )
        elif py_n > len(element.elements):
            evaluation.message("RankedMax", "rank", py_n, len(element.elements))
        else:
            return introselect(
                element.get_mutable_elements(), len(element.elements) - py_n
            )


class RankedMin(Builtin):
    """
    <url>:WMA link:
    https://reference.wolfram.com/language/ref/RankedMin.html</url>

    <dl>
      <dt>'RankedMin'[$list$, $n$]
      <dd>returns the $n$th smallest element of $list$ (with \
          $n$ = 1 yielding the smallest element, $n$ = 2 yielding \
          the second smallest element, and so on).
    </dl>

    >> RankedMin[{482, 17, 181, -12}, 2]
     = 17
    """

    messages = {
        "intpm": "Expected positive integer at position 2 in ``.",
        "rank": "The specified rank `1` is not between 1 and `2`.",
    }
    attributes = A_PROTECTED | A_READ_PROTECTED
    summary_text = "the n-th smallest item"

    def eval(self, element, n: Integer, evaluation: Evaluation):
        "RankedMin[element_List, n_Integer]"
        py_n = n.value
        if py_n < 1:
            evaluation.message(
                "RankedMin", "intpm", Expression(SymbolRankedMin, element, n)
            )
        elif py_n > len(element.elements):
            evaluation.message("RankedMin", "rank", py_n, len(element.elements))
        else:
            return introselect(element.get_mutable_elements(), py_n - 1)


class ReverseSort(Builtin):
    """
    <url>:WMA link:https://reference.wolfram.com/language/ref/ReverseSort.html</url>

    <dl>
      <dt>'ReverseSort'[$list$]
      <dd>sorts $list$ (or the elements of any other expression) according \
          to reverse canonical ordering.

      <dt>'ReverseSort'[$list$, $p$]
      <dd>sorts using $p$ to determine the order of two elements.
    </dl>

    >> ReverseSort[{c, b, d, a}]
     = {d, c, b, a}

    You can specify a binary comparison function:
    >> ReverseSort[{1, 2, 0, 3}, Less]
     = {3, 2, 1, 0}

    Using 'Greater' for the above, reverses the reverse sort:
    >> ReverseSort[{1, 2, 0, 3}, Greater]
     = {0, 1, 2, 3}

    See also <url>:Sort:
    /doc/reference-of-built-in-symbols/descriptive-statistics/order-statistics/sort/</url>.
    """

    attributes = A_PROTECTED
    summary_text = "sort in reverse order"

    rules = {
        "ReverseSort[list_]": "Reverse[Sort[list]]",
        "ReverseSort[list_, p_]": "Reverse[Sort[list, p]]",
    }


# FIXME: there might be a bug in sorting...
#
#    Sort[{
#    "a","b", 1,
#    ByteArray[{1,2,4,1}],
#    2, 1.2, I, 2I-3, A,
#    a+b, a*b, a+1, a*2, b^3, 2/3,
#    A[x], F[2], F[x], F[x_], F[x___], F[x,t], F[x__],
#    Condition[A,b>2], Pattern[expr, A]
#    }]
#
# should be:
#
#     {-3 + 2*I, I, 2/3, 1, 1.2, 2,
#     "a", "b", 2*a,
#      1 + a, A, a*b, b^3, a + b,
#      A[x], A /; b > 2,
#      F[2], F[x], F[x_], F[x___], F[x__], F[x, t],
#      ByteArray["AQIEAQ=="], expr:A}
#
# But this is too complicated a case to run as a test. It needs
# to be isolated. Break this down to smaller pieces,
# and also use Order[] to check smaller components.
# The problem might also be in boxing-order output.


class Sort(Builtin):
    """
    <url>:WMA link:https://reference.wolfram.com/language/ref/Sort.html</url>

    <dl>
      <dt>'Sort'[$list$]
      <dd>sorts $list$ (or the elements of any other expression) according \
          to canonical ordering.

      <dt>'Sort'[$list$, $p$]
      <dd>sorts using $p$ to determine the order of two elements.
    </dl>

    >> Sort[{4, 1.0, a, 3+I}]
     = {1., 3 + I, 4, a}

    Sort uses 'Order' to determine ordering by default.
    You can sort patterns according to their precedence using 'PatternsOrderedQ':
    >> Sort[{items___, item_, OptionsPattern[], item_symbol, item_?test}, PatternsOrderedQ]
     = {item_symbol, item_ ? test, item_, items___, OptionsPattern[]}

    When sorting patterns, values of atoms do not matter:
    >> Sort[{a, b/;t}, PatternsOrderedQ]
     = {b /; t, a}
    >> Sort[{2+c_, 1+b__}, PatternsOrderedQ]
     = {2 + c_, 1 + b__}
    >> Sort[{x_ + n_*y_, x_ + y_}, PatternsOrderedQ]
     = {x_ + n_ y_, x_ + y_}

    See also <url>:ReverseSort:
    /doc/reference-of-built-in-symbols/descriptive-statistics/order-statistics/reversesort/</url>.
    """

    attributes = A_PROTECTED
    summary_text = "sort lexicographically or with a comparison function"

    def eval(self, list, evaluation: Evaluation):
        "Sort[list_]"

        if isinstance(list, Atom):
            evaluation.message("Sort", "normal", Integer1, Expression(SymbolSort, list))
        else:
            new_elements = sorted(list.elements)
            return list.restructure(list.head, new_elements, evaluation)

    def eval_predicate(self, list, p, evaluation: Evaluation):
        "Sort[list_, p_]"

        if isinstance(list, Atom):
            evaluation.message("Sort", "normal", Integer1, Expression(SymbolSort, list))
        else:

            class Key:
                def __init__(self, element):
                    self.element = element

                def __gt__(self, other):
                    return (
                        Expression(p, self.element, other.element).evaluate(evaluation)
                        is not SymbolTrue
                    )

            new_elements = sorted(list.elements, key=Key)
            return list.restructure(list.head, new_elements, evaluation)


class TakeLargest(_RankedTakeLargest):
    """
    <url>
    :WMA link:
    https://reference.wolfram.com/language/ref/TakeLargest.html</url>

    <dl>
      <dt>'TakeLargest'[$list$, $f$, $n$]
      <dd>returns the a sorted list of the $n$ largest items in $list$.
    </dl>

    List the largest two numbers of a list:
    >> TakeLargest[{100, -1, 50, 10}, 2]
     = {100, 50}

    None, Null, Indeterminate and expressions with head Missing are ignored \
    by default:
    >> TakeLargest[{-8, 150, Missing[abc]}, 2]
     = {150, -8}

    You may specify which items are ignored using the option 'ExcludedForms':
    >> TakeLargest[{-8, 150, Missing[abc]}, 2, ExcludedForms -> {}]
     = {Missing[abc], 150}
    """

    attributes = A_PROTECTED
    summary_text = "sublist of n largest elements"

    def eval(self, element, n, evaluation, options):
        "TakeLargest[element_List, n_, OptionsPattern[TakeLargest]]"
        return self._compute(element, n, evaluation, options)


class TakeSmallest(_RankedTakeSmallest):
    """
    <url>:WMA link:https://reference.wolfram.com/language/ref/TakeSmallest.html</url>

    <dl>
      <dt>'TakeSmallest'[$list$, $n$]
      <dd>returns the a sorted list of the $n$ smallest items in $list$.
    </dl>

    List the smallest two numbers of a list:
    >> TakeSmallest[{100, -1, 50, 10}, 2]
     = {-1, 10}

    For details on how to use the 'ExcludedForms' option, see <url>
    :TakeLargest:
    /doc/reference-of-built-in-symbols/descriptive-statistics/order-statistics/takelargest/</url>.
    """

    attributes = A_PROTECTED
    summary_text = "sublist of n smallest elements"

    def eval(self, element, n, evaluation, options):
        "TakeSmallest[element_List, n_, OptionsPattern[TakeSmallest]]"
        return self._compute(element, n, evaluation, options)


# TODO: MinMax
