1# -*- coding: utf-8 -*-
2
3"""
4List Functions
5"""
6
7
8from itertools import chain, permutations
9from typing import Callable
10
11from mathics.version import __version__  # noqa used in loading to check consistency.
12from mathics.builtin.base import (
13    Builtin,
14    Test,
15    InvalidLevelspecError,
16    BinaryOperator,
17    PartError,
18    PartDepthError,
19    PartRangeError,
20    Predefined,
21    SympyFunction,
22)
23from mathics.builtin.scoping import dynamic_scoping
24from mathics.builtin.base import (
25    MessageException,
26    NegativeIntegerException,
27    CountableInteger,
28
29)
30from mathics.core.expression import (
31    Expression,
32    String,
33    ByteArrayAtom,
34    Symbol,
35    SymbolFailed,
36    SymbolNull,
37    SymbolN,
38    SymbolRule,
39    SymbolMakeBoxes,
40    SymbolAssociation,
41    SymbolSequence,
42    Integer,
43    Integer0,
44    Number,
45    Real,
46    strip_context,
47    from_python,
48    SymbolList,
49    SymbolByteArray,
50)
51from mathics.core.expression import min_prec, machine_precision
52from mathics.core.expression import structure
53from mathics.core.evaluation import BreakInterrupt, ContinueInterrupt, ReturnInterrupt
54from mathics.core.rules import Pattern, Rule
55from mathics.core.convert import from_sympy
56from mathics.builtin.numbers.algebra import cancel
57from mathics.algorithm.introselect import introselect
58from mathics.algorithm.clusters import (
59    optimize,
60    agglomerate,
61    kmeans,
62    PrecomputedDistances,
63    LazyDistances,
64)
65from mathics.algorithm.clusters import AutomaticSplitCriterion, AutomaticMergeCriterion
66from mathics.builtin.options import options_to_rules
67
68import sympy
69import heapq
70
71from collections import defaultdict
72import functools
73
74
75def deletecases_with_levelspec(expr, pattern, evaluation, levelspec=1, n=-1):
76    """
77    This function walks the expression `expr` and deleting occurrencies of `pattern`
78
79    If levelspec specifies a number, only those positions with  `levelspec` "coordinates" are return. By default, it just return occurences in the first level.
80
81    If a tuple (nmin, nmax) is provided, it just return those occurences with a number of "coordinates" between nmin and nmax.
82    n indicates the number of occurrences to return. By default, it returns all the occurences.
83    """
84    nothing = Symbol("System`Nothing")
85    from mathics.builtin.patterns import Matcher
86
87    match = Matcher(pattern)
88    match = match.match
89    if type(levelspec) is int:
90        lsmin = 1
91        lsmax = levelspec + 1
92    else:
93        lsmin = levelspec[0]
94        if levelspec[1]:
95            lsmax = levelspec[1] + 1
96        else:
97            lsmax = -1
98    tree = [[expr]]
99    changed_marks = [
100        [False],
101    ]
102    curr_index = [0]
103
104    while curr_index[0] != 1:
105        # If the end of the branch is reached, or no more elements to delete out
106        if curr_index[-1] == len(tree[-1]) or n == 0:
107            leaves = tree[-1]
108            tree.pop()
109            # check if some of the leaves was changed
110            changed = any(changed_marks[-1])
111            changed_marks.pop()
112            if changed:
113                leaves = [leaf for leaf in leaves if leaf is not nothing]
114            curr_index.pop()
115            if len(curr_index) == 0:
116                break
117            idx = curr_index[-1]
118            changed = changed or changed_marks[-1][idx]
119            changed_marks[-1][idx] = changed
120            if changed:
121                head = tree[-1][curr_index[-1]].get_head()
122                tree[-1][idx] = Expression(head, *leaves)
123            if len(curr_index) == 0:
124                break
125            curr_index[-1] = curr_index[-1] + 1
126            continue
127        curr_leave = tree[-1][curr_index[-1]]
128        if match(curr_leave, evaluation) and (len(curr_index) > lsmin):
129            tree[-1][curr_index[-1]] = nothing
130            changed_marks[-1][curr_index[-1]] = True
131            curr_index[-1] = curr_index[-1] + 1
132            n = n - 1
133            continue
134        if curr_leave.is_atom() or lsmax == len(curr_index):
135            curr_index[-1] = curr_index[-1] + 1
136            continue
137        else:
138            tree.append(list(curr_leave.get_leaves()))
139            changed_marks.append([False for l in tree[-1]])
140            curr_index.append(0)
141    return tree[0][0]
142
143
144def find_matching_indices_with_levelspec(expr, pattern, evaluation, levelspec=1, n=-1):
145    """
146    This function walks the expression `expr` looking for a pattern `pattern`
147    and returns the positions of each occurence.
148
149    If levelspec specifies a number, only those positions with  `levelspec` "coordinates" are return. By default, it just return occurences in the first level.
150
151    If a tuple (nmin, nmax) is provided, it just return those occurences with a number of "coordinates" between nmin and nmax.
152    n indicates the number of occurrences to return. By default, it returns all the occurences.
153    """
154    from mathics.builtin.patterns import Matcher
155
156    match = Matcher(pattern)
157    match = match.match
158    if type(levelspec) is int:
159        lsmin = 0
160        lsmax = levelspec
161    else:
162        lsmin = levelspec[0]
163        lsmax = levelspec[1]
164    tree = [expr.get_leaves()]
165    curr_index = [0]
166    found = []
167    while len(tree) > 0:
168        if n == 0:
169            break
170        if curr_index[-1] == len(tree[-1]):
171            curr_index.pop()
172            tree.pop()
173            if len(curr_index) != 0:
174                curr_index[-1] = curr_index[-1] + 1
175            continue
176        curr_leave = tree[-1][curr_index[-1]]
177        if match(curr_leave, evaluation) and (len(curr_index) >= lsmin):
178            found.append([from_python(i) for i in curr_index])
179            curr_index[-1] = curr_index[-1] + 1
180            n = n - 1
181            continue
182        if curr_leave.is_atom() or lsmax == len(curr_index):
183            curr_index[-1] = curr_index[-1] + 1
184            continue
185        else:
186            tree.append(curr_leave.get_leaves())
187            curr_index.append(0)
188    return found
189
190
191class Normal(Builtin):
192    """
193        <dl>
194    <dt>'Normal[expr_]'
195       <dd> Brings especial expressions to a normal expression from
196       different especial forms.
197    </dl>
198    """
199
200
201class ByteArray(Builtin):
202    r"""
203    <dl>
204    <dt>'ByteArray[{$b_1$, $b_2$, ...}]'
205       <dd> Represents a sequence of Bytes $b_1$, $b_2$, ...
206    <dt>'ByteArray["string"]'
207       <dd> Constructs a byte array where bytes comes from decode a b64 encoded String
208    </dl>
209
210    >> A=ByteArray[{1, 25, 3}]
211     = ByteArray["ARkD"]
212    >> A[[2]]
213     = 25
214    >> Normal[A]
215     = {1, 25, 3}
216    >> ToString[A]
217     = ByteArray["ARkD"]
218    >> ByteArray["ARkD"]
219     = ByteArray["ARkD"]
220    >> B=ByteArray["asy"]
221     : The first argument in Bytearray[asy] should be a B64 enconded string or a vector of integers.
222     = $Failed
223    """
224
225    messages = {
226        "aotd": "Elements in `1` are inconsistent with type Byte",
227        "lend": "The first argument in Bytearray[`1`] should "
228        + "be a B64 enconded string or a vector of integers.",
229    }
230
231    def apply_str(self, string, evaluation):
232        "ByteArray[string_String]"
233        try:
234            atom = ByteArrayAtom(string.value)
235        except Exception:
236            evaluation.message("ByteArray", "lend", string)
237            return SymbolFailed
238        return Expression("ByteArray", atom)
239
240    def apply_to_str(self, baa, evaluation):
241        "ToString[ByteArray[baa_ByteArrayAtom]]"
242        return String('ByteArray["' + baa.__str__() + '"]')
243
244    def apply_normal(self, baa, evaluation):
245        "System`Normal[ByteArray[baa_ByteArrayAtom]]"
246        return Expression(SymbolList, *[Integer(x) for x in baa.value])
247
248    def apply_list(self, values, evaluation):
249        "ByteArray[values_List]"
250        if not values.has_form("List", None):
251            return
252        try:
253            ba = bytearray([b.get_int_value() for b in values._leaves])
254        except:
255            evaluation.message("ByteArray", "aotd", values)
256            return
257        return Expression(SymbolByteArray, ByteArrayAtom(ba))
258
259
260class List(Builtin):
261    """
262    <dl>
263    <dt>'List[$e1$, $e2$, ..., $ei$]'
264    <dt>'{$e1$, $e2$, ..., $ei$}'
265        <dd>represents a list containing the elements $e1$...$ei$.
266    </dl>
267
268    'List' is the head of lists:
269    >> Head[{1, 2, 3}]
270     = List
271
272    Lists can be nested:
273    >> {{a, b, {c, d}}}
274     = {{a, b, {c, d}}}
275    """
276
277    attributes = ("Locked",)
278
279    def apply_makeboxes(self, items, f, evaluation):
280        """MakeBoxes[{items___},
281        f:StandardForm|TraditionalForm|OutputForm|InputForm]"""
282
283        items = items.get_sequence()
284        return Expression(
285            "RowBox", Expression(SymbolList, *list_boxes(items, f, "{", "}"))
286        )
287
288
289class ListQ(Test):
290    """
291    <dl>
292    <dt>'ListQ[$expr$]'
293        <dd>tests whether $expr$ is a 'List'.
294    </dl>
295
296    >> ListQ[{1, 2, 3}]
297     = True
298    >> ListQ[{{1, 2}, {3, 4}}]
299     = True
300    >> ListQ[x]
301     = False
302    """
303
304    def test(self, expr):
305        return expr.get_head_name() == "System`List"
306
307
308class NotListQ(Test):
309    """
310    <dl>
311    <dt>'NotListQ[$expr$]'
312        <dd>returns true if $expr$ is not a list.
313    </dl>
314    """
315
316    def test(self, expr):
317        return expr.get_head_name() != "System`List"
318
319
320def list_boxes(items, f, open=None, close=None):
321    result = [Expression(SymbolMakeBoxes, item, f) for item in items]
322    if f.get_name() in ("System`OutputForm", "System`InputForm"):
323        sep = ", "
324    else:
325        sep = ","
326    result = riffle(result, String(sep))
327    if len(items) > 1:
328        result = Expression("RowBox", Expression(SymbolList, *result))
329    elif items:
330        result = result[0]
331    if result:
332        result = [result]
333    else:
334        result = []
335    if open is not None and close is not None:
336        return [String(open)] + result + [String(close)]
337    else:
338        return result
339
340
341class Length(Builtin):
342    """
343    <dl>
344    <dt>'Length[$expr$]'
345        <dd>returns the number of leaves in $expr$.
346    </dl>
347
348    Length of a list:
349    >> Length[{1, 2, 3}]
350     = 3
351
352    'Length' operates on the 'FullForm' of expressions:
353    >> Length[Exp[x]]
354     = 2
355    >> FullForm[Exp[x]]
356     = Power[E, x]
357
358    The length of atoms is 0:
359    >> Length[a]
360     = 0
361
362    Note that rational and complex numbers are atoms, although their
363    'FullForm' might suggest the opposite:
364    >> Length[1/3]
365     = 0
366    >> FullForm[1/3]
367     = Rational[1, 3]
368    """
369
370    def apply(self, expr, evaluation):
371        "Length[expr_]"
372
373        if expr.is_atom():
374            return Integer0
375        else:
376            return Integer(len(expr.leaves))
377
378
379class All(Predefined):
380    """
381    <dl>
382      <dt>'All'
383      <dd>is a possible option value for 'Span', 'Quiet', 'Part' and related functions. 'All' specifies all parts at a particular level.
384    </dl>
385    """
386
387    pass
388
389
390class None_(Predefined):
391    """
392    <dl>
393    <dt>'None'
394        <dd>is a possible value for 'Span' and 'Quiet'.
395    </dl>
396    """
397
398    name = "None"
399
400
401class Span(BinaryOperator):
402    """
403    <dl>
404    <dt>'Span'
405        <dd>is the head of span ranges like '1;;3'.
406    </dl>
407
408    >> ;; // FullForm
409     = Span[1, All]
410    >> 1;;4;;2 // FullForm
411     = Span[1, 4, 2]
412    >> 2;;-2 // FullForm
413     = Span[2, -2]
414    >> ;;3 // FullForm
415     = Span[1, 3]
416
417    ## Parsing: 8 cases to consider
418    #> a ;; b ;; c // FullForm
419     = Span[a, b, c]
420    #>   ;; b ;; c // FullForm
421     = Span[1, b, c]
422    #> a ;;   ;; c // FullForm
423     = Span[a, All, c]
424    #>   ;;   ;; c // FullForm
425     = Span[1, All, c]
426    #> a ;; b      // FullForm
427     = Span[a, b]
428    #>   ;; b      // FullForm
429     = Span[1, b]
430    #> a ;;        // FullForm
431     = Span[a, All]
432    #>   ;;        // FullForm
433     = Span[1, All]
434
435    ## Formatting
436    #> a ;; b ;; c
437     = a ;; b ;; c
438    #> a ;; b
439     = a ;; b
440    #> a ;; b ;; c ;; d
441     = (1 ;; d) (a ;; b ;; c)
442    """
443
444    operator = ";;"
445    precedence = 305
446
447
448def join_lists(lists):
449    new_list = []
450    for list in lists:
451        new_list.extend(list)
452    return new_list
453
454
455def get_part(varlist, indices):
456    " Simple part extraction. indices must be a list of python integers. "
457
458    def rec(cur, rest):
459        if rest:
460            if cur.is_atom():
461                raise PartDepthError(rest[0])
462            pos = rest[0]
463            leaves = cur.get_leaves()
464            try:
465                if pos > 0:
466                    part = leaves[pos - 1]
467                elif pos == 0:
468                    part = cur.get_head()
469                else:
470                    part = leaves[pos]
471            except IndexError:
472                raise PartRangeError
473            return rec(part, rest[1:])
474        else:
475            return cur
476
477    return rec(varlist, indices).copy()
478
479
480def set_part(varlist, indices, newval):
481    " Simple part replacement. indices must be a list of python integers. "
482
483    def rec(cur, rest):
484        if len(rest) > 1:
485            pos = rest[0]
486            if cur.is_atom():
487                raise PartDepthError
488            try:
489                if pos > 0:
490                    part = cur._leaves[pos - 1]
491                elif pos == 0:
492                    part = cur.get_head()
493                else:
494                    part = cur._leaves[pos]
495            except IndexError:
496                raise PartRangeError
497            return rec(part, rest[1:])
498        elif len(rest) == 1:
499            pos = rest[0]
500            if cur.is_atom():
501                raise PartDepthError
502            try:
503                if pos > 0:
504                    cur.set_leaf(pos - 1, newval)
505                elif pos == 0:
506                    cur.set_head(newval)
507                else:
508                    cur.set_leaf(pos, newval)
509            except IndexError:
510                raise PartRangeError
511
512    rec(varlist, indices)
513
514
515def _parts_all_selector():
516    start = 1
517    stop = None
518    step = 1
519
520    def select(inner):
521        if inner.is_atom():
522            raise MessageException("Part", "partd")
523        py_slice = python_seq(start, stop, step, len(inner.leaves))
524        if py_slice is None:
525            raise MessageException("Part", "take", start, stop, inner)
526        return inner.leaves[py_slice]
527
528    return select
529
530
531def _parts_span_selector(pspec):
532    if len(pspec.leaves) > 3:
533        raise MessageException("Part", "span", pspec)
534    start = 1
535    stop = None
536    step = 1
537    if len(pspec.leaves) > 0:
538        start = pspec.leaves[0].get_int_value()
539    if len(pspec.leaves) > 1:
540        stop = pspec.leaves[1].get_int_value()
541        if stop is None:
542            if pspec.leaves[1].get_name() == "System`All":
543                stop = None
544            else:
545                raise MessageException("Part", "span", pspec)
546    if len(pspec.leaves) > 2:
547        step = pspec.leaves[2].get_int_value()
548
549    if start == 0 or stop == 0:
550        # index 0 is undefined
551        raise MessageException("Part", "span", 0)
552
553    if start is None or step is None:
554        raise MessageException("Part", "span", pspec)
555
556    def select(inner):
557        if inner.is_atom():
558            raise MessageException("Part", "partd")
559        py_slice = python_seq(start, stop, step, len(inner.leaves))
560        if py_slice is None:
561            raise MessageException("Part", "take", start, stop, inner)
562        return inner.leaves[py_slice]
563
564    return select
565
566
567def _parts_sequence_selector(pspec):
568    if not isinstance(pspec, (tuple, list)):
569        indices = [pspec]
570    else:
571        indices = pspec
572
573    for index in indices:
574        if not isinstance(index, Integer):
575            raise MessageException("Part", "pspec", pspec)
576
577    def select(inner):
578        if inner.is_atom():
579            raise MessageException("Part", "partd")
580
581        leaves = inner.leaves
582        n = len(leaves)
583
584        for index in indices:
585            int_index = index.value
586
587            if int_index == 0:
588                yield inner.head
589            elif 1 <= int_index <= n:
590                yield leaves[int_index - 1]
591            elif -n <= int_index <= -1:
592                yield leaves[int_index]
593            else:
594                raise MessageException("Part", "partw", index, inner)
595
596    return select
597
598
599def _part_selectors(indices):
600    for index in indices:
601        if index.has_form("Span", None):
602            yield _parts_span_selector(index)
603        elif index.get_name() == "System`All":
604            yield _parts_all_selector()
605        elif index.has_form("List", None):
606            yield _parts_sequence_selector(index.leaves)
607        elif isinstance(index, Integer):
608            yield _parts_sequence_selector(index), lambda x: x[0]
609        else:
610            raise MessageException("Part", "pspec", index)
611
612
613def _list_parts(items, selectors, heads, evaluation, assignment):
614    if not selectors:
615        for item in items:
616            yield item
617    else:
618        selector = selectors[0]
619        if isinstance(selector, tuple):
620            select, unwrap = selector
621        else:
622            select = selector
623            unwrap = None
624
625        for item in items:
626            selected = list(select(item))
627
628            picked = list(
629                _list_parts(selected, selectors[1:], heads, evaluation, assignment)
630            )
631
632            if unwrap is None:
633                if assignment:
634                    expr = Expression(item.head, *picked)
635                    expr.original = None
636                    expr.set_positions()
637                else:
638                    expr = item.restructure(item.head, picked, evaluation)
639
640                yield expr
641            else:
642                yield unwrap(picked)
643
644
645def _parts(items, selectors, evaluation, assignment=False):
646    heads = {}
647    return list(_list_parts([items], list(selectors), heads, evaluation, assignment))[0]
648
649
650def walk_parts(list_of_list, indices, evaluation, assign_list=None):
651    walk_list = list_of_list[0]
652
653    if assign_list is not None:
654        # this double copying is needed to make the current logic in
655        # the assign_list and its access to original work.
656
657        walk_list = walk_list.copy()
658        walk_list.set_positions()
659        list_of_list = [walk_list]
660
661        walk_list = walk_list.copy()
662        walk_list.set_positions()
663
664    indices = [index.evaluate(evaluation) for index in indices]
665
666    try:
667        result = _parts(
668            walk_list, _part_selectors(indices), evaluation, assign_list is not None
669        )
670    except MessageException as e:
671        e.message(evaluation)
672        return False
673
674    if assign_list is not None:
675
676        def replace_item(all, item, new):
677            if item.position is None:
678                all[0] = new
679            else:
680                item.position.replace(new)
681
682        def process_level(item, assignment):
683            if item.is_atom():
684                replace_item(list_of_list, item.original, assignment)
685            elif assignment.get_head_name() != "System`List" or len(item.leaves) != len(
686                assignment.leaves
687            ):
688                if item.original:
689                    replace_item(list_of_list, item.original, assignment)
690                else:
691                    for leaf in item.leaves:
692                        process_level(leaf, assignment)
693            else:
694                for sub_item, sub_assignment in zip(item.leaves, assignment.leaves):
695                    process_level(sub_item, sub_assignment)
696
697        process_level(result, assign_list)
698
699        result = list_of_list[0]
700        result.clear_cache()
701
702    return result
703
704
705def is_in_level(current, depth, start=1, stop=None):
706    if stop is None:
707        stop = current
708    if start < 0:
709        start += current + depth + 1
710    if stop < 0:
711        stop += current + depth + 1
712    return start <= current <= stop
713
714
715def walk_levels(
716    expr,
717    start=1,
718    stop=None,
719    current=0,
720    heads=False,
721    callback=lambda l: l,
722    include_pos=False,
723    cur_pos=[],
724):
725    if expr.is_atom():
726        depth = 0
727        new_expr = expr
728    else:
729        depth = 0
730        if heads:
731            head, head_depth = walk_levels(
732                expr.head,
733                start,
734                stop,
735                current + 1,
736                heads,
737                callback,
738                include_pos,
739                cur_pos + [0],
740            )
741        else:
742            head = expr.head
743        leaves = []
744        for index, leaf in enumerate(expr.leaves):
745            leaf, leaf_depth = walk_levels(
746                leaf,
747                start,
748                stop,
749                current + 1,
750                heads,
751                callback,
752                include_pos,
753                cur_pos + [index + 1],
754            )
755            if leaf_depth + 1 > depth:
756                depth = leaf_depth + 1
757            leaves.append(leaf)
758        new_expr = Expression(head, *leaves)
759    if is_in_level(current, depth, start, stop):
760        if include_pos:
761            new_expr = callback(new_expr, cur_pos)
762        else:
763            new_expr = callback(new_expr)
764    return new_expr, depth
765
766
767def python_levelspec(levelspec):
768    def value_to_level(expr):
769        value = expr.get_int_value()
770        if value is None:
771            if expr == Expression("DirectedInfinity", 1):
772                return None
773            else:
774                raise InvalidLevelspecError
775        else:
776            return value
777
778    if levelspec.has_form("List", None):
779        values = [value_to_level(leaf) for leaf in levelspec.leaves]
780        if len(values) == 1:
781            return values[0], values[0]
782        elif len(values) == 2:
783            return values[0], values[1]
784        else:
785            raise InvalidLevelspecError
786    elif isinstance(levelspec, Symbol) and levelspec.get_name() == "System`All":
787        return 0, None
788    else:
789        return 1, value_to_level(levelspec)
790
791
792class Level(Builtin):
793    """
794    <dl>
795    <dt>'Level[$expr$, $levelspec$]'
796        <dd>gives a list of all subexpressions of $expr$ at the
797        level(s) specified by $levelspec$.
798    </dl>
799
800    Level uses standard level specifications:
801
802    <dl>
803    <dt>$n$
804        <dd>levels 1 through $n$
805    <dt>'Infinity'
806        <dd>all levels from level 1
807    <dt>'{$n$}'
808        <dd>level $n$ only
809    <dt>'{$m$, $n$}'
810        <dd>levels $m$ through $n$
811    </dl>
812
813    Level 0 corresponds to the whole expression.
814
815    A negative level '-$n$' consists of parts with depth $n$.
816
817    Level -1 is the set of atoms in an expression:
818    >> Level[a + b ^ 3 * f[2 x ^ 2], {-1}]
819     = {a, b, 3, 2, x, 2}
820
821    >> Level[{{{{a}}}}, 3]
822     = {{a}, {{a}}, {{{a}}}}
823    >> Level[{{{{a}}}}, -4]
824     = {{{{a}}}}
825    >> Level[{{{{a}}}}, -5]
826     = {}
827
828    >> Level[h0[h1[h2[h3[a]]]], {0, -1}]
829     = {a, h3[a], h2[h3[a]], h1[h2[h3[a]]], h0[h1[h2[h3[a]]]]}
830
831    Use the option 'Heads -> True' to include heads:
832    >> Level[{{{{a}}}}, 3, Heads -> True]
833     = {List, List, List, {a}, {{a}}, {{{a}}}}
834    >> Level[x^2 + y^3, 3, Heads -> True]
835     = {Plus, Power, x, 2, x ^ 2, Power, y, 3, y ^ 3}
836
837    >> Level[a ^ 2 + 2 * b, {-1}, Heads -> True]
838     = {Plus, Power, a, 2, Times, 2, b}
839    >> Level[f[g[h]][x], {-1}, Heads -> True]
840     = {f, g, h, x}
841    >> Level[f[g[h]][x], {-2, -1}, Heads -> True]
842     = {f, g, h, g[h], x, f[g[h]][x]}
843    """
844
845    options = {
846        "Heads": "False",
847    }
848
849    def apply(self, expr, ls, evaluation, options={}):
850        "Level[expr_, ls_, OptionsPattern[Level]]"
851
852        try:
853            start, stop = python_levelspec(ls)
854        except InvalidLevelspecError:
855            evaluation.message("Level", "level", ls)
856            return
857        result = []
858
859        def callback(level):
860            result.append(level)
861            return level
862
863        heads = self.get_option(options, "Heads", evaluation).is_true()
864        walk_levels(expr, start, stop, heads=heads, callback=callback)
865        return Expression(SymbolList, *result)
866
867
868class LevelQ(Test):
869    """
870    <dl>
871    <dt>'LevelQ[$expr$]'
872        <dd>tests whether $expr$ is a valid level specification.
873    </dl>
874
875    >> LevelQ[2]
876     = True
877    >> LevelQ[{2, 4}]
878     = True
879    >> LevelQ[Infinity]
880     = True
881    >> LevelQ[a + b]
882     = False
883    """
884
885    def test(self, ls):
886        try:
887            start, stop = python_levelspec(ls)
888            return True
889        except InvalidLevelspecError:
890            return False
891
892
893def python_seq(start, stop, step, length):
894    """
895    Converts mathematica sequence tuple to python slice object.
896
897    Based on David Mashburn's generic slice:
898    https://gist.github.com/davidmashburn/9764309
899    """
900    if step == 0:
901        return None
902
903    # special empty case
904    if stop is None and length is not None:
905        empty_stop = length
906    else:
907        empty_stop = stop
908    if start is not None and empty_stop + 1 == start and step > 0:
909        return slice(0, 0, 1)
910
911    if start == 0 or stop == 0:
912        return None
913
914    # wrap negative values to postive and convert from 1-based to 0-based
915    if start < 0:
916        start += length
917    else:
918        start -= 1
919
920    if stop is None:
921        if step < 0:
922            stop = 0
923        else:
924            stop = length - 1
925    elif stop < 0:
926        stop += length
927    else:
928        assert stop > 0
929        stop -= 1
930
931    # check bounds
932    if (
933        not 0 <= start < length
934        or not 0 <= stop < length
935        or step > 0
936        and start - stop > 1
937        or step < 0
938        and stop - start > 1
939    ):  # nopep8
940        return None
941
942    # include the stop value
943    if step > 0:
944        stop += 1
945    else:
946        stop -= 1
947        if stop == -1:
948            stop = None
949        if start == 0:
950            start = None
951
952    return slice(start, stop, step)
953
954
955def convert_seq(seq):
956    """
957    converts a sequence specification into a (start, stop, step) tuple.
958    returns None on failure
959    """
960    start, stop, step = 1, None, 1
961    name = seq.get_name()
962    value = seq.get_int_value()
963    if name == "System`All":
964        pass
965    elif name == "System`None":
966        stop = 0
967    elif value is not None:
968        if value > 0:
969            stop = value
970        else:
971            start = value
972    elif seq.has_form("List", 1, 2, 3):
973        if len(seq.leaves) == 1:
974            start = stop = seq.leaves[0].get_int_value()
975            if stop is None:
976                return None
977        else:
978            start = seq.leaves[0].get_int_value()
979            stop = seq.leaves[1].get_int_value()
980            if start is None or stop is None:
981                return None
982        if len(seq.leaves) == 3:
983            step = seq.leaves[2].get_int_value()
984            if step is None:
985                return None
986    else:
987        return None
988    return (start, stop, step)
989
990
991class Part(Builtin):
992    """
993    <dl>
994      <dt>'Part[$expr$, $i$]'
995      <dd>returns part $i$ of $expr$.
996    </dl>
997
998    Extract an element from a list:
999    >> A = {a, b, c, d};
1000    >> A[[3]]
1001     = c
1002
1003    Negative indices count from the end:
1004    >> {a, b, c}[[-2]]
1005     = b
1006
1007    'Part' can be applied on any expression, not necessarily lists.
1008    >> (a + b + c)[[2]]
1009     = b
1010    '$expr$[[0]]' gives the head of $expr$:
1011    >> (a + b + c)[[0]]
1012     = Plus
1013
1014    Parts of nested lists:
1015    >> M = {{a, b}, {c, d}};
1016    >> M[[1, 2]]
1017     = b
1018
1019    You can use 'Span' to specify a range of parts:
1020    >> {1, 2, 3, 4}[[2;;4]]
1021     = {2, 3, 4}
1022    >> {1, 2, 3, 4}[[2;;-1]]
1023     = {2, 3, 4}
1024
1025    A list of parts extracts elements at certain indices:
1026    >> {a, b, c, d}[[{1, 3, 3}]]
1027     = {a, c, c}
1028
1029    Get a certain column of a matrix:
1030    >> B = {{a, b, c}, {d, e, f}, {g, h, i}};
1031    >> B[[;;, 2]]
1032     = {b, e, h}
1033
1034    Extract a submatrix of 1st and 3rd row and the two last columns:
1035    >> B = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};
1036
1037    >> B[[{1, 3}, -2;;-1]]
1038     = {{2, 3}, {8, 9}}
1039
1040    The 3d column of a matrix:
1041    >> {{a, b, c}, {d, e, f}, {g, h, i}}[[All, 3]]
1042     = {c, f, i}
1043
1044    Further examples:
1045    >> (a+b+c+d)[[-1;;-2]]
1046     = 0
1047    >> x[[2]]
1048     : Part specification is longer than depth of object.
1049     = x[[2]]
1050
1051    Assignments to parts are possible:
1052    >> B[[;;, 2]] = {10, 11, 12}
1053     = {10, 11, 12}
1054    >> B
1055     = {{1, 10, 3}, {4, 11, 6}, {7, 12, 9}}
1056    >> B[[;;, 3]] = 13
1057     = 13
1058    >> B
1059     = {{1, 10, 13}, {4, 11, 13}, {7, 12, 13}}
1060    >> B[[1;;-2]] = t;
1061    >> B
1062     = {t, t, {7, 12, 13}}
1063
1064    >> F = Table[i*j*k, {i, 1, 3}, {j, 1, 3}, {k, 1, 3}];
1065    >> F[[;; All, 2 ;; 3, 2]] = t;
1066    >> F
1067     = {{{1, 2, 3}, {2, t, 6}, {3, t, 9}}, {{2, 4, 6}, {4, t, 12}, {6, t, 18}}, {{3, 6, 9}, {6, t, 18}, {9, t, 27}}}
1068    >> F[[;; All, 1 ;; 2, 3 ;; 3]] = k;
1069    >> F
1070     = {{{1, 2, k}, {2, t, k}, {3, t, 9}}, {{2, 4, k}, {4, t, k}, {6, t, 18}}, {{3, 6, k}, {6, t, k}, {9, t, 27}}}
1071
1072    Of course, part specifications have precedence over most arithmetic operations:
1073    >> A[[1]] + B[[2]] + C[[3]] // Hold // FullForm
1074     = Hold[Plus[Part[A, 1], Part[B, 2], Part[C, 3]]]
1075
1076    #> a = {2,3,4}; i = 1; a[[i]] = 0; a
1077     = {0, 3, 4}
1078
1079    ## Negative step
1080    #> {1,2,3,4,5}[[3;;1;;-1]]
1081     = {3, 2, 1}
1082
1083    #> {1, 2, 3, 4, 5}[[;; ;; -1]]      (* MMA bug *)
1084     = {5, 4, 3, 2, 1}
1085
1086    #> Range[11][[-3 ;; 2 ;; -2]]
1087     = {9, 7, 5, 3}
1088    #> Range[11][[-3 ;; -7 ;; -3]]
1089     = {9, 6}
1090    #> Range[11][[7 ;; -7;; -2]]
1091     = {7, 5}
1092
1093    #> {1, 2, 3, 4}[[1;;3;;-1]]
1094     : Cannot take positions 1 through 3 in {1, 2, 3, 4}.
1095     = {1, 2, 3, 4}[[1 ;; 3 ;; -1]]
1096    #> {1, 2, 3, 4}[[3;;1]]
1097     : Cannot take positions 3 through 1 in {1, 2, 3, 4}.
1098     = {1, 2, 3, 4}[[3 ;; 1]]
1099    """
1100
1101    attributes = ("NHoldRest", "ReadProtected")
1102
1103    def apply_makeboxes(self, list, i, f, evaluation):
1104        """MakeBoxes[Part[list_, i___],
1105        f:StandardForm|TraditionalForm|OutputForm|InputForm]"""
1106
1107        i = i.get_sequence()
1108        list = Expression(SymbolMakeBoxes, list, f)
1109        if f.get_name() in ("System`OutputForm", "System`InputForm"):
1110            open, close = "[[", "]]"
1111        else:
1112            open, close = "\u301a", "\u301b"
1113        indices = list_boxes(i, f, open, close)
1114        result = Expression("RowBox", Expression(SymbolList, list, *indices))
1115        return result
1116
1117    def apply(self, list, i, evaluation):
1118        "Part[list_, i___]"
1119
1120        indices = i.get_sequence()
1121        # How to deal with ByteArrays
1122        if list.get_head_name() == "System`ByteArray":
1123            list = list.evaluate(evaluation)
1124            if len(indices) > 1:
1125                print(
1126                    "Part::partd1: Depth of object ByteArray[<3>] "
1127                    + "is not sufficient for the given part specification."
1128                )
1129                return
1130            idx = indices[0]
1131            if idx.get_head_name() == "System`Integer":
1132                idx = idx.get_int_value()
1133                if idx == 0:
1134                    return Symbol("System`ByteArray")
1135                data = list._leaves[0].value
1136                lendata = len(data)
1137                if idx < 0:
1138                    idx = data - idx
1139                    if idx < 0:
1140                        evaluation.message("Part", "partw", i, list)
1141                        return
1142                else:
1143                    idx = idx - 1
1144                    if idx > lendata:
1145                        evaluation.message("Part", "partw", i, list)
1146                        return
1147                return Integer(data[idx])
1148            if idx == Symbol("System`All"):
1149                return list
1150            # TODO: handling ranges and lists...
1151            evaluation.message("Part", "notimplemented")
1152            return
1153
1154        # Otherwise...
1155        result = walk_parts([list], indices, evaluation)
1156        if result:
1157            return result
1158
1159
1160class Partition(Builtin):
1161    """
1162    <dl>
1163    <dt>'Partition[$list$, $n$]'
1164        <dd>partitions $list$ into sublists of length $n$.
1165    <dt>'Parition[$list$, $n$, $d$]'
1166        <dd>partitions $list$ into sublists of length $n$ which
1167        overlap $d$ indicies.
1168    </dl>
1169
1170    >> Partition[{a, b, c, d, e, f}, 2]
1171     = {{a, b}, {c, d}, {e, f}}
1172
1173    >> Partition[{a, b, c, d, e, f}, 3, 1]
1174     = {{a, b, c}, {b, c, d}, {c, d, e}, {d, e, f}}
1175
1176    #> Partition[{a, b, c, d, e}, 2]
1177     = {{a, b}, {c, d}}
1178    """
1179
1180    # TODO: Nested list length specifications
1181    """
1182    >> Partition[{{11, 12, 13}, {21, 22, 23}, {31, 32, 33}}, {2, 2}, 1]
1183     = {{{{11, 12}, {21, 22}}, {{12, 13}, {22, 23}}}, {{{21, 22}, {31, 32}}, {{22, 23}, {32, 33}}}}
1184    """
1185
1186    rules = {
1187        "Parition[list_, n_, d_, k]": "Partition[list, n, d, {k, k}]",
1188    }
1189
1190    def _partition(self, expr, n, d, evaluation):
1191        assert n > 0 and d > 0
1192
1193        inner = structure("List", expr, evaluation)
1194        outer = structure("List", inner, evaluation)
1195
1196        make_slice = inner.slice
1197
1198        def slices():
1199            leaves = expr.leaves
1200            for lower in range(0, len(leaves), d):
1201                upper = lower + n
1202
1203                chunk = leaves[lower:upper]
1204                if len(chunk) != n:
1205                    continue
1206
1207                yield make_slice(expr, slice(lower, upper))
1208
1209        return outer(slices())
1210
1211    def apply_no_overlap(self, l, n, evaluation):
1212        "Partition[l_List, n_Integer]"
1213        # TODO: Error checking
1214        return self._partition(l, n.get_int_value(), n.get_int_value(), evaluation)
1215
1216    def apply(self, l, n, d, evaluation):
1217        "Partition[l_List, n_Integer, d_Integer]"
1218        # TODO: Error checking
1219        return self._partition(l, n.get_int_value(), d.get_int_value(), evaluation)
1220
1221
1222class Extract(Builtin):
1223    """
1224    <dl>
1225    <dt>'Extract[$expr$, $list$]'
1226        <dd>extracts parts of $expr$ specified by $list$.
1227    <dt>'Extract[$expr$, {$list1$, $list2$, ...}]'
1228        <dd>extracts a list of parts.
1229    </dl>
1230
1231    'Extract[$expr$, $i$, $j$, ...]' is equivalent to 'Part[$expr$, {$i$, $j$, ...}]'.
1232
1233    >> Extract[a + b + c, {2}]
1234     = b
1235    >> Extract[{{a, b}, {c, d}}, {{1}, {2, 2}}]
1236     = {{a, b}, d}
1237    """
1238
1239    attributes = ("NHoldRest",)
1240
1241    rules = {
1242        "Extract[expr_, list_List]": "Part[expr, Sequence @@ list]",
1243        "Extract[expr_, {lists___List}]": "Extract[expr, #]& /@ {lists}",
1244    }
1245
1246
1247class First(Builtin):
1248    """
1249    <dl>
1250    <dt>'First[$expr$]'
1251        <dd>returns the first element in $expr$.
1252    </dl>
1253
1254    'First[$expr$]' is equivalent to '$expr$[[1]]'.
1255
1256    >> First[{a, b, c}]
1257     = a
1258    >> First[a + b + c]
1259     = a
1260    >> First[x]
1261     : Nonatomic expression expected.
1262     = First[x]
1263    """
1264
1265    def apply(self, expr, evaluation):
1266        "First[expr_]"
1267
1268        if expr.is_atom():
1269            evaluation.message("First", "normal")
1270            return
1271        return expr.leaves[0]
1272
1273
1274class Last(Builtin):
1275    """
1276    <dl>
1277    <dt>'Last[$expr$]'
1278        <dd>returns the last element in $expr$.
1279    </dl>
1280
1281    'Last[$expr$]' is equivalent to '$expr$[[-1]]'.
1282
1283    >> Last[{a, b, c}]
1284     = c
1285    >> Last[x]
1286     : Nonatomic expression expected.
1287     = Last[x]
1288    """
1289
1290    def apply(self, expr, evaluation):
1291        "Last[expr_]"
1292
1293        if expr.is_atom():
1294            evaluation.message("Last", "normal")
1295            return
1296        return expr.leaves[-1]
1297
1298
1299class Most(Builtin):
1300    """
1301    <dl>
1302    <dt>'Most[$expr$]'
1303        <dd>returns $expr$ with the last element removed.
1304    </dl>
1305
1306    'Most[$expr$]' is equivalent to '$expr$[[;;-2]]'.
1307
1308    >> Most[{a, b, c}]
1309     = {a, b}
1310    >> Most[a + b + c]
1311     = a + b
1312    >> Most[x]
1313     : Nonatomic expression expected.
1314     = Most[x]
1315
1316    #> A[x__] := 7 /; Length[{x}] == 3;
1317    #> Most[A[1, 2, 3, 4]]
1318     = 7
1319    #> ClearAll[A];
1320    """
1321
1322    def apply(self, expr, evaluation):
1323        "Most[expr_]"
1324
1325        if expr.is_atom():
1326            evaluation.message("Most", "normal")
1327            return
1328        return expr.slice(expr.head, slice(0, -1), evaluation)
1329
1330
1331class Rest(Builtin):
1332    """
1333    <dl>
1334    <dt>'Rest[$expr$]'
1335        <dd>returns $expr$ with the first element removed.
1336    </dl>
1337
1338    'Rest[$expr$]' is equivalent to '$expr$[[2;;]]'.
1339
1340    >> Rest[{a, b, c}]
1341     = {b, c}
1342    >> Rest[a + b + c]
1343     = b + c
1344    >> Rest[x]
1345     : Nonatomic expression expected.
1346     = Rest[x]
1347    """
1348
1349    def apply(self, expr, evaluation):
1350        "Rest[expr_]"
1351
1352        if expr.is_atom():
1353            evaluation.message("Rest", "normal")
1354            return
1355        return expr.slice(expr.head, slice(1, len(expr.leaves)), evaluation)
1356
1357
1358class ReplacePart(Builtin):
1359    """
1360    <dl>
1361    <dt>'ReplacePart[$expr$, $i$ -> $new$]'
1362        <dd>replaces part $i$ in $expr$ with $new$.
1363    <dt>'ReplacePart[$expr$, {{$i$, $j$} -> $e1$, {$k$, $l$} -> $e2$}]'
1364        <dd>replaces parts $i$ and $j$ with $e1$, and parts $k$ and
1365        $l$ with $e2$.
1366    </dl>
1367
1368    >> ReplacePart[{a, b, c}, 1 -> t]
1369     = {t, b, c}
1370    >> ReplacePart[{{a, b}, {c, d}}, {2, 1} -> t]
1371     = {{a, b}, {t, d}}
1372    >> ReplacePart[{{a, b}, {c, d}}, {{2, 1} -> t, {1, 1} -> t}]
1373     = {{t, b}, {t, d}}
1374    >> ReplacePart[{a, b, c}, {{1}, {2}} -> t]
1375     = {t, t, c}
1376
1377    Delayed rules are evaluated once for each replacement:
1378    >> n = 1;
1379    >> ReplacePart[{a, b, c, d}, {{1}, {3}} :> n++]
1380     = {1, b, 2, d}
1381
1382    Non-existing parts are simply ignored:
1383    >> ReplacePart[{a, b, c}, 4 -> t]
1384     = {a, b, c}
1385    You can replace heads by replacing part 0:
1386    >> ReplacePart[{a, b, c}, 0 -> Times]
1387     = a b c
1388    (This is equivalent to 'Apply'.)
1389
1390    Negative part numbers count from the end:
1391    >> ReplacePart[{a, b, c}, -1 -> t]
1392     = {a, b, t}
1393    """
1394
1395    messages = {
1396        "reps": "`1` is not a list of replacement rules.",
1397    }
1398
1399    rules = {
1400        "ReplacePart[expr_, (Rule|RuleDelayed)[i_, new_]]": (
1401            "ReplacePart[expr, {i -> new}]"
1402        ),
1403        "ReplacePart[expr_, Pattern[rule, "
1404        "Rule|RuleDelayed][{indices___?(Head[#]===List&)}, new_]]": (
1405            "ReplacePart[expr, rule[#, new]& /@ {indices}]"
1406        ),
1407    }
1408
1409    def apply(self, expr, replacements, evaluation):
1410        "ReplacePart[expr_, {replacements___}]"
1411
1412        new_expr = expr.copy()
1413        replacements = replacements.get_sequence()
1414        for replacement in replacements:
1415            if not replacement.has_form("Rule", 2) and not replacement.has_form(  # noqa
1416                "RuleDelayed", 2
1417            ):
1418                evaluation.message(
1419                    "ReplacePart", "reps", Expression(SymbolList, *replacements)
1420                )
1421                return
1422            position = replacement.leaves[0]
1423            replace = replacement.leaves[1]
1424            if position.has_form("List", None):
1425                position = position.get_mutable_leaves()
1426            else:
1427                position = [position]
1428            for index, pos in enumerate(position):
1429                value = pos.get_int_value()
1430                if value is None:
1431                    position = None
1432                    break
1433                else:
1434                    position[index] = value
1435            if position is None:
1436                continue
1437            try:
1438                if replacement.get_head_name() == "System`RuleDelayed":
1439                    replace_value = replace.evaluate(evaluation)
1440                else:
1441                    replace_value = replace
1442                set_part(new_expr, position, replace_value)
1443            except PartError:
1444                pass
1445
1446        return new_expr
1447
1448
1449class FirstPosition(Builtin):
1450    """
1451    <dl>
1452    <dt>'FirstPosition[$expr$, $pattern$]'
1453        <dd>gives the position of the first element in $expr$ that matches $pattern$, or Missing["NotFound"] if no such element is found.
1454    <dt>'FirstPosition[$expr$, $pattern$, $default$]'
1455        <dd>gives default if no element matching $pattern$ is found.
1456    <dt>'FirstPosition[$expr$, $pattern$, $default$, $levelspec$]'
1457        <dd>finds only objects that appear on levels specified by $levelspec$.
1458    </dl>
1459
1460    >> FirstPosition[{a, b, a, a, b, c, b}, b]
1461     = {2}
1462
1463    >> FirstPosition[{{a, a, b}, {b, a, a}, {a, b, a}}, b]
1464     = {1, 3}
1465
1466    >> FirstPosition[{x, y, z}, b]
1467     = Missing[NotFound]
1468
1469    Find the first position at which x^2 to appears:
1470    >> FirstPosition[{1 + x^2, 5, x^4, a + (1 + x^2)^2}, x^2]
1471     = {1, 2}
1472
1473    #> FirstPosition[{1, 2, 3}, _?StringQ, "NoStrings"]
1474     = NoStrings
1475
1476    #> FirstPosition[a, a]
1477     = {}
1478
1479    #> FirstPosition[{{{1, 2}, {2, 3}, {3, 1}}, {{1, 2}, {2, 3}, {3, 1}}},3]
1480     = {1, 2, 2}
1481
1482    #> FirstPosition[{{1, {2, 1}}, {2, 3}, {3, 1}}, 2, Missing["NotFound"],2]
1483     = {2, 1}
1484
1485    #> FirstPosition[{{1, {2, 1}}, {2, 3}, {3, 1}}, 2, Missing["NotFound"],4]
1486     = {1, 2, 1}
1487
1488    #> FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing["NotFound"], {1}]
1489     = Missing[NotFound]
1490
1491    #> FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing["NotFound"], 0]
1492     = Missing[NotFound]
1493
1494    #> FirstPosition[{{1, 2}, {1, {2, 1}}, {2, 3}}, 2, Missing["NotFound"], {3}]
1495     = {2, 2, 1}
1496
1497    #> FirstPosition[{{1, 2}, {1, {2, 1}}, {2, 3}}, 2, Missing["NotFound"], 3]
1498     = {1, 2}
1499
1500    #> FirstPosition[{{1, 2}, {1, {2, 1}}, {2, 3}}, 2,  Missing["NotFound"], {}]
1501     = {1, 2}
1502
1503    #> FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing["NotFound"], {1, 2, 3}]
1504     : Level specification {1, 2, 3} is not of the form n, {n}, or {m, n}.
1505     = FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing[NotFound], {1, 2, 3}]
1506
1507    #> FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing["NotFound"], a]
1508     : Level specification a is not of the form n, {n}, or {m, n}.
1509     = FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing[NotFound], a]
1510
1511    #> FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing["NotFound"], {1, a}]
1512     : Level specification {1, a} is not of the form n, {n}, or {m, n}.
1513     = FirstPosition[{{1, 2}, {2, 3}, {3, 1}}, 3, Missing[NotFound], {1, a}]
1514
1515    """
1516
1517    messages = {
1518        "level": "Level specification `1` is not of the form n, {n}, or {m, n}.",
1519    }
1520
1521    def apply(
1522        self, expr, pattern, evaluation, default=None, minLevel=None, maxLevel=None
1523    ):
1524        "FirstPosition[expr_, pattern_]"
1525
1526        if expr == pattern:
1527            return Expression(SymbolList)
1528
1529        result = []
1530
1531        def check_pattern(input_list, pat, result, beginLevel):
1532            for i in range(0, len(input_list.leaves)):
1533                nested_level = beginLevel
1534                result.append(i + 1)
1535                if input_list.leaves[i] == pat:
1536                    # found the pattern
1537                    if minLevel is None or nested_level >= minLevel:
1538                        return True
1539
1540                else:
1541                    if isinstance(input_list.leaves[i], Expression) and (
1542                        maxLevel is None or maxLevel > nested_level
1543                    ):
1544                        nested_level = nested_level + 1
1545                        if check_pattern(
1546                            input_list.leaves[i], pat, result, nested_level
1547                        ):
1548                            return True
1549
1550                result.pop()
1551            return False
1552
1553        is_found = False
1554        if isinstance(expr, Expression) and (maxLevel is None or maxLevel > 0):
1555            is_found = check_pattern(expr, pattern, result, 1)
1556        if is_found:
1557            return Expression(SymbolList, *result)
1558        else:
1559            return Expression("Missing", "NotFound") if default is None else default
1560
1561    def apply_default(self, expr, pattern, default, evaluation):
1562        "FirstPosition[expr_, pattern_, default_]"
1563        return self.apply(expr, pattern, evaluation, default=default)
1564
1565    def apply_level(self, expr, pattern, default, level, evaluation):
1566        "FirstPosition[expr_, pattern_, default_, level_]"
1567
1568        def is_interger_list(expr_list):
1569            return all(
1570                isinstance(expr_list.leaves[i], Integer)
1571                for i in range(len(expr_list.leaves))
1572            )
1573
1574        if level.has_form("List", None):
1575            len_list = len(level.leaves)
1576            if len_list > 2 or not is_interger_list(level):
1577                return evaluation.message("FirstPosition", "level", level)
1578            elif len_list == 0:
1579                min_Level = max_Level = None
1580            elif len_list == 1:
1581                min_Level = max_Level = level.leaves[0].get_int_value()
1582            elif len_list == 2:
1583                min_Level = level.leaves[0].get_int_value()
1584                max_Level = level.leaves[1].get_int_value()
1585        elif isinstance(level, Integer):
1586            min_Level = 0
1587            max_Level = level.get_int_value()
1588        else:
1589            return evaluation.message("FirstPosition", "level", level)
1590
1591        return self.apply(
1592            expr,
1593            pattern,
1594            evaluation,
1595            default=default,
1596            minLevel=min_Level,
1597            maxLevel=max_Level,
1598        )
1599
1600
1601def _drop_take_selector(name, seq, sliced):
1602    seq_tuple = convert_seq(seq)
1603    if seq_tuple is None:
1604        raise MessageException(name, "seqs", seq)
1605
1606    def select(inner):
1607        start, stop, step = seq_tuple
1608        if inner.is_atom():
1609            py_slice = None
1610        else:
1611            py_slice = python_seq(start, stop, step, len(inner.leaves))
1612        if py_slice is None:
1613            if stop is None:
1614                stop = Symbol("Infinity")
1615            raise MessageException(name, name.lower(), start, stop, inner)
1616        return sliced(inner.leaves, py_slice)
1617
1618    return select
1619
1620
1621def _take_span_selector(seq):
1622    return _drop_take_selector("Take", seq, lambda x, s: x[s])
1623
1624
1625def _drop_span_selector(seq):
1626    def sliced(x, s):
1627        y = list(x[:])
1628        del y[s]
1629        return y
1630
1631    return _drop_take_selector("Drop", seq, sliced)
1632
1633
1634class Take(Builtin):
1635    """
1636    <dl>
1637    <dt>'Take[$expr$, $n$]'
1638        <dd>returns $expr$ with all but the first $n$ leaves removed.
1639    </dl>
1640
1641    >> Take[{a, b, c, d}, 3]
1642     = {a, b, c}
1643    >> Take[{a, b, c, d}, -2]
1644     = {c, d}
1645    >> Take[{a, b, c, d, e}, {2, -2}]
1646     = {b, c, d}
1647
1648    Take a submatrix:
1649    >> A = {{a, b, c}, {d, e, f}};
1650    >> Take[A, 2, 2]
1651     = {{a, b}, {d, e}}
1652
1653    Take a single column:
1654    >> Take[A, All, {2}]
1655     = {{b}, {e}}
1656
1657    #> Take[Range[10], {8, 2, -1}]
1658     = {8, 7, 6, 5, 4, 3, 2}
1659    #> Take[Range[10], {-3, -7, -2}]
1660     = {8, 6, 4}
1661
1662    #> Take[Range[6], {-5, -2, -2}]
1663     : Cannot take positions -5 through -2 in {1, 2, 3, 4, 5, 6}.
1664     = Take[{1, 2, 3, 4, 5, 6}, {-5, -2, -2}]
1665
1666    #> Take[l, {-1}]
1667     : Nonatomic expression expected at position 1 in Take[l, {-1}].
1668     = Take[l, {-1}]
1669
1670    ## Empty case
1671    #> Take[{1, 2, 3, 4, 5}, {-1, -2}]
1672     = {}
1673    #> Take[{1, 2, 3, 4, 5}, {0, -1}]
1674     = {}
1675    #> Take[{1, 2, 3, 4, 5}, {1, 0}]
1676     = {}
1677    #> Take[{1, 2, 3, 4, 5}, {2, 1}]
1678     = {}
1679    #> Take[{1, 2, 3, 4, 5}, {1, 0, 2}]
1680     = {}
1681    #> Take[{1, 2, 3, 4, 5}, {1, 0, -1}]
1682     : Cannot take positions 1 through 0 in {1, 2, 3, 4, 5}.
1683     = Take[{1, 2, 3, 4, 5}, {1, 0, -1}]
1684    """
1685
1686    messages = {
1687        "normal": "Nonatomic expression expected at position `1` in `2`.",
1688    }
1689
1690    def apply(self, items, seqs, evaluation):
1691        "Take[items_, seqs___]"
1692
1693        seqs = seqs.get_sequence()
1694
1695        if items.is_atom():
1696            return evaluation.message(
1697                "Take", "normal", 1, Expression("Take", items, *seqs)
1698            )
1699
1700        try:
1701            return _parts(items, [_take_span_selector(seq) for seq in seqs], evaluation)
1702        except MessageException as e:
1703            e.message(evaluation)
1704
1705
1706class Drop(Builtin):
1707    """
1708    <dl>
1709    <dt>'Drop[$expr$, $n$]'
1710        <dd>returns $expr$ with the first $n$ leaves removed.
1711    </dl>
1712
1713    >> Drop[{a, b, c, d}, 3]
1714     = {d}
1715    >> Drop[{a, b, c, d}, -2]
1716     = {a, b}
1717    >> Drop[{a, b, c, d, e}, {2, -2}]
1718     = {a, e}
1719
1720    Drop a submatrix:
1721    >> A = Table[i*10 + j, {i, 4}, {j, 4}]
1722     = {{11, 12, 13, 14}, {21, 22, 23, 24}, {31, 32, 33, 34}, {41, 42, 43, 44}}
1723    >> Drop[A, {2, 3}, {2, 3}]
1724     = {{11, 14}, {41, 44}}
1725
1726    #> Drop[Range[10], {-2, -6, -3}]
1727     = {1, 2, 3, 4, 5, 7, 8, 10}
1728    #> Drop[Range[10], {10, 1, -3}]
1729     = {2, 3, 5, 6, 8, 9}
1730
1731    #> Drop[Range[6], {-5, -2, -2}]
1732     : Cannot drop positions -5 through -2 in {1, 2, 3, 4, 5, 6}.
1733     = Drop[{1, 2, 3, 4, 5, 6}, {-5, -2, -2}]
1734    """
1735
1736    messages = {
1737        "normal": "Nonatomic expression expected at position `1` in `2`.",
1738        "drop": "Cannot drop positions `1` through `2` in `3`.",
1739    }
1740
1741    def apply(self, items, seqs, evaluation):
1742        "Drop[items_, seqs___]"
1743
1744        seqs = seqs.get_sequence()
1745
1746        if items.is_atom():
1747            return evaluation.message(
1748                "Drop", "normal", 1, Expression("Drop", items, *seqs)
1749            )
1750
1751        try:
1752            return _parts(items, [_drop_span_selector(seq) for seq in seqs], evaluation)
1753        except MessageException as e:
1754            e.message(evaluation)
1755
1756
1757class Select(Builtin):
1758    """
1759    <dl>
1760    <dt>'Select[{$e1$, $e2$, ...}, $f$]'
1761        <dd>returns a list of the elements $ei$ for which $f$[$ei$]
1762        returns 'True'.
1763    </dl>
1764
1765    Find numbers greater than zero:
1766    >> Select[{-3, 0, 1, 3, a}, #>0&]
1767     = {1, 3}
1768
1769    'Select' works on an expression with any head:
1770    >> Select[f[a, 2, 3], NumberQ]
1771     = f[2, 3]
1772
1773    >> Select[a, True]
1774     : Nonatomic expression expected.
1775     = Select[a, True]
1776
1777    #> A[x__] := 31415 /; Length[{x}] == 3;
1778    #> Select[A[5, 2, 7, 1], OddQ]
1779     = 31415
1780    #> ClearAll[A];
1781    """
1782
1783    def apply(self, items, expr, evaluation):
1784        "Select[items_, expr_]"
1785
1786        if items.is_atom():
1787            evaluation.message("Select", "normal")
1788            return
1789
1790        def cond(leaf):
1791            test = Expression(expr, leaf)
1792            return test.evaluate(evaluation).is_true()
1793
1794        return items.filter(items.head, cond, evaluation)
1795
1796
1797class Split(Builtin):
1798    """
1799    <dl>
1800    <dt>'Split[$list$]'
1801        <dd>splits $list$ into collections of consecutive identical elements.
1802    <dt>'Split[$list$, $test$]'
1803        <dd>splits $list$ based on whether the function $test$ yields
1804        'True' on consecutive elements.
1805    </dl>
1806
1807    >> Split[{x, x, x, y, x, y, y, z}]
1808     = {{x, x, x}, {y}, {x}, {y, y}, {z}}
1809
1810    #> Split[{x, x, x, y, x, y, y, z}, x]
1811     = {{x}, {x}, {x}, {y}, {x}, {y}, {y}, {z}}
1812
1813    Split into increasing or decreasing runs of elements
1814    >> Split[{1, 5, 6, 3, 6, 1, 6, 3, 4, 5, 4}, Less]
1815     = {{1, 5, 6}, {3, 6}, {1, 6}, {3, 4, 5}, {4}}
1816
1817    >> Split[{1, 5, 6, 3, 6, 1, 6, 3, 4, 5, 4}, Greater]
1818     = {{1}, {5}, {6, 3}, {6, 1}, {6, 3}, {4}, {5, 4}}
1819
1820    Split based on first element
1821    >> Split[{x -> a, x -> y, 2 -> a, z -> c, z -> a}, First[#1] === First[#2] &]
1822     = {{x -> a, x -> y}, {2 -> a}, {z -> c, z -> a}}
1823
1824    #> Split[{}]
1825     = {}
1826
1827    #> A[x__] := 321 /; Length[{x}] == 5;
1828    #> Split[A[x, x, x, y, x, y, y, z]]
1829     = 321
1830    #> ClearAll[A];
1831    """
1832
1833    rules = {
1834        "Split[list_]": "Split[list, SameQ]",
1835    }
1836
1837    messages = {
1838        "normal": "Nonatomic expression expected at position `1` in `2`.",
1839    }
1840
1841    def apply(self, mlist, test, evaluation):
1842        "Split[mlist_, test_]"
1843
1844        expr = Expression("Split", mlist, test)
1845
1846        if mlist.is_atom():
1847            evaluation.message("Select", "normal", 1, expr)
1848            return
1849
1850        if not mlist.leaves:
1851            return Expression(mlist.head)
1852
1853        result = [[mlist.leaves[0]]]
1854        for leaf in mlist.leaves[1:]:
1855            applytest = Expression(test, result[-1][-1], leaf)
1856            if applytest.evaluate(evaluation).is_true():
1857                result[-1].append(leaf)
1858            else:
1859                result.append([leaf])
1860
1861        inner = structure("List", mlist, evaluation)
1862        outer = structure(mlist.head, inner, evaluation)
1863        return outer([inner(l) for l in result])
1864
1865
1866class SplitBy(Builtin):
1867    """
1868    <dl>
1869    <dt>'SplitBy[$list$, $f$]'
1870        <dd>splits $list$ into collections of consecutive elements
1871        that give the same result when $f$ is applied.
1872    </dl>
1873
1874    >> SplitBy[Range[1, 3, 1/3], Round]
1875     = {{1, 4 / 3}, {5 / 3, 2, 7 / 3}, {8 / 3, 3}}
1876
1877    >> SplitBy[{1, 2, 1, 1.2}, {Round, Identity}]
1878     = {{{1}}, {{2}}, {{1}, {1.2}}}
1879
1880    #> SplitBy[Tuples[{1, 2}, 3], First]
1881     = {{{1, 1, 1}, {1, 1, 2}, {1, 2, 1}, {1, 2, 2}}, {{2, 1, 1}, {2, 1, 2}, {2, 2, 1}, {2, 2, 2}}}
1882    """
1883
1884    rules = {
1885        "SplitBy[list_]": "SplitBy[list, Identity]",
1886    }
1887
1888    messages = {
1889        "normal": "Nonatomic expression expected at position `1` in `2`.",
1890    }
1891
1892    def apply(self, mlist, func, evaluation):
1893        "SplitBy[mlist_, func_?NotListQ]"
1894
1895        expr = Expression("Split", mlist, func)
1896
1897        if mlist.is_atom():
1898            evaluation.message("Select", "normal", 1, expr)
1899            return
1900
1901        plist = [l for l in mlist.leaves]
1902
1903        result = [[plist[0]]]
1904        prev = Expression(func, plist[0]).evaluate(evaluation)
1905        for leaf in plist[1:]:
1906            curr = Expression(func, leaf).evaluate(evaluation)
1907            if curr == prev:
1908                result[-1].append(leaf)
1909            else:
1910                result.append([leaf])
1911            prev = curr
1912
1913        inner = structure("List", mlist, evaluation)
1914        outer = structure(mlist.head, inner, evaluation)
1915        return outer([inner(l) for l in result])
1916
1917    def apply_multiple(self, mlist, funcs, evaluation):
1918        "SplitBy[mlist_, funcs_?ListQ]"
1919        expr = Expression("Split", mlist, funcs)
1920
1921        if mlist.is_atom():
1922            evaluation.message("Select", "normal", 1, expr)
1923            return
1924
1925        result = mlist
1926        for f in funcs.leaves[::-1]:
1927            result = self.apply(result, f, evaluation)
1928
1929        return result
1930
1931
1932class Pick(Builtin):
1933    """
1934    <dl>
1935    <dt>'Pick[$list$, $sel$]'
1936        <dd>returns those items in $list$ that are True in $sel$.
1937    <dt>'Pick[$list$, $sel$, $patt$]'
1938        <dd>returns those items in $list$ that match $patt$ in $sel$.
1939    </dl>
1940
1941    >> Pick[{a, b, c}, {False, True, False}]
1942     = {b}
1943
1944    >> Pick[f[g[1, 2], h[3, 4]], {{True, False}, {False, True}}]
1945     = f[g[1], h[4]]
1946
1947    >> Pick[{a, b, c, d, e}, {1, 2, 3.5, 4, 5.5}, _Integer]
1948     = {a, b, d}
1949    """
1950
1951    def _do(self, items0, sel0, match, evaluation):
1952        def pick(items, sel):
1953            for x, s in zip(items, sel):
1954                if match(s):
1955                    yield x
1956                elif not x.is_atom() and not s.is_atom():
1957                    yield x.restructure(x.head, pick(x.leaves, s.leaves), evaluation)
1958
1959        r = list(pick([items0], [sel0]))
1960        if not r:
1961            return Expression(SymbolSequence)
1962        else:
1963            return r[0]
1964
1965    def apply(self, items, sel, evaluation):
1966        "Pick[items_, sel_]"
1967        return self._do(items, sel, lambda s: s.is_true(), evaluation)
1968
1969    def apply_pattern(self, items, sel, pattern, evaluation):
1970        "Pick[items_, sel_, pattern_]"
1971        from mathics.builtin.patterns import Matcher
1972
1973        match = Matcher(pattern).match
1974        return self._do(items, sel, lambda s: match(s, evaluation), evaluation)
1975
1976
1977class Cases(Builtin):
1978    r"""
1979    <dl>
1980      <dt>'Cases[$list$, $pattern$]'
1981      <dd>returns the elements of $list$ that match $pattern$.
1982
1983      <dt>'Cases[$list$, $pattern$, $ls$]'
1984      <dd>returns the elements matching at levelspec $ls$.
1985
1986      <dt>'Cases[$list$, $pattern$, Heads->$bool$]'
1987      <dd>Match including the head of the expression in the search.
1988    </dl>
1989
1990    >> Cases[{a, 1, 2.5, "string"}, _Integer|_Real]
1991     = {1, 2.5}
1992    >> Cases[_Complex][{1, 2I, 3, 4-I, 5}]
1993     = {2 I, 4 - I}
1994
1995    Find symbols among the elements of an expression:
1996    >> Cases[{b, 6, \[Pi]}, _Symbol]
1997     = {b, Pi}
1998
1999    Also include the head of the expression in the previous search:
2000    >> Cases[{b, 6, \[Pi]}, _Symbol, Heads -> True]
2001     = {List, b, Pi}
2002
2003    #> Cases[1, 2]
2004     = {}
2005
2006    #> Cases[f[1, 2], 2]
2007     = {2}
2008
2009    #> Cases[f[f[1, 2], f[2]], 2]
2010     = {}
2011    #> Cases[f[f[1, 2], f[2]], 2, 2]
2012     = {2, 2}
2013    #> Cases[f[f[1, 2], f[2], 2], 2, Infinity]
2014     = {2, 2, 2}
2015
2016    #> Cases[{1, f[2], f[3, 3, 3], 4, f[5, 5]}, f[x__] :> Plus[x]]
2017     = {2, 9, 10}
2018    #> Cases[{1, f[2], f[3, 3, 3], 4, f[5, 5]}, f[x__] -> Plus[x]]
2019     = {2, 3, 3, 3, 5, 5}
2020
2021    ## Issue 531
2022    #> z = f[x, y]; x = 1; Cases[z, _Symbol, Infinity]
2023     = {y}
2024    """
2025
2026    rules = {
2027        "Cases[pattern_][list_]": "Cases[list, pattern]",
2028    }
2029
2030    options = {
2031        "Heads": "False",
2032    }
2033
2034    def apply(self, items, pattern, ls, evaluation, options):
2035        "Cases[items_, pattern_, ls_:{1}, OptionsPattern[]]"
2036        if items.is_atom():
2037            return Expression(SymbolList)
2038
2039        from mathics.builtin.patterns import Matcher
2040        if ls.has_form("Rule", 2):
2041            if ls.leaves[0].get_name() == "System`Heads":
2042                heads = ls.leaves[1].is_true()
2043                ls = Expression("List", 1)
2044            else:
2045                return evaluation.message("Position", "level", ls)
2046        else:
2047            heads = self.get_option(options, "Heads", evaluation).is_true()
2048
2049        try:
2050            start, stop = python_levelspec(ls)
2051        except InvalidLevelspecError:
2052            return evaluation.message("Position", "level", ls)
2053
2054        results = []
2055
2056        if pattern.has_form("Rule", 2) or pattern.has_form("RuleDelayed", 2):
2057
2058            match = Matcher(pattern.leaves[0]).match
2059            rule = Rule(pattern.leaves[0], pattern.leaves[1])
2060
2061            def callback(level):
2062                if match(level, evaluation):
2063                    result = rule.apply(level, evaluation)
2064                    result = result.evaluate(evaluation)
2065                    results.append(result)
2066                return level
2067
2068        else:
2069            match = Matcher(pattern).match
2070
2071            def callback(level):
2072                if match(level, evaluation):
2073                    results.append(level)
2074                return level
2075
2076        walk_levels(items, start, stop, heads=heads, callback=callback)
2077
2078        return Expression(SymbolList, *results)
2079
2080
2081class DeleteCases(Builtin):
2082    """
2083    <dl>
2084    <dt>'DeleteCases[$list$, $pattern$]'
2085        <dd>returns the elements of $list$ that do not match $pattern$.
2086
2087    <dt>'DeleteCases[$list$, $pattern$, $levelspec$]'
2088        <dd> removes all parts of $list on levels specified by $levelspec$
2089             that match pattern (not fully implemented).
2090
2091    <dt>'DeleteCases[$list$, $pattern$, $levelspec$, $n$]'
2092        <dd> removes the first $n$ parts of $list$ that match $pattern$.
2093    </dl>
2094
2095    >> DeleteCases[{a, 1, 2.5, "string"}, _Integer|_Real]
2096     = {a, string}
2097
2098    >> DeleteCases[{a, b, 1, c, 2, 3}, _Symbol]
2099     = {1, 2, 3}
2100
2101    ## Issue 531
2102    #> z = {x, y}; x = 1; DeleteCases[z, _Symbol]
2103     = {1}
2104    """
2105
2106    messages = {
2107        "level": "Level specification `1` is not of the form n, {n}, or {m, n}.",
2108        "innf": "Non-negative integer or Infinity expected at position 4 in `1`",
2109    }
2110
2111    def apply_ls_n(self, items, pattern, levelspec, n, evaluation):
2112        "DeleteCases[items_, pattern_, levelspec_:1, n_:System`Infinity]"
2113
2114        if items.is_atom():
2115            evaluation.message("Select", "normal")
2116            return
2117        # If levelspec is specified to a non-trivial value,
2118        # we need to proceed with this complicate procedure
2119        # involving 1) decode what is the levelspec means
2120        # 2) find all the occurences
2121        # 3) Set all the occurences to ```System`Nothing```
2122
2123        levelspec = python_levelspec(levelspec)
2124
2125        if n == Symbol("Infinity"):
2126            n = -1
2127        elif n.get_head_name() == "System`Integer":
2128            n = n.get_int_value()
2129            if n < 0:
2130                evaluation.message(
2131                    "DeleteCases",
2132                    "innf",
2133                    Expression("DeleteCases", items, pattern, levelspec, n),
2134                )
2135        else:
2136            evaluation.message(
2137                "DeleteCases",
2138                "innf",
2139                Expression("DeleteCases", items, pattern, levelspec, n),
2140            )
2141            return SymbolNull
2142
2143        if levelspec[0] != 1 or levelspec[1] != 1:
2144            return deletecases_with_levelspec(items, pattern, evaluation, levelspec, n)
2145        # A more efficient way to proceed if levelspec == 1
2146        from mathics.builtin.patterns import Matcher
2147
2148        match = Matcher(pattern).match
2149        if n == -1:
2150
2151            def cond(leaf):
2152                return not match(leaf, evaluation)
2153
2154            return items.filter("List", cond, evaluation)
2155        else:
2156
2157            def condn(leaf):
2158                nonlocal n
2159                if n == 0:
2160                    return True
2161                elif match(leaf, evaluation):
2162                    n = n - 1
2163                    return False
2164                else:
2165                    return True
2166
2167            return items.filter("List", condn, evaluation)
2168
2169
2170class Count(Builtin):
2171    """
2172    <dl>
2173    <dt>'Count[$list$, $pattern$]'
2174        <dd>returns the number of times $pattern$ appears in $list$.
2175    <dt>'Count[$list$, $pattern$, $ls$]'
2176        <dd>counts the elements matching at levelspec $ls$.
2177    </dl>
2178
2179    >> Count[{3, 7, 10, 7, 5, 3, 7, 10}, 3]
2180     = 2
2181
2182    >> Count[{{a, a}, {a, a, a}, a}, a, {2}]
2183     = 5
2184    """
2185
2186    rules = {
2187        "Count[pattern_][list_]": "Count[list, pattern]",
2188        "Count[list_, arguments__]": "Length[Cases[list, arguments]]",
2189    }
2190
2191
2192class LeafCount(Builtin):
2193    """
2194    <dl>
2195    <dt>'LeafCount[$expr$]'
2196        <dd>returns the total number of indivisible subexpressions in $expr$.
2197    </dl>
2198
2199    >> LeafCount[1 + x + y^a]
2200     = 6
2201
2202    >> LeafCount[f[x, y]]
2203     = 3
2204
2205    >> LeafCount[{1 / 3, 1 + I}]
2206     = 7
2207
2208    >> LeafCount[Sqrt[2]]
2209     = 5
2210
2211    >> LeafCount[100!]
2212     = 1
2213
2214    #> LeafCount[f[a, b][x, y]]
2215     = 5
2216
2217    #> NestList[# /. s[x_][y_][z_] -> x[z][y[z]] &, s[s][s][s[s]][s][s], 4];
2218    #> LeafCount /@ %
2219     = {7, 8, 8, 11, 11}
2220
2221    #> LeafCount[1 / 3, 1 + I]
2222     : LeafCount called with 2 arguments; 1 argument is expected.
2223     = LeafCount[1 / 3, 1 + I]
2224    """
2225
2226    messages = {
2227        "argx": "LeafCount called with `1` arguments; 1 argument is expected.",
2228    }
2229
2230    def apply(self, expr, evaluation):
2231        "LeafCount[expr___]"
2232
2233        from mathics.core.expression import Rational, Complex
2234
2235        leaves = []
2236
2237        def callback(level):
2238            if isinstance(level, Rational):
2239                leaves.extend(
2240                    [level.get_head(), level.numerator(), level.denominator()]
2241                )
2242            elif isinstance(level, Complex):
2243                leaves.extend([level.get_head(), level.real, level.imag])
2244            else:
2245                leaves.append(level)
2246            return level
2247
2248        expr = expr.get_sequence()
2249        if len(expr) != 1:
2250            return evaluation.message("LeafCount", "argx", Integer(len(expr)))
2251
2252        walk_levels(expr[0], start=-1, stop=-1, heads=True, callback=callback)
2253        return Integer(len(leaves))
2254
2255
2256class Position(Builtin):
2257    """
2258    <dl>
2259    <dt>'Position[$expr$, $patt$]'
2260        <dd>returns the list of positions for which $expr$ matches $patt$.
2261    <dt>'Position[$expr$, $patt$, $ls$]'
2262        <dd>returns the positions on levels specified by levelspec $ls$.
2263    </dl>
2264
2265    >> Position[{1, 2, 2, 1, 2, 3, 2}, 2]
2266     = {{2}, {3}, {5}, {7}}
2267
2268    Find positions upto 3 levels deep
2269    >> Position[{1 + Sin[x], x, (Tan[x] - y)^2}, x, 3]
2270     = {{1, 2, 1}, {2}}
2271
2272    Find all powers of x
2273    >> Position[{1 + x^2, x y ^ 2,  4 y,  x ^ z}, x^_]
2274     = {{1, 2}, {4}}
2275
2276    Use Position as an operator
2277    >> Position[_Integer][{1.5, 2, 2.5}]
2278     = {{2}}
2279    """
2280
2281    options = {"Heads": "True"}
2282
2283    rules = {
2284        "Position[pattern_][expr_]": "Position[expr, pattern]",
2285    }
2286
2287    def apply_invalidlevel(self, patt, expr, ls, evaluation, options={}):
2288        "Position[expr_, patt_, ls_, OptionsPattern[Position]]"
2289
2290        return evaluation.message("Position", "level", ls)
2291
2292    def apply_level(self, expr, patt, ls, evaluation, options={}):
2293        """Position[expr_, patt_, Optional[Pattern[ls, _?LevelQ], {0, DirectedInfinity[1]}],
2294        OptionsPattern[Position]]"""
2295
2296        try:
2297            start, stop = python_levelspec(ls)
2298        except InvalidLevelspecError:
2299            return evaluation.message("Position", "level", ls)
2300
2301        from mathics.builtin.patterns import Matcher
2302
2303        match = Matcher(patt).match
2304        result = []
2305
2306        def callback(level, pos):
2307            if match(level, evaluation):
2308                result.append(pos)
2309            return level
2310
2311        heads = self.get_option(options, "Heads", evaluation).is_true()
2312        walk_levels(expr, start, stop, heads=heads, callback=callback, include_pos=True)
2313        return from_python(result)
2314
2315
2316class MemberQ(Builtin):
2317    """
2318    <dl>
2319    <dt>'MemberQ[$list$, $pattern$]'
2320        <dd>returns 'True' if $pattern$ matches any element of $list$,
2321        or 'False' otherwise.
2322    </dl>
2323
2324    >> MemberQ[{a, b, c}, b]
2325     = True
2326    >> MemberQ[{a, b, c}, d]
2327     = False
2328    >> MemberQ[{"a", b, f[x]}, _?NumericQ]
2329     = False
2330    >> MemberQ[_List][{{}}]
2331     = True
2332    """
2333
2334    rules = {
2335        "MemberQ[list_, pattern_]": ("Length[Select[list, MatchQ[#, pattern]&]] > 0"),
2336        "MemberQ[pattern_][expr_]": "MemberQ[expr, pattern]",
2337    }
2338
2339
2340class Range(Builtin):
2341    """
2342    <dl>
2343    <dt>'Range[$n$]'
2344        <dd>returns a list of integers from 1 to $n$.
2345    <dt>'Range[$a$, $b$]'
2346        <dd>returns a list of integers from $a$ to $b$.
2347    </dl>
2348    >> Range[5]
2349     = {1, 2, 3, 4, 5}
2350    >> Range[-3, 2]
2351     = {-3, -2, -1, 0, 1, 2}
2352    >> Range[0, 2, 1/3]
2353     = {0, 1 / 3, 2 / 3, 1, 4 / 3, 5 / 3, 2}
2354    """
2355
2356    rules = {
2357        "Range[imax_?RealNumberQ]": "Range[1, imax, 1]",
2358        "Range[imin_?RealNumberQ, imax_?RealNumberQ]": "Range[imin, imax, 1]",
2359    }
2360
2361    attributes = ("Listable", "Protected")
2362
2363    def apply(self, imin, imax, di, evaluation):
2364        "Range[imin_?RealNumberQ, imax_?RealNumberQ, di_?RealNumberQ]"
2365
2366        imin = imin.to_sympy()
2367        imax = imax.to_sympy()
2368        di = di.to_sympy()
2369        index = imin
2370        result = []
2371        while index <= imax:
2372            evaluation.check_stopped()
2373            result.append(from_sympy(index))
2374            index += di
2375        return Expression(SymbolList, *result)
2376
2377
2378class _IterationFunction(Builtin):
2379    """
2380    >> Sum[k, {k, Range[5]}]
2381     = 15
2382    """
2383
2384    attributes = ("HoldAll",)
2385    allow_loopcontrol = False
2386    throw_iterb = True
2387
2388    def get_result(self, items):
2389        pass
2390
2391    def apply_symbol(self, expr, iterator, evaluation):
2392        "%(name)s[expr_, iterator_Symbol]"
2393        iterator = iterator.evaluate(evaluation)
2394        if iterator.has_form(["List", "Range", "Sequence"], None):
2395            leaves = iterator.leaves
2396            if len(leaves) == 1:
2397                return self.apply_max(expr, *leaves, evaluation)
2398            elif len(leaves) == 2:
2399                if leaves[1].has_form(["List", "Sequence"], None):
2400                    seq = Expression(SymbolSequence, *(leaves[1].leaves))
2401                    return self.apply_list(expr, leaves[0], seq, evaluation)
2402                else:
2403                    return self.apply_range(expr, *leaves, evaluation)
2404            elif len(leaves) == 3:
2405                return self.apply_iter_nostep(expr, *leaves, evaluation)
2406            elif len(leaves) == 4:
2407                return self.apply_iter(expr, *leaves, evaluation)
2408
2409        if self.throw_iterb:
2410            evaluation.message(self.get_name(), "iterb")
2411        return
2412
2413    def apply_range(self, expr, i, imax, evaluation):
2414        "%(name)s[expr_, {i_Symbol, imax_}]"
2415        imax = imax.evaluate(evaluation)
2416        if imax.has_form("Range", None):
2417            # Fixme: this should work as an iterator in python3, not
2418            # building the sequence explicitly...
2419            seq = Expression(SymbolSequence, *(imax.evaluate(evaluation).leaves))
2420            return self.apply_list(expr, i, seq, evaluation)
2421        elif imax.has_form("List", None):
2422            seq = Expression(SymbolSequence, *(imax.leaves))
2423            return self.apply_list(expr, i, seq, evaluation)
2424        else:
2425            return self.apply_iter(expr, i, Integer(1), imax, Integer(1), evaluation)
2426
2427    def apply_max(self, expr, imax, evaluation):
2428        "%(name)s[expr_, {imax_}]"
2429
2430        index = 0
2431        imax = imax.evaluate(evaluation)
2432        imax = imax.numerify(evaluation)
2433        if isinstance(imax, Number):
2434            imax = imax.round()
2435        imax = imax.get_float_value()
2436        if imax is None:
2437            if self.throw_iterb:
2438                evaluation.message(self.get_name(), "iterb")
2439            return
2440        result = []
2441        while index < imax:
2442            evaluation.check_stopped()
2443            try:
2444                result.append(expr.evaluate(evaluation))
2445            except ContinueInterrupt:
2446                if self.allow_loopcontrol:
2447                    pass
2448                else:
2449                    raise
2450            except BreakInterrupt:
2451                if self.allow_loopcontrol:
2452                    break
2453                else:
2454                    raise
2455            except ReturnInterrupt as e:
2456                if self.allow_loopcontrol:
2457                    return e.expr
2458                else:
2459                    raise
2460            index += 1
2461        return self.get_result(result)
2462
2463    def apply_iter_nostep(self, expr, i, imin, imax, evaluation):
2464        "%(name)s[expr_, {i_Symbol, imin_, imax_}]"
2465        return self.apply_iter(expr, i, imin, imax, Integer(1), evaluation)
2466
2467    def apply_iter(self, expr, i, imin, imax, di, evaluation):
2468        "%(name)s[expr_, {i_Symbol, imin_, imax_, di_}]"
2469
2470        if isinstance(self, SympyFunction) and di.get_int_value() == 1:
2471            whole_expr = Expression(
2472                self.get_name(), expr, Expression(SymbolList, i, imin, imax)
2473            )
2474            sympy_expr = whole_expr.to_sympy(evaluation=evaluation)
2475            if sympy_expr is None:
2476                return None
2477
2478            # apply Together to produce results similar to Mathematica
2479            result = sympy.together(sympy_expr)
2480            result = from_sympy(result)
2481            result = cancel(result)
2482
2483            if not result.sameQ(whole_expr):
2484                return result
2485            return
2486
2487        index = imin.evaluate(evaluation)
2488        imax = imax.evaluate(evaluation)
2489        di = di.evaluate(evaluation)
2490
2491        result = []
2492        compare_type = (
2493            "GreaterEqual"
2494            if Expression("Less", di, Integer0).evaluate(evaluation).to_python()
2495            else "LessEqual"
2496        )
2497        while True:
2498            cont = Expression(compare_type, index, imax).evaluate(evaluation)
2499            if cont == Symbol("False"):
2500                break
2501            if not cont.is_true():
2502                if self.throw_iterb:
2503                    evaluation.message(self.get_name(), "iterb")
2504                return
2505
2506            evaluation.check_stopped()
2507            try:
2508                item = dynamic_scoping(expr.evaluate, {i.name: index}, evaluation)
2509                result.append(item)
2510            except ContinueInterrupt:
2511                if self.allow_loopcontrol:
2512                    pass
2513                else:
2514                    raise
2515            except BreakInterrupt:
2516                if self.allow_loopcontrol:
2517                    break
2518                else:
2519                    raise
2520            except ReturnInterrupt as e:
2521                if self.allow_loopcontrol:
2522                    return e.expr
2523                else:
2524                    raise
2525            index = Expression("Plus", index, di).evaluate(evaluation)
2526        return self.get_result(result)
2527
2528    def apply_list(self, expr, i, items, evaluation):
2529        "%(name)s[expr_, {i_Symbol, {items___}}]"
2530        items = items.evaluate(evaluation).get_sequence()
2531        result = []
2532        for item in items:
2533            evaluation.check_stopped()
2534            try:
2535                item = dynamic_scoping(expr.evaluate, {i.name: item}, evaluation)
2536                result.append(item)
2537            except ContinueInterrupt:
2538                if self.allow_loopcontrol:
2539                    pass
2540                else:
2541                    raise
2542            except BreakInterrupt:
2543                if self.allow_loopcontrol:
2544                    break
2545                else:
2546                    raise
2547            except ReturnInterrupt as e:
2548                if self.allow_loopcontrol:
2549                    return e.expr
2550                else:
2551                    raise
2552        return self.get_result(result)
2553
2554    def apply_multi(self, expr, first, sequ, evaluation):
2555        "%(name)s[expr_, first_, sequ__]"
2556
2557        sequ = sequ.get_sequence()
2558        name = self.get_name()
2559        return Expression(name, Expression(name, expr, *sequ), first)
2560
2561
2562class ConstantArray(Builtin):
2563    """
2564    <dl>
2565    <dt>'ConstantArray[$expr$, $n$]'
2566        <dd>returns a list of $n$ copies of $expr$.
2567    </dl>
2568
2569    >> ConstantArray[a, 3]
2570     = {a, a, a}
2571    >> ConstantArray[a, {2, 3}]
2572     = {{a, a, a}, {a, a, a}}
2573    """
2574
2575    rules = {
2576        "ConstantArray[c_, dims_]": "Apply[Table[c, ##]&, List /@ dims]",
2577        "ConstantArray[c_, n_Integer]": "ConstantArray[c, {n}]",
2578    }
2579
2580
2581class Array(Builtin):
2582    """
2583    <dl>
2584    <dt>'Array[$f$, $n$]'
2585        <dd>returns the $n$-element list '{$f$[1], ..., $f$[$n$]}'.
2586    <dt>'Array[$f$, $n$, $a$]'
2587        <dd>returns the $n$-element list '{$f$[$a$], ..., $f$[$a$ + $n$]}'.
2588    <dt>'Array[$f$, {$n$, $m$}, {$a$, $b$}]'
2589        <dd>returns an $n$-by-$m$ matrix created by applying $f$ to
2590        indices ranging from '($a$, $b$)' to '($a$ + $n$, $b$ + $m$)'.
2591    <dt>'Array[$f$, $dims$, $origins$, $h$]'
2592        <dd>returns an expression with the specified dimensions and
2593        index origins, with head $h$ (instead of 'List').
2594    </dl>
2595
2596    >> Array[f, 4]
2597     = {f[1], f[2], f[3], f[4]}
2598    >> Array[f, {2, 3}]
2599     = {{f[1, 1], f[1, 2], f[1, 3]}, {f[2, 1], f[2, 2], f[2, 3]}}
2600    >> Array[f, {2, 3}, 3]
2601     = {{f[3, 3], f[3, 4], f[3, 5]}, {f[4, 3], f[4, 4], f[4, 5]}}
2602    >> Array[f, {2, 3}, {4, 6}]
2603     = {{f[4, 6], f[4, 7], f[4, 8]}, {f[5, 6], f[5, 7], f[5, 8]}}
2604    >> Array[f, {2, 3}, 1, Plus]
2605     = f[1, 1] + f[1, 2] + f[1, 3] + f[2, 1] + f[2, 2] + f[2, 3]
2606
2607    #> Array[f, {2, 3}, {1, 2, 3}]
2608     : {2, 3} and {1, 2, 3} should have the same length.
2609     = Array[f, {2, 3}, {1, 2, 3}]
2610    #> Array[f, a]
2611     : Single or list of non-negative integers expected at position 2.
2612     = Array[f, a]
2613    #> Array[f, 2, b]
2614     : Single or list of non-negative integers expected at position 3.
2615     = Array[f, 2, b]
2616    """
2617
2618    messages = {
2619        "plen": "`1` and `2` should have the same length.",
2620    }
2621
2622    def apply(self, f, dimsexpr, origins, head, evaluation):
2623        "Array[f_, dimsexpr_, origins_:1, head_:List]"
2624
2625        if dimsexpr.has_form("List", None):
2626            dims = dimsexpr.get_mutable_leaves()
2627        else:
2628            dims = [dimsexpr]
2629        for index, dim in enumerate(dims):
2630            value = dim.get_int_value()
2631            if value is None:
2632                evaluation.message("Array", "ilsnn", 2)
2633                return
2634            dims[index] = value
2635        if origins.has_form("List", None):
2636            if len(origins.leaves) != len(dims):
2637                evaluation.message("Array", "plen", dimsexpr, origins)
2638                return
2639            origins = origins.get_mutable_leaves()
2640        else:
2641            origins = [origins] * len(dims)
2642        for index, origin in enumerate(origins):
2643            value = origin.get_int_value()
2644            if value is None:
2645                evaluation.message("Array", "ilsnn", 3)
2646                return
2647            origins[index] = value
2648
2649        dims = list(zip(dims, origins))
2650
2651        def rec(rest_dims, current):
2652            evaluation.check_stopped()
2653            if rest_dims:
2654                level = []
2655                count, origin = rest_dims[0]
2656                for index in range(origin, origin + count):
2657                    level.append(rec(rest_dims[1:], current + [index]))
2658                return Expression(head, *level)
2659            else:
2660                return Expression(f, *(Integer(index) for index in current))
2661
2662        return rec(dims, [])
2663
2664
2665class Table(_IterationFunction):
2666    """
2667    <dl>
2668      <dt>'Table[$expr$, $n$]'
2669      <dd>generates a list of $n$ copies of $expr$.
2670
2671      <dt>'Table[$expr$, {$i$, $n$}]'
2672      <dd>generates a list of the values of expr when $i$ runs from 1 to $n$.
2673
2674      <dt>'Table[$expr$, {$i$, $start$, $stop$, $step$}]'
2675      <dd>evaluates $expr$ with $i$ ranging from $start$ to $stop$,
2676        incrementing by $step$.
2677
2678      <dt>'Table[$expr$, {$i$, {$e1$, $e2$, ..., $ei$}}]'
2679      <dd>evaluates $expr$ with $i$ taking on the values $e1$, $e2$,
2680        ..., $ei$.
2681    </dl>
2682    >> Table[x, 3]
2683     = {x, x, x}
2684    >> n = 0; Table[n = n + 1, {5}]
2685     = {1, 2, 3, 4, 5}
2686    >> Table[i, {i, 4}]
2687     = {1, 2, 3, 4}
2688    >> Table[i, {i, 2, 5}]
2689     = {2, 3, 4, 5}
2690    >> Table[i, {i, 2, 6, 2}]
2691     = {2, 4, 6}
2692    >> Table[i, {i, Pi, 2 Pi, Pi / 2}]
2693     = {Pi, 3 Pi / 2, 2 Pi}
2694    >> Table[x^2, {x, {a, b, c}}]
2695     = {a ^ 2, b ^ 2, c ^ 2}
2696
2697    'Table' supports multi-dimensional tables:
2698    >> Table[{i, j}, {i, {a, b}}, {j, 1, 2}]
2699     = {{{a, 1}, {a, 2}}, {{b, 1}, {b, 2}}}
2700
2701    #> Table[x, {x,0,1/3}]
2702     = {0}
2703    #> Table[x, {x, -0.2, 3.9}]
2704     = {-0.2, 0.8, 1.8, 2.8, 3.8}
2705    """
2706
2707    rules = {
2708        "Table[expr_, n_Integer]": "Table[expr, {n}]",
2709    }
2710
2711    def get_result(self, items):
2712        return Expression(SymbolList, *items)
2713
2714
2715class Join(Builtin):
2716    """
2717    <dl>
2718    <dt>'Join[$l1$, $l2$]'
2719        <dd>concatenates the lists $l1$ and $l2$.
2720    </dl>
2721
2722    'Join' concatenates lists:
2723    >> Join[{a, b}, {c, d, e}]
2724     = {a, b, c, d, e}
2725    >> Join[{{a, b}, {c, d}}, {{1, 2}, {3, 4}}]
2726     = {{a, b}, {c, d}, {1, 2}, {3, 4}}
2727
2728    The concatenated expressions may have any head:
2729    >> Join[a + b, c + d, e + f]
2730     = a + b + c + d + e + f
2731
2732    However, it must be the same for all expressions:
2733    >> Join[a + b, c * d]
2734     : Heads Plus and Times are expected to be the same.
2735     = Join[a + b, c d]
2736
2737    #> Join[x, y]
2738     = Join[x, y]
2739    #> Join[x + y, z]
2740     = Join[x + y, z]
2741    #> Join[x + y, y z, a]
2742     : Heads Plus and Times are expected to be the same.
2743     = Join[x + y, y z, a]
2744    #> Join[x, y + z, y z]
2745     = Join[x, y + z, y z]
2746    """
2747
2748    attributes = ("Flat", "OneIdentity")
2749
2750    def apply(self, lists, evaluation):
2751        "Join[lists___]"
2752
2753        result = []
2754        head = None
2755        sequence = lists.get_sequence()
2756
2757        for list in sequence:
2758            if list.is_atom():
2759                return
2760            if head is not None and list.get_head() != head:
2761                evaluation.message("Join", "heads", head, list.get_head())
2762                return
2763            head = list.get_head()
2764            result.extend(list.leaves)
2765
2766        if result:
2767            return sequence[0].restructure(head, result, evaluation, deps=sequence)
2768        else:
2769            return Expression(SymbolList)
2770
2771
2772class Catenate(Builtin):
2773    """
2774    <dl>
2775    <dt>'Catenate[{$l1$, $l2$, ...}]'
2776        <dd>concatenates the lists $l1$, $l2$, ...
2777    </dl>
2778
2779    >> Catenate[{{1, 2, 3}, {4, 5}}]
2780     = {1, 2, 3, 4, 5}
2781    """
2782
2783    messages = {"invrp": "`1` is not a list."}
2784
2785    def apply(self, lists, evaluation):
2786        "Catenate[lists_List]"
2787
2788        def parts():
2789            for l in lists.leaves:
2790                head_name = l.get_head_name()
2791                if head_name == "System`List":
2792                    yield l.leaves
2793                elif head_name != "System`Missing":
2794                    raise MessageException("Catenate", "invrp", l)
2795
2796        try:
2797            result = list(chain(*list(parts())))
2798            if result:
2799                return lists.leaves[0].restructure(
2800                    "List", result, evaluation, deps=lists.leaves
2801                )
2802            else:
2803                return Expression(SymbolList)
2804        except MessageException as e:
2805            e.message(evaluation)
2806
2807
2808class Insert(Builtin):
2809    """
2810    <dl>
2811      <dt>'Insert[$list$, $elem$, $n$]'
2812      <dd>inserts $elem$ at position $n$ in $list$. When $n$ is negative, the position is counted from the end.
2813    </dl>
2814
2815    >> Insert[{a,b,c,d,e}, x, 3]
2816     = {a, b, x, c, d, e}
2817
2818    >> Insert[{a,b,c,d,e}, x, -2]
2819     = {a, b, c, d, x, e}
2820    """
2821
2822    def apply(self, expr, elem, n, evaluation):
2823        "Insert[expr_List, elem_, n_Integer]"
2824
2825        py_n = n.to_python()
2826        new_list = list(expr.get_leaves())
2827
2828        position = py_n - 1 if py_n > 0 else py_n + 1
2829        new_list.insert(position, elem)
2830        return expr.restructure(expr.head, new_list, evaluation, deps=(expr, elem))
2831
2832
2833class Append(Builtin):
2834    """
2835    <dl>
2836      <dt>'Append[$expr$, $elem$]'
2837      <dd>returns $expr$ with $elem$ appended.
2838    </dl>
2839
2840    >> Append[{1, 2, 3}, 4]
2841     = {1, 2, 3, 4}
2842
2843    'Append' works on expressions with heads other than 'List':
2844    >> Append[f[a, b], c]
2845     = f[a, b, c]
2846
2847    Unlike 'Join', 'Append' does not flatten lists in $item$:
2848    >> Append[{a, b}, {c, d}]
2849     = {a, b, {c, d}}
2850
2851    #> Append[a, b]
2852     : Nonatomic expression expected.
2853     = Append[a, b]
2854    """
2855
2856    def apply(self, expr, item, evaluation):
2857        "Append[expr_, item_]"
2858
2859        if expr.is_atom():
2860            return evaluation.message("Append", "normal")
2861
2862        return expr.restructure(
2863            expr.head,
2864            list(chain(expr.get_leaves(), [item])),
2865            evaluation,
2866            deps=(expr, item),
2867        )
2868
2869
2870class AppendTo(Builtin):
2871    """
2872    <dl>
2873    <dt>'AppendTo[$s$, $item$]'
2874        <dd>append $item$ to value of $s$ and sets $s$ to the result.
2875    </dl>
2876
2877    >> s = {};
2878    >> AppendTo[s, 1]
2879     = {1}
2880    >> s
2881     = {1}
2882
2883    'Append' works on expressions with heads other than 'List':
2884    >> y = f[];
2885    >> AppendTo[y, x]
2886     = f[x]
2887    >> y
2888     = f[x]
2889
2890    #> AppendTo[{}, 1]
2891     : {} is not a variable with a value, so its value cannot be changed.
2892     = AppendTo[{}, 1]
2893
2894    #> AppendTo[a, b]
2895     : a is not a variable with a value, so its value cannot be changed.
2896     = AppendTo[a, b]
2897    """
2898
2899    attributes = ("HoldFirst",)
2900
2901    messages = {
2902        "rvalue": "`1` is not a variable with a value, so its value cannot be changed.",
2903    }
2904
2905    def apply(self, s, item, evaluation):
2906        "AppendTo[s_, item_]"
2907        resolved_s = s.evaluate(evaluation)
2908        if s == resolved_s:
2909            return evaluation.message("AppendTo", "rvalue", s)
2910
2911        if not resolved_s.is_atom():
2912            result = Expression("Set", s, Expression("Append", resolved_s, item))
2913            return result.evaluate(evaluation)
2914
2915        return evaluation.message("AppendTo", "normal", Expression("AppendTo", s, item))
2916
2917
2918class Prepend(Builtin):
2919    """
2920    <dl>
2921     <dt>'Prepend[$expr$, $item$]'
2922     <dd>returns $expr$ with $item$ prepended to its leaves.
2923
2924     <dt>'Prepend[$expr$]'
2925     <dd>'Prepend[$elem$][$expr$]' is equivalent to 'Prepend[$expr$,$elem$]'.
2926    </dl>
2927
2928    'Prepend' is similar to 'Append', but adds $item$ to the beginning
2929    of $expr$:
2930    >> Prepend[{2, 3, 4}, 1]
2931     = {1, 2, 3, 4}
2932
2933    'Prepend' works on expressions with heads other than 'List':
2934    >> Prepend[f[b, c], a]
2935     = f[a, b, c]
2936
2937    Unlike 'Join', 'Prepend' does not flatten lists in $item$:
2938    >> Prepend[{c, d}, {a, b}]
2939     = {{a, b}, c, d}
2940
2941    #> Prepend[a, b]
2942     : Nonatomic expression expected.
2943     = Prepend[a, b]
2944    """
2945
2946    def apply(self, expr, item, evaluation):
2947        "Prepend[expr_, item_]"
2948
2949        if expr.is_atom():
2950            return evaluation.message("Prepend", "normal")
2951
2952        return expr.restructure(
2953            expr.head,
2954            list(chain([item], expr.get_leaves())),
2955            evaluation,
2956            deps=(expr, item),
2957        )
2958
2959
2960class PrependTo(Builtin):
2961    """
2962    <dl>
2963    <dt>'PrependTo[$s$, $item$]'
2964        <dd>prepends $item$ to value of $s$ and sets $s$ to the result.
2965    </dl>
2966
2967    Assign s to a list
2968    >> s = {1, 2, 4, 9}
2969     = {1, 2, 4, 9}
2970
2971    Add a new value at the beginning of the list:
2972    >> PrependTo[s, 0]
2973     = {0, 1, 2, 4, 9}
2974
2975    The value assigned to s has changed:
2976    >> s
2977     = {0, 1, 2, 4, 9}
2978
2979    'PrependTo' works with a head other than 'List':
2980    >> y = f[a, b, c];
2981    >> PrependTo[y, x]
2982     = f[x, a, b, c]
2983    >> y
2984     = f[x, a, b, c]
2985
2986    #> PrependTo[{a, b}, 1]
2987     :  {a, b} is not a variable with a value, so its value cannot be changed.
2988     = PrependTo[{a, b}, 1]
2989
2990    #> PrependTo[a, b]
2991     : a is not a variable with a value, so its value cannot be changed.
2992     = PrependTo[a, b]
2993
2994    #> x = 1 + 2;
2995    #> PrependTo[x, {3, 4}]
2996     : Nonatomic expression expected at position 1 in PrependTo[x, {3, 4}].
2997     =  PrependTo[x, {3, 4}]
2998    """
2999
3000    attributes = ("HoldFirst",)
3001
3002    messages = {
3003        "rvalue": "`1` is not a variable with a value, so its value cannot be changed.",
3004        "normal": "Nonatomic expression expected at position 1 in `1`.",
3005    }
3006
3007    def apply(self, s, item, evaluation):
3008        "PrependTo[s_, item_]"
3009        resolved_s = s.evaluate(evaluation)
3010        if s == resolved_s:
3011            return evaluation.message("PrependTo", "rvalue", s)
3012
3013        if not resolved_s.is_atom():
3014            result = Expression("Set", s, Expression("Prepend", resolved_s, item))
3015            return result.evaluate(evaluation)
3016
3017        return evaluation.message(
3018            "PrependTo", "normal", Expression("PrependTo", s, item)
3019        )
3020
3021
3022def get_tuples(items):
3023    if not items:
3024        yield []
3025    else:
3026        for item in items[0]:
3027            for rest in get_tuples(items[1:]):
3028                yield [item] + rest
3029
3030
3031class Tuples(Builtin):
3032    """
3033    <dl>
3034    <dt>'Tuples[$list$, $n$]'
3035        <dd>returns a list of all $n$-tuples of elements in $list$.
3036    <dt>'Tuples[{$list1$, $list2$, ...}]'
3037        <dd>returns a list of tuples with elements from the given lists.
3038    </dl>
3039
3040    >> Tuples[{a, b, c}, 2]
3041     = {{a, a}, {a, b}, {a, c}, {b, a}, {b, b}, {b, c}, {c, a}, {c, b}, {c, c}}
3042    >> Tuples[{}, 2]
3043     = {}
3044    >> Tuples[{a, b, c}, 0]
3045     = {{}}
3046
3047    >> Tuples[{{a, b}, {1, 2, 3}}]
3048     = {{a, 1}, {a, 2}, {a, 3}, {b, 1}, {b, 2}, {b, 3}}
3049
3050    The head of $list$ need not be 'List':
3051    >> Tuples[f[a, b, c], 2]
3052     = {f[a, a], f[a, b], f[a, c], f[b, a], f[b, b], f[b, c], f[c, a], f[c, b], f[c, c]}
3053    However, when specifying multiple expressions, 'List' is always used:
3054    >> Tuples[{f[a, b], g[c, d]}]
3055     = {{a, c}, {a, d}, {b, c}, {b, d}}
3056    """
3057
3058    def apply_n(self, expr, n, evaluation):
3059        "Tuples[expr_, n_Integer]"
3060
3061        if expr.is_atom():
3062            evaluation.message("Tuples", "normal")
3063            return
3064        n = n.get_int_value()
3065        if n is None or n < 0:
3066            evaluation.message("Tuples", "intnn")
3067            return
3068        items = expr.leaves
3069
3070        def iterate(n_rest):
3071            evaluation.check_stopped()
3072            if n_rest <= 0:
3073                yield []
3074            else:
3075                for item in items:
3076                    for rest in iterate(n_rest - 1):
3077                        yield [item] + rest
3078
3079        return Expression(
3080            "List", *(Expression(expr.head, *leaves) for leaves in iterate(n))
3081        )
3082
3083    def apply_lists(self, exprs, evaluation):
3084        "Tuples[{exprs___}]"
3085
3086        exprs = exprs.get_sequence()
3087        items = []
3088        for expr in exprs:
3089            evaluation.check_stopped()
3090            if expr.is_atom():
3091                evaluation.message("Tuples", "normal")
3092                return
3093            items.append(expr.leaves)
3094
3095        return Expression(
3096            "List", *(Expression(SymbolList, *leaves) for leaves in get_tuples(items))
3097        )
3098
3099
3100class Reap(Builtin):
3101    """
3102    <dl>
3103    <dt>'Reap[$expr$]'
3104        <dd>gives the result of evaluating $expr$, together with all
3105        values sown during this evaluation. Values sown with different
3106        tags are given in different lists.
3107    <dt>'Reap[$expr$, $pattern$]'
3108        <dd>only yields values sown with a tag matching $pattern$.
3109        'Reap[$expr$]' is equivalent to 'Reap[$expr$, _]'.
3110    <dt>'Reap[$expr$, {$pattern1$, $pattern2$, ...}]'
3111        <dd>uses multiple patterns.
3112    <dt>'Reap[$expr$, $pattern$, $f$]'
3113        <dd>applies $f$ on each tag and the corresponding values sown
3114        in the form '$f$[tag, {e1, e2, ...}]'.
3115    </dl>
3116
3117    >> Reap[Sow[3]; Sow[1]]
3118     = {1, {{3, 1}}}
3119
3120    >> Reap[Sow[2, {x, x, x}]; Sow[3, x]; Sow[4, y]; Sow[4, 1], {_Symbol, _Integer, x}, f]
3121     = {4, {{f[x, {2, 2, 2, 3}], f[y, {4}]}, {f[1, {4}]}, {f[x, {2, 2, 2, 3}]}}}
3122
3123    Find the unique elements of a list, keeping their order:
3124    >> Reap[Sow[Null, {a, a, b, d, c, a}], _, # &][[2]]
3125     = {a, b, d, c}
3126
3127    Sown values are reaped by the innermost matching 'Reap':
3128    >> Reap[Reap[Sow[a, x]; Sow[b, 1], _Symbol, Print["Inner: ", #1]&];, _, f]
3129     | Inner: x
3130     = {Null, {f[1, {b}]}}
3131
3132    When no value is sown, an empty list is returned:
3133    >> Reap[x]
3134     = {x, {}}
3135    """
3136
3137    attributes = ("HoldFirst",)
3138
3139    rules = {
3140        "Reap[expr_, pattern_, f_]": (
3141            "{#[[1]], #[[2, 1]]}& [Reap[expr, {pattern}, f]]"
3142        ),
3143        "Reap[expr_, pattern_]": "Reap[expr, pattern, #2&]",
3144        "Reap[expr_]": "Reap[expr, _]",
3145    }
3146
3147    def apply(self, expr, patterns, f, evaluation):
3148        "Reap[expr_, {patterns___}, f_]"
3149
3150        patterns = patterns.get_sequence()
3151        sown = [(Pattern.create(pattern), []) for pattern in patterns]
3152
3153        def listener(e, tag):
3154            result = False
3155            for pattern, items in sown:
3156                if pattern.does_match(tag, evaluation):
3157                    for item in items:
3158                        if item[0].sameQ(tag):
3159                            item[1].append(e)
3160                            break
3161                    else:
3162                        items.append((tag, [e]))
3163                    result = True
3164            return result
3165
3166        evaluation.add_listener("sow", listener)
3167        try:
3168            result = expr.evaluate(evaluation)
3169            items = []
3170            for pattern, tags in sown:
3171                leaves = []
3172                for tag, elements in tags:
3173                    leaves.append(Expression(f, tag, Expression(SymbolList, *elements)))
3174                items.append(Expression(SymbolList, *leaves))
3175            return Expression(SymbolList, result, Expression(SymbolList, *items))
3176        finally:
3177            evaluation.remove_listener("sow", listener)
3178
3179
3180class Sow(Builtin):
3181    """
3182    <dl>
3183    <dt>'Sow[$e$]'
3184        <dd>sends the value $e$ to the innermost 'Reap'.
3185    <dt>'Sow[$e$, $tag$]'
3186        <dd>sows $e$ using $tag$. 'Sow[$e$]' is equivalent to 'Sow[$e$, Null]'.
3187    <dt>'Sow[$e$, {$tag1$, $tag2$, ...}]'
3188        <dd>uses multiple tags.
3189    </dl>
3190    """
3191
3192    rules = {
3193        "Sow[e_]": "Sow[e, {Null}]",
3194        "Sow[e_, tag_]": "Sow[e, {tag}]",
3195    }
3196
3197    def apply(self, e, tags, evaluation):
3198        "Sow[e_, {tags___}]"
3199
3200        tags = tags.get_sequence()
3201        for tag in tags:
3202            evaluation.publish("sow", e, tag)
3203        return e
3204
3205
3206class UnitVector(Builtin):
3207    """
3208    <dl>
3209    <dt>'UnitVector[$n$, $k$]'
3210        <dd>returns the $n$-dimensional unit vector with a 1 in position $k$.
3211    <dt>'UnitVector[$k$]'
3212        <dd>is equivalent to 'UnitVector[2, $k$]'.
3213    </dl>
3214    >> UnitVector[2]
3215     = {0, 1}
3216    >> UnitVector[4, 3]
3217     = {0, 0, 1, 0}
3218    """
3219
3220    messages = {
3221        "nokun": "There is no unit vector in direction `1` in `2` dimensions.",
3222    }
3223
3224    rules = {
3225        "UnitVector[k_Integer]": "UnitVector[2, k]",
3226    }
3227
3228    def apply(self, n, k, evaluation):
3229        "UnitVector[n_Integer, k_Integer]"
3230
3231        n = n.get_int_value()
3232        k = k.get_int_value()
3233        if n is None or k is None:
3234            return
3235        if not 1 <= k <= n:
3236            evaluation.message("UnitVector", "nokun", k, n)
3237            return
3238
3239        def item(i):
3240            if i == k:
3241                return Integer(1)
3242            else:
3243                return Integer0
3244
3245        return Expression(SymbolList, *(item(i) for i in range(1, n + 1)))
3246
3247
3248def riffle(items, sep):
3249    result = items[:1]
3250    for item in items[1:]:
3251        result.append(sep)
3252        result.append(item)
3253    return result
3254
3255
3256def riffle_lists(items, seps):
3257    if len(seps) == 0:  # special case
3258        seps = [Expression(SymbolList)]
3259
3260    i = 0
3261    while i < len(items):
3262        yield items[i]
3263        if i == len(items) - 1 and len(items) != len(seps):
3264            return
3265        yield seps[i % len(seps)]
3266        i += 1
3267
3268
3269class Riffle(Builtin):
3270    """
3271    <dl>
3272    <dt>'Riffle[$list$, $x$]'
3273        <dd>inserts a copy of $x$ between each element of $list$.
3274    <dt>'Riffle[{$a1$, $a2$, ...}, {$b1$, $b2$, ...}]'
3275        <dd>interleaves the elements of both lists, returning
3276        '{$a1$, $b1$, $a2$, $b2$, ...}'.
3277    </dl>
3278
3279    >> Riffle[{a, b, c}, x]
3280     = {a, x, b, x, c}
3281    >> Riffle[{a, b, c}, {x, y, z}]
3282     = {a, x, b, y, c, z}
3283    >> Riffle[{a, b, c, d, e, f}, {x, y, z}]
3284     = {a, x, b, y, c, z, d, x, e, y, f}
3285
3286    #> Riffle[{1, 2, 3, 4}, {x, y, z, t}]
3287     = {1, x, 2, y, 3, z, 4, t}
3288    #> Riffle[{1, 2}, {1, 2, 3}]
3289     = {1, 1, 2}
3290    #> Riffle[{1, 2}, {1, 2}]
3291     = {1, 1, 2, 2}
3292
3293    #> Riffle[{a,b,c}, {}]
3294     = {a, {}, b, {}, c}
3295    #> Riffle[{}, {}]
3296     = {}
3297    #> Riffle[{}, {a,b}]
3298     = {}
3299    """
3300
3301    def apply(self, list, sep, evaluation):
3302        "Riffle[list_List, sep_]"
3303
3304        if sep.has_form("List", None):
3305            result = riffle_lists(list.get_leaves(), sep.leaves)
3306        else:
3307            result = riffle_lists(list.get_leaves(), [sep])
3308
3309        return list.restructure("List", result, evaluation, deps=(list, sep))
3310
3311
3312def _is_sameq(same_test):
3313    # System`SameQ is protected, so nobody should ever be able to change
3314    # it (see Set::wrsym). We just check for its name here thus.
3315    return same_test.is_symbol() and same_test.get_name() == "System`SameQ"
3316
3317
3318def _test_pair(test, a, b, evaluation, name):
3319    test_expr = Expression(test, a, b)
3320    result = test_expr.evaluate(evaluation)
3321    if not (
3322        result.is_symbol() and (result.has_symbol("True") or result.has_symbol("False"))
3323    ):
3324        evaluation.message(name, "smtst", test_expr, result)
3325    return result.is_true()
3326
3327
3328class _SlowEquivalence:
3329    # models an equivalence relation through a user defined test function. for n
3330    # distinct elements (each in its own bin), we need sum(1, .., n - 1) = O(n^2)
3331    # comparisons.
3332
3333    def __init__(self, test, evaluation, name):
3334        self._groups = []
3335        self._test = test
3336        self._evaluation = evaluation
3337        self._name = name
3338
3339    def select(self, elem):
3340        return self._groups
3341
3342    def sameQ(self, a, b) -> bool:
3343        """Mathics SameQ"""
3344        return _test_pair(self._test, a, b, self._evaluation, self._name)
3345
3346
3347class _FastEquivalence:
3348    # models an equivalence relation through SameQ. for n distinct elements (each
3349    # in its own bin), we expect to make O(n) comparisons (if the hash function
3350    # does not fail us by distributing items very unevenly).
3351
3352    # IMPORTANT NOTE ON ATOM'S HASH FUNCTIONS / this code relies on this assumption:
3353    #
3354    # if SameQ[a, b] == true then hash(a) == hash(b)
3355    #
3356    # more specifically, this code bins items based on their hash code, and only if
3357    # the hash code matches, is SameQ evoked.
3358    #
3359    # this assumption has been checked for these types: Integer, Real, Complex,
3360    # String, Rational (*), Expression, Image; new atoms need proper hash functions
3361    #
3362    # (*) Rational values are sympy Rationals which are always held in reduced form
3363    # and thus are hashed correctly (see sympy/core/numbers.py:Rational.__eq__()).
3364
3365    def __init__(self):
3366        self._hashes = defaultdict(list)
3367
3368    def select(self, elem):
3369        return self._hashes[hash(elem)]
3370
3371    def sameQ(self, a, b) -> bool:
3372        """Mathics SameQ"""
3373        return a.sameQ(b)
3374
3375
3376class _GatherBin:
3377    def __init__(self, item):
3378        self._items = [item]
3379        self.add_to = self._items.append
3380
3381    def from_python(self):
3382        return Expression(SymbolList, *self._items)
3383
3384
3385class _TallyBin:
3386    def __init__(self, item):
3387        self._item = item
3388        self._count = 1
3389
3390    def add_to(self, item):
3391        self._count += 1
3392
3393    def from_python(self):
3394        return Expression(SymbolList, self._item, Integer(self._count))
3395
3396
3397class _DeleteDuplicatesBin:
3398    def __init__(self, item):
3399        self._item = item
3400        self.add_to = lambda elem: None
3401
3402    def from_python(self):
3403        return self._item
3404
3405
3406class _GatherOperation(Builtin):
3407    rules = {"%(name)s[list_]": "%(name)s[list, SameQ]"}
3408
3409    messages = {
3410        "normal": "Nonatomic expression expected at position `1` in `2`.",
3411        "list": "List expected at position `2` in `1`.",
3412        "smtst": (
3413            "Application of the SameTest yielded `1`, which evaluates "
3414            "to `2`. The SameTest must evaluate to True or False at "
3415            "every pair of elements."
3416        ),
3417    }
3418
3419    def apply(self, values, test, evaluation):
3420        "%(name)s[values_, test_]"
3421        if not self._check_list(values, test, evaluation):
3422            return
3423
3424        if _is_sameq(test):
3425            return self._gather(values, values, _FastEquivalence())
3426        else:
3427            return self._gather(
3428                values, values, _SlowEquivalence(test, evaluation, self.get_name())
3429            )
3430
3431    def _check_list(self, values, arg2, evaluation):
3432        if values.is_atom():
3433            expr = Expression(self.get_name(), values, arg2)
3434            evaluation.message(self.get_name(), "normal", 1, expr)
3435            return False
3436
3437        if values.get_head_name() != "System`List":
3438            expr = Expression(self.get_name(), values, arg2)
3439            evaluation.message(self.get_name(), "list", expr, 1)
3440            return False
3441
3442        return True
3443
3444    def _gather(self, keys, values, equivalence):
3445        bins = []
3446        Bin = self._bin
3447
3448        for key, value in zip(keys.leaves, values.leaves):
3449            selection = equivalence.select(key)
3450            for prototype, add_to_bin in selection:  # find suitable bin
3451                if equivalence.sameQ(prototype, key):
3452                    add_to_bin(value)  # add to existing bin
3453                    break
3454            else:
3455                new_bin = Bin(value)  # create new bin
3456                selection.append((key, new_bin.add_to))
3457                bins.append(new_bin)
3458
3459        return Expression(SymbolList, *[b.from_python() for b in bins])
3460
3461
3462class Gather(_GatherOperation):
3463    """
3464    <dl>
3465    <dt>'Gather[$list$, $test$]'
3466    <dd>gathers leaves of $list$ into sub lists of items that are the same according to $test$.
3467
3468    <dt>'Gather[$list$]'
3469    <dd>gathers leaves of $list$ into sub lists of items that are the same.
3470    </dl>
3471
3472    The order of the items inside the sub lists is the same as in the original list.
3473
3474    >> Gather[{1, 7, 3, 7, 2, 3, 9}]
3475     = {{1}, {7, 7}, {3, 3}, {2}, {9}}
3476
3477    >> Gather[{1/3, 2/6, 1/9}]
3478     = {{1 / 3, 1 / 3}, {1 / 9}}
3479    """
3480
3481    _bin = _GatherBin
3482
3483
3484class GatherBy(_GatherOperation):
3485    """
3486    <dl>
3487    <dt>'GatherBy[$list$, $f$]'
3488    <dd>gathers leaves of $list$ into sub lists of items whose image under $f identical.
3489
3490    <dt>'GatherBy[$list$, {$f$, $g$, ...}]'
3491    <dd>gathers leaves of $list$ into sub lists of items whose image under $f identical.
3492    Then, gathers these sub lists again into sub sub lists, that are identical under $g.
3493    </dl>
3494
3495    >> GatherBy[{{1, 3}, {2, 2}, {1, 1}}, Total]
3496     = {{{1, 3}, {2, 2}}, {{1, 1}}}
3497
3498    >> GatherBy[{"xy", "abc", "ab"}, StringLength]
3499     = {{xy, ab}, {abc}}
3500
3501    >> GatherBy[{{2, 0}, {1, 5}, {1, 0}}, Last]
3502     = {{{2, 0}, {1, 0}}, {{1, 5}}}
3503
3504    >> GatherBy[{{1, 2}, {2, 1}, {3, 5}, {5, 1}, {2, 2, 2}}, {Total, Length}]
3505     = {{{{1, 2}, {2, 1}}}, {{{3, 5}}}, {{{5, 1}}, {{2, 2, 2}}}}
3506    """
3507
3508    rules = {
3509        "GatherBy[l_]": "GatherBy[l, Identity]",
3510        "GatherBy[l_, {r__, f_}]": "Map[GatherBy[#, f]&, GatherBy[l, {r}], {Length[{r}]}]",
3511        "GatherBy[l_, {f_}]": "GatherBy[l, f]",
3512    }
3513
3514    _bin = _GatherBin
3515
3516    def apply(self, values, func, evaluation):
3517        "%(name)s[values_, func_]"
3518
3519        if not self._check_list(values, func, evaluation):
3520            return
3521
3522        keys = Expression("Map", func, values).evaluate(evaluation)
3523        if len(keys.leaves) != len(values.leaves):
3524            return
3525
3526        return self._gather(keys, values, _FastEquivalence())
3527
3528
3529class Tally(_GatherOperation):
3530    """
3531    <dl>
3532    <dt>'Tally[$list$]'
3533    <dd>counts and returns the number of occurences of objects and returns
3534    the result as a list of pairs {object, count}.
3535    <dt>'Tally[$list$, $test$]'
3536    <dd>counts the number of occurences of  objects and uses $test to
3537    determine if two objects should be counted in the same bin.
3538    </dl>
3539
3540    >> Tally[{a, b, c, b, a}]
3541     = {{a, 2}, {b, 2}, {c, 1}}
3542
3543    Tally always returns items in the order as they first appear in $list$:
3544    >> Tally[{b, b, a, a, a, d, d, d, d, c}]
3545     = {{b, 2}, {a, 3}, {d, 4}, {c, 1}}
3546    """
3547
3548    _bin = _TallyBin
3549
3550
3551class DeleteDuplicates(_GatherOperation):
3552    """
3553    <dl>
3554    <dt>'DeleteDuplicates[$list$]'
3555        <dd>deletes duplicates from $list$.
3556    <dt>'DeleteDuplicates[$list$, $test$]'
3557        <dd>deletes elements from $list$ based on whether the function
3558        $test$ yields 'True' on pairs of elements.
3559    DeleteDuplicates does not change the order of the remaining elements.
3560    </dl>
3561
3562    >> DeleteDuplicates[{1, 7, 8, 4, 3, 4, 1, 9, 9, 2, 1}]
3563     = {1, 7, 8, 4, 3, 9, 2}
3564
3565    >> DeleteDuplicates[{3,2,1,2,3,4}, Less]
3566     = {3, 2, 1}
3567
3568    #> DeleteDuplicates[{3,2,1,2,3,4}, Greater]
3569     = {3, 3, 4}
3570
3571    #> DeleteDuplicates[{}]
3572     = {}
3573    """
3574
3575    _bin = _DeleteDuplicatesBin
3576
3577
3578class _SetOperation(Builtin):
3579    messages = {
3580        "normal": "Non-atomic expression expected at position `1` in `2`.",
3581        "heads": (
3582            "Heads `1` and `2` at positions `3` and `4` are expected " "to be the same."
3583        ),
3584        "smtst": (
3585            "Application of the SameTest yielded `1`, which evaluates "
3586            "to `2`. The SameTest must evaluate to True or False at "
3587            "every pair of elements."
3588        ),
3589    }
3590
3591    options = {
3592        "SameTest": "SameQ",
3593    }
3594
3595    @staticmethod
3596    def _remove_duplicates(arg, same_test):
3597        "removes duplicates from a single operand"
3598        result = []
3599        for a in arg:
3600            if not any(same_test(a, b) for b in result):
3601                result.append(a)
3602        return result
3603
3604    def apply(self, lists, evaluation, options={}):
3605        "%(name)s[lists__, OptionsPattern[%(name)s]]"
3606
3607        seq = lists.get_sequence()
3608
3609        for pos, e in enumerate(seq):
3610            if e.is_atom():
3611                return evaluation.message(
3612                    self.get_name(),
3613                    "normal",
3614                    pos + 1,
3615                    Expression(self.get_name(), *seq),
3616                )
3617
3618        for pos, e in enumerate(zip(seq, seq[1:])):
3619            e1, e2 = e
3620            if e1.head != e2.head:
3621                return evaluation.message(
3622                    self.get_name(), "heads", e1.head, e2.head, pos + 1, pos + 2
3623                )
3624
3625        same_test = self.get_option(options, "SameTest", evaluation)
3626        operands = [l.leaves for l in seq]
3627        if not _is_sameq(same_test):
3628            sameQ = lambda a, b: _test_pair(
3629                same_test, a, b, evaluation, self.get_name()
3630            )
3631            operands = [self._remove_duplicates(op, sameQ) for op in operands]
3632            items = functools.reduce(
3633                lambda a, b: [e for e in self._elementwise(a, b, sameQ)], operands
3634            )
3635        else:
3636            items = list(
3637                functools.reduce(getattr(set, self._operation), map(set, operands))
3638            )
3639
3640        return Expression(seq[0].get_head(), *sorted(items))
3641
3642
3643class Union(_SetOperation):
3644    """
3645    <dl>
3646    <dt>'Union[$a$, $b$, ...]'
3647    <dd>gives the union of the given set or sets. The resulting list will be sorted
3648    and each element will only occur once.
3649    </dl>
3650
3651    >> Union[{5, 1, 3, 7, 1, 8, 3}]
3652     = {1, 3, 5, 7, 8}
3653
3654    >> Union[{a, b, c}, {c, d, e}]
3655     = {a, b, c, d, e}
3656
3657    >> Union[{c, b, a}]
3658     = {a, b, c}
3659
3660    >> Union[{{a, 1}, {b, 2}}, {{c, 1}, {d, 3}}, SameTest->(SameQ[Last[#1],Last[#2]]&)]
3661     = {{b, 2}, {c, 1}, {d, 3}}
3662
3663    >> Union[{1, 2, 3}, {2, 3, 4}, SameTest->Less]
3664     = {1, 2, 2, 3, 4}
3665
3666    #> Union[{1, -1, 2}, {-2, 3}, SameTest -> (Abs[#1] == Abs[#2] &)]
3667     = {-2, 1, 3}
3668    """
3669
3670    _operation = "union"
3671
3672    def _elementwise(self, a, b, sameQ: Callable[..., bool]):
3673        for eb in b:
3674            yield eb
3675        for ea in a:
3676            if not any(sameQ(eb, ea) for eb in b):
3677                yield ea
3678
3679
3680class Intersection(_SetOperation):
3681    """
3682    <dl>
3683    <dt>'Intersection[$a$, $b$, ...]'
3684    <dd>gives the intersection of the sets. The resulting list will be sorted
3685    and each element will only occur once.
3686    </dl>
3687
3688    >> Intersection[{1000, 100, 10, 1}, {1, 5, 10, 15}]
3689     = {1, 10}
3690
3691    >> Intersection[{{a, b}, {x, y}}, {{x, x}, {x, y}, {x, z}}]
3692     = {{x, y}}
3693
3694    >> Intersection[{c, b, a}]
3695     = {a, b, c}
3696
3697    >> Intersection[{1, 2, 3}, {2, 3, 4}, SameTest->Less]
3698     = {3}
3699
3700    #> Intersection[{1, -1, -2, 2, -3}, {1, -2, 2, 3}, SameTest -> (Abs[#1] == Abs[#2] &)]
3701     = {-3, -2, 1}
3702    """
3703
3704    _operation = "intersection"
3705
3706    def _elementwise(self, a, b, sameQ: Callable[..., bool]):
3707        for ea in a:
3708            if any(sameQ(eb, ea) for eb in b):
3709                yield ea
3710
3711
3712class Complement(_SetOperation):
3713    """
3714    <dl>
3715    <dt>'Complement[$all$, $e1$, $e2$, ...]'
3716        <dd>returns an expression containing the elements in the set
3717        $all$ that are not in any of $e1$, $e2$, etc.
3718    <dt>'Complement[$all$, $e1$, $e2$, ..., SameTest->$test$]'
3719        <dd>applies $test$ to the elements in $all$ and each of the
3720        $ei$ to determine equality.
3721    </dl>
3722
3723    The sets $all$, $e1$, etc can have any head, which must all match.
3724    The returned expression has the same head as the input
3725    expressions. The expression will be sorted and each element will
3726    only occur once.
3727
3728    >> Complement[{a, b, c}, {a, c}]
3729     = {b}
3730    >> Complement[{a, b, c}, {a, c}, {b}]
3731     = {}
3732    >> Complement[f[z, y, x, w], f[x], f[x, z]]
3733     = f[w, y]
3734    >> Complement[{c, b, a}]
3735     = {a, b, c}
3736
3737    #> Complement[a, b]
3738     : Non-atomic expression expected at position 1 in Complement[a, b].
3739     = Complement[a, b]
3740    #> Complement[f[a], g[b]]
3741     : Heads f and g at positions 1 and 2 are expected to be the same.
3742     = Complement[f[a], g[b]]
3743    #> Complement[{a, b, c}, {a, c}, SameTest->(True&)]
3744     = {}
3745    #> Complement[{a, b, c}, {a, c}, SameTest->(False&)]
3746     = {a, b, c}
3747    """
3748
3749    _operation = "difference"
3750
3751    def _elementwise(self, a, b, sameQ: Callable[..., bool]):
3752        for ea in a:
3753            if not any(sameQ(eb, ea) for eb in b):
3754                yield ea
3755
3756
3757class IntersectingQ(Builtin):
3758    """
3759    <dl>
3760    <dt>'IntersectingQ[$a$, $b$]'
3761    <dd>gives True if there are any common elements in $a and $b, or False if $a and $b are disjoint.
3762    </dl>
3763    """
3764
3765    rules = {"IntersectingQ[a_List, b_List]": "Length[Intersect[a, b]] > 0"}
3766
3767
3768class DisjointQ(Test):
3769    """
3770    <dl>
3771    <dt>'DisjointQ[$a$, $b$]'
3772    <dd>gives True if $a and $b are disjoint, or False if $a and $b have any common elements.
3773    </dl>
3774    """
3775
3776    rules = {"DisjointQ[a_List, b_List]": "Not[IntersectingQ[a, b]]"}
3777
3778
3779class Fold(Builtin):
3780    """
3781    <dl>
3782    <dt>'Fold[$f$, $x$, $list$]'
3783        <dd>returns the result of iteratively applying the binary
3784        operator $f$ to each element of $list$, starting with $x$.
3785    <dt>'Fold[$f$, $list$]'
3786        <dd>is equivalent to 'Fold[$f$, First[$list$], Rest[$list$]]'.
3787    </dl>
3788
3789    >> Fold[Plus, 5, {1, 1, 1}]
3790     = 8
3791    >> Fold[f, 5, {1, 2, 3}]
3792     = f[f[f[5, 1], 2], 3]
3793    """
3794
3795    rules = {
3796        "Fold[exp_, x_, head_]": "Module[{list = Level[head, 1], res = x, i = 1}, Do[res = exp[res, list[[i]]], {i, 1, Length[list]}]; res]",
3797        "Fold[exp_, head_] /; Length[head] > 0": "Fold[exp, First[head], Rest[head]]",
3798    }
3799
3800
3801class FoldList(Builtin):
3802    """
3803    <dl>
3804    <dt>'FoldList[$f$, $x$, $list$]'
3805        <dd>returns a list starting with $x$, where each element is
3806        the result of applying the binary operator $f$ to the previous
3807        result and the next element of $list$.
3808    <dt>'FoldList[$f$, $list$]'
3809        <dd>is equivalent to 'FoldList[$f$, First[$list$], Rest[$list$]]'.
3810    </dl>
3811
3812    >> FoldList[f, x, {1, 2, 3}]
3813     = {x, f[x, 1], f[f[x, 1], 2], f[f[f[x, 1], 2], 3]}
3814    >> FoldList[Times, {1, 2, 3}]
3815     = {1, 2, 6}
3816    """
3817
3818    rules = {
3819        "FoldList[exp_, x_, head_]": "Module[{i = 1}, Head[head] @@ Prepend[Table[Fold[exp, x, Take[head, i]], {i, 1, Length[head]}], x]]",
3820        "FoldList[exp_, head_]": "If[Length[head] == 0, head, FoldList[exp, First[head], Rest[head]]]",
3821    }
3822
3823
3824class Accumulate(Builtin):
3825    """
3826    <dl>
3827    <dt>'Accumulate[$list$]'
3828        <dd>accumulates the values of $list$, returning a new list.
3829    </dl>
3830
3831    >> Accumulate[{1, 2, 3}]
3832     = {1, 3, 6}
3833    """
3834
3835    rules = {"Accumulate[head_]": "FoldList[Plus, head]"}
3836
3837
3838class Total(Builtin):
3839    """
3840    <dl>
3841    <dt>'Total[$list$]'
3842        <dd>adds all values in $list$.
3843    <dt>'Total[$list$, $n$]'
3844        <dd>adds all values up to level $n$.
3845    <dt>'Total[$list$, {$n$}]'
3846        <dd>totals only the values at level {$n$}.
3847    <dt>'Total[$list$, {$n_1$, $n_2$}]'
3848        <dd>totals at levels {$n_1$, $n_2$}.
3849    </dl>
3850
3851    >> Total[{1, 2, 3}]
3852     = 6
3853    >> Total[{{1, 2, 3}, {4, 5, 6}, {7, 8 ,9}}]
3854     = {12, 15, 18}
3855
3856    Total over rows and columns
3857    >> Total[{{1, 2, 3}, {4, 5, 6}, {7, 8 ,9}}, 2]
3858     = 45
3859
3860    Total over rows instead of columns
3861    >> Total[{{1, 2, 3}, {4, 5, 6}, {7, 8 ,9}}, {2}]
3862     = {6, 15, 24}
3863    """
3864
3865    rules = {
3866        "Total[head_]": "Apply[Plus, head]",
3867        "Total[head_, n_]": "Apply[Plus, Flatten[head, n]]",
3868    }
3869
3870
3871class Reverse(Builtin):
3872    """
3873    <dl>
3874    <dt>'Reverse[$expr$]'
3875        <dd>reverses the order of $expr$'s items (on the top level)
3876    <dt>'Reverse[$expr$, $n$]'
3877        <dd>reverses the order of items in $expr$ on level $n$
3878    <dt>'Reverse[$expr$, {$n1$, $n2$, ...}]'
3879        <dd>reverses the order of items in $expr$ on levels $n1$, $n2$, ...
3880    </dl>
3881
3882    >> Reverse[{1, 2, 3}]
3883     = {3, 2, 1}
3884    >> Reverse[x[a, b, c]]
3885     = x[c, b, a]
3886    >> Reverse[{{1, 2}, {3, 4}}, 1]
3887     = {{3, 4}, {1, 2}}
3888    >> Reverse[{{1, 2}, {3, 4}}, 2]
3889     = {{2, 1}, {4, 3}}
3890    >> Reverse[{{1, 2}, {3, 4}}, {1, 2}]
3891     = {{4, 3}, {2, 1}}
3892    """
3893
3894    messages = {
3895        "ilsmp": "Positive integer or list of positive integers expected at position 2 of ``."
3896    }
3897
3898    @staticmethod
3899    def _reverse(
3900        expr, level, levels, evaluation
3901    ):  # depth >= 1, levels are expected to be unique and sorted
3902        if not isinstance(expr, Expression):
3903            return expr
3904
3905        if levels[0] == level:
3906            expr = expr.restructure(expr.head, reversed(expr.leaves), evaluation)
3907
3908            if len(levels) > 1:
3909                expr = expr.restructure(
3910                    expr.head,
3911                    [
3912                        Reverse._reverse(leaf, level + 1, levels[1:], evaluation)
3913                        for leaf in expr.leaves
3914                    ],
3915                    evaluation,
3916                )
3917        else:
3918            expr = expr.restructure(
3919                expr.head,
3920                [
3921                    Reverse._reverse(leaf, level + 1, levels, evaluation)
3922                    for leaf in expr.leaves
3923                ],
3924                evaluation,
3925            )
3926
3927        return expr
3928
3929    def apply_top_level(self, expr, evaluation):
3930        "Reverse[expr_]"
3931        return Reverse._reverse(expr, 1, (1,), evaluation)
3932
3933    def apply(self, expr, levels, evaluation):
3934        "Reverse[expr_, levels_]"
3935        if isinstance(levels, Integer):
3936            py_levels = [levels.get_int_value()]
3937        elif levels.get_head_name() == "System`List":
3938            if not levels.leaves:
3939                return expr
3940            if any(not isinstance(level, Integer) for level in levels.leaves):
3941                py_levels = None
3942            else:
3943                py_levels = sorted(
3944                    list(set(level.get_int_value() for level in levels.leaves))
3945                )
3946        else:
3947            py_levels = None
3948        if py_levels and py_levels[0] < 1:  # if py_level is not None, it's sorted
3949            py_levels = None
3950        if py_levels is None:
3951            evaluation.message("Reverse", "ilsmp", Expression("Reverse", expr, levels))
3952        else:
3953            return Reverse._reverse(expr, 1, py_levels, evaluation)
3954
3955
3956class CentralMoment(Builtin):  # see https://en.wikipedia.org/wiki/Central_moment
3957    """
3958    <dl>
3959    <dt>'CentralMoment[$list$, $r$]'
3960      <dd>gives the the $r$th central moment (i.e. the $r$th moment about the mean) of $list$.
3961    </dl>
3962
3963    >> CentralMoment[{1.1, 1.2, 1.4, 2.1, 2.4}, 4]
3964     = 0.100845
3965    """
3966
3967    rules = {
3968        "CentralMoment[list_List, r_]": "Total[(list - Mean[list]) ^ r] / Length[list]",
3969    }
3970
3971
3972class Skewness(Builtin):  # see https://en.wikipedia.org/wiki/Skewness
3973    """
3974    <dl>
3975    <dt>'Skewness[$list$]'
3976      <dd>gives Pearson's moment coefficient of skewness for $list$ (a measure for estimating
3977      the symmetry of a distribution).
3978    </dl>
3979
3980    >> Skewness[{1.1, 1.2, 1.4, 2.1, 2.4}]
3981     = 0.407041
3982    """
3983
3984    rules = {
3985        "Skewness[list_List]": "CentralMoment[list, 3] / (CentralMoment[list, 2] ^ (3 / 2))",
3986    }
3987
3988
3989class Kurtosis(Builtin):  # see https://en.wikipedia.org/wiki/Kurtosis
3990    """
3991    <dl>
3992    <dt>'Kurtosis[$list$]'
3993      <dd>gives the Pearson measure of kurtosis for $list$ (a measure of existing outliers).
3994    </dl>
3995
3996    >> Kurtosis[{1.1, 1.2, 1.4, 2.1, 2.4}]
3997     = 1.42098
3998    """
3999
4000    rules = {
4001        "Kurtosis[list_List]": "CentralMoment[list, 4] / (CentralMoment[list, 2] ^ 2)",
4002    }
4003
4004
4005class Mean(Builtin):
4006    """
4007    <dl>
4008    <dt>'Mean[$list$]'
4009      <dd>returns the statistical mean of $list$.
4010    </dl>
4011
4012    >> Mean[{26, 64, 36}]
4013     = 42
4014
4015    >> Mean[{1, 1, 2, 3, 5, 8}]
4016     = 10 / 3
4017
4018    >> Mean[{a, b}]
4019     = (a + b) / 2
4020    """
4021
4022    rules = {
4023        "Mean[list_]": "Total[list] / Length[list]",
4024    }
4025
4026
4027class _NotRectangularException(Exception):
4028    pass
4029
4030
4031class _Rectangular(Builtin):
4032    # A helper for Builtins X that allow X[{a1, a2, ...}, {b1, b2, ...}, ...] to be evaluated
4033    # as {X[{a1, b1, ...}, {a1, b2, ...}, ...]}.
4034
4035    def rect(self, l):
4036        lengths = [len(leaf.leaves) for leaf in l.leaves]
4037        if all(length == 0 for length in lengths):
4038            return  # leave as is, without error
4039
4040        n_columns = lengths[0]
4041        if any(length != n_columns for length in lengths[1:]):
4042            raise _NotRectangularException()
4043
4044        transposed = [[leaf.leaves[i] for leaf in l.leaves] for i in range(n_columns)]
4045
4046        return Expression(
4047            "List",
4048            *[
4049                Expression(self.get_name(), Expression(SymbolList, *items))
4050                for items in transposed
4051            ],
4052        )
4053
4054
4055class Variance(_Rectangular):
4056    """
4057    <dl>
4058    <dt>'Variance[$list$]'
4059      <dd>computes the variance of $list. $list$ may consist of numerical values
4060      or symbols. Numerical values may be real or complex.
4061
4062      Variance[{{$a1$, $a2$, ...}, {$b1$, $b2$, ...}, ...}] will yield
4063      {Variance[{$a1$, $b1$, ...}, Variance[{$a2$, $b2$, ...}], ...}.
4064    </dl>
4065
4066    >> Variance[{1, 2, 3}]
4067     = 1
4068
4069    >> Variance[{7, -5, 101, 3}]
4070     = 7475 / 3
4071
4072    >> Variance[{1 + 2I, 3 - 10I}]
4073     = 74
4074
4075    >> Variance[{a, a}]
4076     = 0
4077
4078    >> Variance[{{1, 3, 5}, {4, 10, 100}}]
4079     = {9 / 2, 49 / 2, 9025 / 2}
4080    """
4081
4082    messages = {
4083        "shlen": "`` must contain at least two elements.",
4084        "rectt": "Expected a rectangular array at position 1 in ``.",
4085    }
4086
4087    # for the general formulation of real and complex variance below, see for example
4088    # https://en.wikipedia.org/wiki/Variance#Generalizations
4089
4090    def apply(self, l, evaluation):
4091        "Variance[l_List]"
4092        if len(l.leaves) <= 1:
4093            evaluation.message("Variance", "shlen", l)
4094        elif all(leaf.get_head_name() == "System`List" for leaf in l.leaves):
4095            try:
4096                return self.rect(l)
4097            except _NotRectangularException:
4098                evaluation.message("Variance", "rectt", Expression("Variance", l))
4099        else:
4100            d = Expression("Subtract", l, Expression("Mean", l))
4101            return Expression(
4102                "Divide",
4103                Expression("Dot", d, Expression("Conjugate", d)),
4104                len(l.leaves) - 1,
4105            )
4106
4107
4108class StandardDeviation(_Rectangular):
4109    """
4110    <dl>
4111    <dt>'StandardDeviation[$list$]'
4112      <dd>computes the standard deviation of $list. $list$ may consist of numerical values
4113      or symbols. Numerical values may be real or complex.
4114
4115      StandardDeviation[{{$a1$, $a2$, ...}, {$b1$, $b2$, ...}, ...}] will yield
4116      {StandardDeviation[{$a1$, $b1$, ...}, StandardDeviation[{$a2$, $b2$, ...}], ...}.
4117    </dl>
4118
4119    >> StandardDeviation[{1, 2, 3}]
4120     = 1
4121
4122    >> StandardDeviation[{7, -5, 101, 100}]
4123     = Sqrt[13297] / 2
4124
4125    >> StandardDeviation[{a, a}]
4126     = 0
4127
4128    >> StandardDeviation[{{1, 10}, {-1, 20}}]
4129     = {Sqrt[2], 5 Sqrt[2]}
4130    """
4131
4132    messages = {
4133        "shlen": "`` must contain at least two elements.",
4134        "rectt": "Expected a rectangular array at position 1 in ``.",
4135    }
4136
4137    def apply(self, l, evaluation):
4138        "StandardDeviation[l_List]"
4139        if len(l.leaves) <= 1:
4140            evaluation.message("StandardDeviation", "shlen", l)
4141        elif all(leaf.get_head_name() == "System`List" for leaf in l.leaves):
4142            try:
4143                return self.rect(l)
4144            except _NotRectangularException:
4145                evaluation.message(
4146                    "StandardDeviation", "rectt", Expression("StandardDeviation", l)
4147                )
4148        else:
4149            return Expression("Sqrt", Expression("Variance", l))
4150
4151
4152class Covariance(Builtin):
4153    """
4154    <dl>
4155    <dt>'Covariance[$a$, $b$]'
4156      <dd>computes the covariance between the equal-sized vectors $a$ and $b$.
4157    </dl>
4158
4159    >> Covariance[{0.2, 0.3, 0.1}, {0.3, 0.3, -0.2}]
4160     = 0.025
4161    """
4162
4163    messages = {
4164        "shlen": "`` must contain at least two elements.",
4165        "vctmat": "`1` and `2` need to be of equal length.",
4166    }
4167
4168    def apply(self, a, b, evaluation):
4169        "Covariance[a_List, b_List]"
4170
4171        if len(a.leaves) != len(b.leaves):
4172            evaluation.message("Covariance", "vctmat", a, b)
4173        elif len(a.leaves) < 2:
4174            evaluation.message("Covariance", "shlen", a)
4175        elif len(b.leaves) < 2:
4176            evaluation.message("Covariance", "shlen", b)
4177        else:
4178            ma = Expression("Subtract", a, Expression("Mean", a))
4179            mb = Expression("Subtract", b, Expression("Mean", b))
4180            return Expression(
4181                "Divide",
4182                Expression("Dot", ma, Expression("Conjugate", mb)),
4183                len(a.leaves) - 1,
4184            )
4185
4186
4187class Correlation(Builtin):
4188    """
4189    <dl>
4190    <dt>'Correlation[$a$, $b$]'
4191      <dd>computes Pearson's correlation of two equal-sized vectors $a$ and $b$.
4192    </dl>
4193
4194    An example from Wikipedia:
4195
4196    >> Correlation[{10, 8, 13, 9, 11, 14, 6, 4, 12, 7, 5}, {8.04, 6.95, 7.58, 8.81, 8.33, 9.96, 7.24, 4.26, 10.84, 4.82, 5.68}]
4197     = 0.816421
4198    """
4199
4200    messages = {
4201        "shlen": "`` must contain at least two elements.",
4202        "vctmat": "`1` and `2` need to be of equal length.",
4203    }
4204
4205    def apply(self, a, b, evaluation):
4206        "Correlation[a_List, b_List]"
4207
4208        if len(a.leaves) != len(b.leaves):
4209            evaluation.message("Correlation", "vctmat", a, b)
4210        elif len(a.leaves) < 2:
4211            evaluation.message("Correlation", "shlen", a)
4212        elif len(b.leaves) < 2:
4213            evaluation.message("Correlation", "shlen", b)
4214        else:
4215            da = Expression("StandardDeviation", a)
4216            db = Expression("StandardDeviation", b)
4217            return Expression(
4218                "Divide", Expression("Covariance", a, b), Expression("Times", da, db)
4219            )
4220
4221
4222class _Rotate(Builtin):
4223    messages = {"rspec": "`` should be an integer or a list of integers."}
4224
4225    def _rotate(self, expr, n, evaluation):
4226        if not isinstance(expr, Expression):
4227            return expr
4228
4229        leaves = expr.leaves
4230        if not leaves:
4231            return expr
4232
4233        index = (self._sign * n[0]) % len(leaves)  # with Python's modulo: index >= 1
4234        new_leaves = chain(leaves[index:], leaves[:index])
4235
4236        if len(n) > 1:
4237            new_leaves = [self._rotate(item, n[1:], evaluation) for item in new_leaves]
4238
4239        return expr.restructure(expr.head, new_leaves, evaluation)
4240
4241    def apply_one(self, expr, evaluation):
4242        "%(name)s[expr_]"
4243        return self._rotate(expr, [1], evaluation)
4244
4245    def apply(self, expr, n, evaluation):
4246        "%(name)s[expr_, n_]"
4247        if isinstance(n, Integer):
4248            py_cycles = [n.get_int_value()]
4249        elif n.get_head_name() == "System`List" and all(
4250            isinstance(x, Integer) for x in n.leaves
4251        ):
4252            py_cycles = [x.get_int_value() for x in n.leaves]
4253            if not py_cycles:
4254                return expr
4255        else:
4256            evaluation.message(self.get_name(), "rspec", n)
4257            return
4258
4259        return self._rotate(expr, py_cycles, evaluation)
4260
4261
4262class RotateLeft(_Rotate):
4263    """
4264    <dl>
4265    <dt>'RotateLeft[$expr$]'
4266        <dd>rotates the items of $expr$' by one item to the left.
4267    <dt>'RotateLeft[$expr$, $n$]'
4268        <dd>rotates the items of $expr$' by $n$ items to the left.
4269    <dt>'RotateLeft[$expr$, {$n1$, $n2$, ...}]'
4270        <dd>rotates the items of $expr$' by $n1$ items to the left at the first level, by $n2$ items to the left at
4271        the second level, and so on.
4272    </dl>
4273
4274    >> RotateLeft[{1, 2, 3}]
4275     = {2, 3, 1}
4276    >> RotateLeft[Range[10], 3]
4277     = {4, 5, 6, 7, 8, 9, 10, 1, 2, 3}
4278    >> RotateLeft[x[a, b, c], 2]
4279     = x[c, a, b]
4280    >> RotateLeft[{{a, b, c}, {d, e, f}, {g, h, i}}, {1, 2}]
4281     = {{f, d, e}, {i, g, h}, {c, a, b}}
4282    """
4283
4284    _sign = 1
4285
4286
4287class RotateRight(_Rotate):
4288    """
4289    <dl>
4290    <dt>'RotateRight[$expr$]'
4291        <dd>rotates the items of $expr$' by one item to the right.
4292    <dt>'RotateRight[$expr$, $n$]'
4293        <dd>rotates the items of $expr$' by $n$ items to the right.
4294    <dt>'RotateRight[$expr$, {$n1$, $n2$, ...}]'
4295        <dd>rotates the items of $expr$' by $n1$ items to the right at the first level, by $n2$ items to the right at
4296        the second level, and so on.
4297    </dl>
4298
4299    >> RotateRight[{1, 2, 3}]
4300     = {3, 1, 2}
4301    >> RotateRight[Range[10], 3]
4302     = {8, 9, 10, 1, 2, 3, 4, 5, 6, 7}
4303    >> RotateRight[x[a, b, c], 2]
4304     = x[b, c, a]
4305    >> RotateRight[{{a, b, c}, {d, e, f}, {g, h, i}}, {1, 2}]
4306     = {{h, i, g}, {b, c, a}, {e, f, d}}
4307    """
4308
4309    _sign = -1
4310
4311
4312class Median(_Rectangular):
4313    """
4314    <dl>
4315    <dt>'Median[$list$]'
4316      <dd>returns the median of $list$.
4317    </dl>
4318
4319    >> Median[{26, 64, 36}]
4320     = 36
4321
4322    For lists with an even number of elements, Median returns the mean of the two middle values:
4323    >> Median[{-11, 38, 501, 1183}]
4324     = 539 / 2
4325
4326    Passing a matrix returns the medians of the respective columns:
4327    >> Median[{{100, 1, 10, 50}, {-1, 1, -2, 2}}]
4328     = {99 / 2, 1, 4, 26}
4329    """
4330
4331    messages = {"rectn": "Expected a rectangular array of numbers at position 1 in ``."}
4332
4333    def apply(self, l, evaluation):
4334        "Median[l_List]"
4335        if not l.leaves:
4336            return
4337        if all(leaf.get_head_name() == "System`List" for leaf in l.leaves):
4338            try:
4339                return self.rect(l)
4340            except _NotRectangularException:
4341                evaluation.message("Median", "rectn", Expression("Median", l))
4342        elif all(leaf.is_numeric() for leaf in l.leaves):
4343            v = l.get_mutable_leaves()  # copy needed for introselect
4344            n = len(v)
4345            if n % 2 == 0:  # even number of elements?
4346                i = n // 2
4347                a = introselect(v, i)
4348                b = introselect(v, i - 1)
4349                return Expression("Divide", Expression("Plus", a, b), 2)
4350            else:
4351                i = n // 2
4352                return introselect(v, i)
4353        else:
4354            evaluation.message("Median", "rectn", Expression("Median", l))
4355
4356
4357class RankedMin(Builtin):
4358    """
4359    <dl>
4360    <dt>'RankedMin[$list$, $n$]'
4361      <dd>returns the $n$th smallest element of $list$ (with $n$ = 1 yielding the smallest element,
4362      $n$ = 2 yielding the second smallest element, and so on).
4363    </dl>
4364
4365    >> RankedMin[{482, 17, 181, -12}, 2]
4366     = 17
4367    """
4368
4369    messages = {
4370        "intpm": "Expected positive integer at position 2 in ``.",
4371        "rank": "The specified rank `1` is not between 1 and `2`.",
4372    }
4373
4374    def apply(self, l, n, evaluation):
4375        "RankedMin[l_List, n_Integer]"
4376        py_n = n.get_int_value()
4377        if py_n < 1:
4378            evaluation.message("RankedMin", "intpm", Expression("RankedMin", l, n))
4379        elif py_n > len(l.leaves):
4380            evaluation.message("RankedMin", "rank", py_n, len(l.leaves))
4381        else:
4382            return introselect(l.get_mutable_leaves(), py_n - 1)
4383
4384
4385class RankedMax(Builtin):
4386    """
4387    <dl>
4388    <dt>'RankedMax[$list$, $n$]'
4389      <dd>returns the $n$th largest element of $list$ (with $n$ = 1 yielding the largest element,
4390      $n$ = 2 yielding the second largest element, and so on).
4391    </dl>
4392
4393    >> RankedMax[{482, 17, 181, -12}, 2]
4394     = 181
4395    """
4396
4397    messages = {
4398        "intpm": "Expected positive integer at position 2 in ``.",
4399        "rank": "The specified rank `1` is not between 1 and `2`.",
4400    }
4401
4402    def apply(self, l, n, evaluation):
4403        "RankedMax[l_List, n_Integer]"
4404        py_n = n.get_int_value()
4405        if py_n < 1:
4406            evaluation.message("RankedMax", "intpm", Expression("RankedMax", l, n))
4407        elif py_n > len(l.leaves):
4408            evaluation.message("RankedMax", "rank", py_n, len(l.leaves))
4409        else:
4410            return introselect(l.get_mutable_leaves(), len(l.leaves) - py_n)
4411
4412
4413class Quantile(Builtin):
4414    """
4415    <dl>
4416    <dt>'Quantile[$list$, $q$]'
4417      <dd>returns the $q$th quantile of $list$.
4418    </dl>
4419
4420    >> Quantile[Range[11], 1/3]
4421     = 4
4422
4423    >> Quantile[Range[16], 1/4]
4424     = 5
4425    """
4426
4427    rules = {
4428        "Quantile[list_List, q_, abcd_]": "Quantile[list, {q}, abcd]",
4429        "Quantile[list_List, q_]": "Quantile[list, q, {{0, 1}, {1, 0}}]",
4430    }
4431
4432    messages = {
4433        "nquan": "The quantile `1` has to be between 0 and 1.",
4434    }
4435
4436    def apply(self, l, qs, a, b, c, d, evaluation):
4437        """Quantile[l_List, qs_List, {{a_, b_}, {c_, d_}}]"""
4438
4439        n = len(l.leaves)
4440        partially_sorted = l.get_mutable_leaves()
4441
4442        def ranked(i):
4443            return introselect(partially_sorted, min(max(0, i - 1), n - 1))
4444
4445        numeric_qs = qs.evaluate(evaluation).numerify(evaluation)
4446        results = []
4447
4448        for q in numeric_qs.leaves:
4449            py_q = q.to_mpmath()
4450
4451            if py_q is None or not 0.0 <= py_q <= 1.0:
4452                evaluation.message("Quantile", "nquan", q)
4453                return
4454
4455            x = Expression(
4456                "Plus", a, Expression("Times", Expression("Plus", Integer(n), b), q)
4457            )
4458
4459            numeric_x = x.evaluate(evaluation).numerify(evaluation)
4460
4461            if isinstance(numeric_x, Integer):
4462                results.append(ranked(numeric_x.get_int_value()))
4463            else:
4464                py_x = numeric_x.to_mpmath()
4465
4466                if py_x is None:
4467                    return
4468
4469                from mpmath import floor as mpfloor, ceil as mpceil
4470
4471                if c.get_int_value() == 1 and d.get_int_value() == 0:  # k == 1?
4472                    results.append(ranked(int(mpceil(py_x))))
4473                else:
4474                    py_floor_x = mpfloor(py_x)
4475                    s0 = ranked(int(py_floor_x))
4476                    s1 = ranked(int(mpceil(py_x)))
4477
4478                    k = Expression(
4479                        "Plus",
4480                        c,
4481                        Expression(
4482                            "Times",
4483                            d,
4484                            Expression("Subtract", x, Expression("Floor", x)),
4485                        ),
4486                    )
4487
4488                    results.append(
4489                        Expression(
4490                            "Plus",
4491                            s0,
4492                            Expression("Times", k, Expression("Subtract", s1, s0)),
4493                        )
4494                    )
4495
4496        if len(results) == 1:
4497            return results[0]
4498        else:
4499            return Expression(SymbolList, *results)
4500
4501
4502class Quartiles(Builtin):
4503    """
4504    <dl>
4505    <dt>'Quartiles[$list$]'
4506      <dd>returns the 1/4, 1/2, and 3/4 quantiles of $list$.
4507    </dl>
4508
4509    >> Quartiles[Range[25]]
4510     = {27 / 4, 13, 77 / 4}
4511    """
4512
4513    rules = {
4514        "Quartiles[list_List]": "Quantile[list, {1/4, 1/2, 3/4}, {{1/2, 0}, {0, 1}}]",
4515    }
4516
4517
4518class _RankedTake(Builtin):
4519    messages = {
4520        "intpm": "Expected non-negative integer at position `1` in `2`.",
4521        "rank": "The specified rank `1` is not between 1 and `2`.",
4522    }
4523
4524    options = {
4525        "ExcludedForms": "Automatic",
4526    }
4527
4528    def _compute(self, l, n, evaluation, options, f=None):
4529        try:
4530            limit = CountableInteger.from_expression(n)
4531        except MessageException as e:
4532            e.message(evaluation)
4533            return
4534        except NegativeIntegerException:
4535            if f:
4536                args = (3, Expression(self.get_name(), l, f, n))
4537            else:
4538                args = (2, Expression(self.get_name(), l, n))
4539            evaluation.message(self.get_name(), "intpm", *args)
4540            return
4541
4542        if limit is None:
4543            return
4544
4545        if limit == 0:
4546            return Expression(SymbolList)
4547        else:
4548            excluded = self.get_option(options, "ExcludedForms", evaluation)
4549            if excluded:
4550                if (
4551                    isinstance(excluded, Symbol)
4552                    and excluded.get_name() == "System`Automatic"
4553                ):
4554
4555                    def exclude(item):
4556                        if isinstance(item, Symbol) and item.get_name() in (
4557                            "System`None",
4558                            "System`Null",
4559                            "System`Indeterminate",
4560                        ):
4561                            return True
4562                        elif item.get_head_name() == "System`Missing":
4563                            return True
4564                        else:
4565                            return False
4566
4567                else:
4568                    excluded = Expression("Alternatives", *excluded.leaves)
4569
4570                    def exclude(item):
4571                        return (
4572                            Expression("MatchQ", item, excluded)
4573                            .evaluate(evaluation)
4574                            .is_true()
4575                        )
4576
4577                filtered = [leaf for leaf in l.leaves if not exclude(leaf)]
4578            else:
4579                filtered = l.leaves
4580
4581            if limit > len(filtered):
4582                if not limit.is_upper_limit():
4583                    evaluation.message(
4584                        self.get_name(), "rank", limit.get_int_value(), len(filtered)
4585                    )
4586                    return
4587                else:
4588                    py_n = len(filtered)
4589            else:
4590                py_n = limit.get_int_value()
4591
4592            if py_n < 1:
4593                return Expression(SymbolList)
4594
4595            if f:
4596                heap = [
4597                    (Expression(f, leaf).evaluate(evaluation), leaf, i)
4598                    for i, leaf in enumerate(filtered)
4599                ]
4600                leaf_pos = 1  # in tuple above
4601            else:
4602                heap = [(leaf, i) for i, leaf in enumerate(filtered)]
4603                leaf_pos = 0  # in tuple above
4604
4605            if py_n == 1:
4606                result = [self._get_1(heap)]
4607            else:
4608                result = self._get_n(py_n, heap)
4609
4610            return l.restructure("List", [x[leaf_pos] for x in result], evaluation)
4611
4612
4613class _RankedTakeSmallest(_RankedTake):
4614    def _get_1(self, a):
4615        return min(a)
4616
4617    def _get_n(self, n, heap):
4618        return heapq.nsmallest(n, heap)
4619
4620
4621class _RankedTakeLargest(_RankedTake):
4622    def _get_1(self, a):
4623        return max(a)
4624
4625    def _get_n(self, n, heap):
4626        return heapq.nlargest(n, heap)
4627
4628
4629class TakeLargest(_RankedTakeLargest):
4630    """
4631    <dl>
4632    <dt>'TakeLargest[$list$, $f$, $n$]'
4633        <dd>returns the a sorted list of the $n$ largest items in $list$.
4634    </dl>
4635
4636    >> TakeLargest[{100, -1, 50, 10}, 2]
4637     = {100, 50}
4638
4639    None, Null, Indeterminate and expressions with head Missing are ignored
4640    by default:
4641    >> TakeLargest[{-8, 150, Missing[abc]}, 2]
4642     = {150, -8}
4643
4644    You may specify which items are ignored using the option ExcludedForms:
4645    >> TakeLargest[{-8, 150, Missing[abc]}, 2, ExcludedForms -> {}]
4646     = {Missing[abc], 150}
4647    """
4648
4649    def apply(self, l, n, evaluation, options):
4650        "TakeLargest[l_List, n_, OptionsPattern[TakeLargest]]"
4651        return self._compute(l, n, evaluation, options)
4652
4653
4654class TakeLargestBy(_RankedTakeLargest):
4655    """
4656    <dl>
4657    <dt>'TakeLargestBy[$list$, $f$, $n$]'
4658        <dd>returns the a sorted list of the $n$ largest items in $list$
4659        using $f$ to retrieve the items' keys to compare them.
4660    </dl>
4661
4662    For details on how to use the ExcludedForms option, see TakeLargest[].
4663
4664    >> TakeLargestBy[{{1, -1}, {10, 100}, {23, 7, 8}, {5, 1}}, Total, 2]
4665     = {{10, 100}, {23, 7, 8}}
4666
4667    >> TakeLargestBy[{"abc", "ab", "x"}, StringLength, 1]
4668     = {abc}
4669    """
4670
4671    def apply(self, l, f, n, evaluation, options):
4672        "TakeLargestBy[l_List, f_, n_, OptionsPattern[TakeLargestBy]]"
4673        return self._compute(l, n, evaluation, options, f=f)
4674
4675
4676class TakeSmallest(_RankedTakeSmallest):
4677    """
4678    <dl>
4679    <dt>'TakeSmallest[$list$, $f$, $n$]'
4680        <dd>returns the a sorted list of the $n$ smallest items in $list$.
4681    </dl>
4682
4683    For details on how to use the ExcludedForms option, see TakeLargest[].
4684
4685    >> TakeSmallest[{100, -1, 50, 10}, 2]
4686     = {-1, 10}
4687    """
4688
4689    def apply(self, l, n, evaluation, options):
4690        "TakeSmallest[l_List, n_, OptionsPattern[TakeSmallest]]"
4691        return self._compute(l, n, evaluation, options)
4692
4693
4694class TakeSmallestBy(_RankedTakeSmallest):
4695    """
4696    <dl>
4697    <dt>'TakeSmallestBy[$list$, $f$, $n$]'
4698        <dd>returns the a sorted list of the $n$ smallest items in $list$
4699        using $f$ to retrieve the items' keys to compare them.
4700    </dl>
4701
4702    For details on how to use the ExcludedForms option, see TakeLargest[].
4703
4704    >> TakeSmallestBy[{{1, -1}, {10, 100}, {23, 7, 8}, {5, 1}}, Total, 2]
4705     = {{1, -1}, {5, 1}}
4706
4707    >> TakeSmallestBy[{"abc", "ab", "x"}, StringLength, 1]
4708     = {x}
4709    """
4710
4711    def apply(self, l, f, n, evaluation, options):
4712        "TakeSmallestBy[l_List, f_, n_, OptionsPattern[TakeSmallestBy]]"
4713        return self._compute(l, n, evaluation, options, f=f)
4714
4715
4716class _IllegalPaddingDepth(Exception):
4717    def __init__(self, level):
4718        self.level = level
4719
4720
4721class _Pad(Builtin):
4722    messages = {
4723        "normal": "Expression at position 1 in `` must not be an atom.",
4724        "level": "Cannot pad list `3` which has `4` using padding `1` which specifies `2`.",
4725        "ilsm": "Expected an integer or a list of integers at position `1` in `2`.",
4726    }
4727
4728    rules = {"%(name)s[l_]": "%(name)s[l, Automatic]"}
4729
4730    @staticmethod
4731    def _find_dims(expr):
4732        def dive(expr, level):
4733            if isinstance(expr, Expression):
4734                if expr.leaves:
4735                    return max(dive(x, level + 1) for x in expr.leaves)
4736                else:
4737                    return level + 1
4738            else:
4739                return level
4740
4741        def calc(expr, dims, level):
4742            if isinstance(expr, Expression):
4743                for x in expr.leaves:
4744                    calc(x, dims, level + 1)
4745                dims[level] = max(dims[level], len(expr.leaves))
4746
4747        dims = [0] * dive(expr, 0)
4748        calc(expr, dims, 0)
4749        return dims
4750
4751    @staticmethod
4752    def _build(l, n, x, m, level, mode):  # mode < 0 for left pad, > 0 for right pad
4753        if not n:
4754            return l
4755        if not isinstance(l, Expression):
4756            raise _IllegalPaddingDepth(level)
4757
4758        if isinstance(m, (list, tuple)):
4759            current_m = m[0] if m else 0
4760            next_m = m[1:]
4761        else:
4762            current_m = m
4763            next_m = m
4764
4765        def clip(a, d, s):
4766            assert d != 0
4767            if s < 0:
4768                return a[-d:]  # end with a[-1]
4769            else:
4770                return a[:d]  # start with a[0]
4771
4772        def padding(amount, sign):
4773            if amount == 0:
4774                return []
4775            elif len(n) > 1:
4776                return [
4777                    _Pad._build(
4778                        Expression(SymbolList), n[1:], x, next_m, level + 1, mode
4779                    )
4780                ] * amount
4781            else:
4782                return clip(x * (1 + amount // len(x)), amount, sign)
4783
4784        leaves = l.leaves
4785        d = n[0] - len(leaves)
4786        if d < 0:
4787            new_leaves = clip(leaves, d, mode)
4788            padding_main = []
4789        elif d >= 0:
4790            new_leaves = leaves
4791            padding_main = padding(d, mode)
4792
4793        if current_m > 0:
4794            padding_margin = padding(
4795                min(current_m, len(new_leaves) + len(padding_main)), -mode
4796            )
4797
4798            if len(padding_margin) > len(padding_main):
4799                padding_main = []
4800                new_leaves = clip(
4801                    new_leaves, -(len(padding_margin) - len(padding_main)), mode
4802                )
4803            elif len(padding_margin) > 0:
4804                padding_main = clip(padding_main, -len(padding_margin), mode)
4805        else:
4806            padding_margin = []
4807
4808        if len(n) > 1:
4809            new_leaves = (
4810                _Pad._build(e, n[1:], x, next_m, level + 1, mode) for e in new_leaves
4811            )
4812
4813        if mode < 0:
4814            parts = (padding_main, new_leaves, padding_margin)
4815        else:
4816            parts = (padding_margin, new_leaves, padding_main)
4817
4818        return Expression(l.get_head(), *list(chain(*parts)))
4819
4820    def _pad(self, in_l, in_n, in_x, in_m, evaluation, expr):
4821        if not isinstance(in_l, Expression):
4822            evaluation.message(self.get_name(), "normal", expr())
4823            return
4824
4825        py_n = None
4826        if isinstance(in_n, Symbol) and in_n.get_name() == "System`Automatic":
4827            py_n = _Pad._find_dims(in_l)
4828        elif in_n.get_head_name() == "System`List":
4829            if all(isinstance(leaf, Integer) for leaf in in_n.leaves):
4830                py_n = [leaf.get_int_value() for leaf in in_n.leaves]
4831        elif isinstance(in_n, Integer):
4832            py_n = [in_n.get_int_value()]
4833
4834        if py_n is None:
4835            evaluation.message(self.get_name(), "ilsm", 2, expr())
4836            return
4837
4838        if in_x.get_head_name() == "System`List":
4839            py_x = in_x.leaves
4840        else:
4841            py_x = [in_x]
4842
4843        if isinstance(in_m, Integer):
4844            py_m = in_m.get_int_value()
4845        else:
4846            if not all(isinstance(x, Integer) for x in in_m.leaves):
4847                evaluation.message(self.get_name(), "ilsm", 4, expr())
4848                return
4849            py_m = [x.get_int_value() for x in in_m.leaves]
4850
4851        try:
4852            return _Pad._build(in_l, py_n, py_x, py_m, 1, self._mode)
4853        except _IllegalPaddingDepth as e:
4854
4855            def levels(k):
4856                if k == 1:
4857                    return "1 level"
4858                else:
4859                    return "%d levels" % k
4860
4861            evaluation.message(
4862                self.get_name(),
4863                "level",
4864                in_n,
4865                levels(len(py_n)),
4866                in_l,
4867                levels(e.level - 1),
4868            )
4869            return None
4870
4871    def apply_zero(self, l, n, evaluation):
4872        "%(name)s[l_, n_]"
4873        return self._pad(
4874            l,
4875            n,
4876            Integer0,
4877            Integer0,
4878            evaluation,
4879            lambda: Expression(self.get_name(), l, n),
4880        )
4881
4882    def apply(self, l, n, x, evaluation):
4883        "%(name)s[l_, n_, x_]"
4884        return self._pad(
4885            l,
4886            n,
4887            x,
4888            Integer0,
4889            evaluation,
4890            lambda: Expression(self.get_name(), l, n, x),
4891        )
4892
4893    def apply_margin(self, l, n, x, m, evaluation):
4894        "%(name)s[l_, n_, x_, m_]"
4895        return self._pad(
4896            l, n, x, m, evaluation, lambda: Expression(self.get_name(), l, n, x, m)
4897        )
4898
4899
4900class PadLeft(_Pad):
4901    """
4902    <dl>
4903    <dt>'PadLeft[$list$, $n$]'
4904        <dd>pads $list$ to length $n$ by adding 0 on the left.
4905    <dt>'PadLeft[$list$, $n$, $x$]'
4906        <dd>pads $list$ to length $n$ by adding $x$ on the left.
4907    <dt>'PadLeft[$list$, {$n1$, $n2, ...}, $x$]'
4908        <dd>pads $list$ to lengths $n1$, $n2$ at levels 1, 2, ... respectively by adding $x$ on the left.
4909    <dt>'PadLeft[$list$, $n$, $x$, $m$]'
4910        <dd>pads $list$ to length $n$ by adding $x$ on the left and adding a margin of $m$ on the right.
4911    <dt>'PadLeft[$list$, $n$, $x$, {$m1$, $m2$, ...}]'
4912        <dd>pads $list$ to length $n$ by adding $x$ on the left and adding margins of $m1$, $m2$, ...
4913         on levels 1, 2, ... on the right.
4914    <dt>'PadLeft[$list$]'
4915        <dd>turns the ragged list $list$ into a regular list by adding 0 on the left.
4916    </dl>
4917
4918    >> PadLeft[{1, 2, 3}, 5]
4919     = {0, 0, 1, 2, 3}
4920    >> PadLeft[x[a, b, c], 5]
4921     = x[0, 0, a, b, c]
4922    >> PadLeft[{1, 2, 3}, 2]
4923     = {2, 3}
4924    >> PadLeft[{{}, {1, 2}, {1, 2, 3}}]
4925     = {{0, 0, 0}, {0, 1, 2}, {1, 2, 3}}
4926    >> PadLeft[{1, 2, 3}, 10, {a, b, c}, 2]
4927     = {b, c, a, b, c, 1, 2, 3, a, b}
4928    >> PadLeft[{{1, 2, 3}}, {5, 2}, x, 1]
4929     = {{x, x}, {x, x}, {x, x}, {3, x}, {x, x}}
4930    """
4931
4932    _mode = -1
4933
4934
4935class PadRight(_Pad):
4936    """
4937    <dl>
4938    <dt>'PadRight[$list$, $n$]'
4939        <dd>pads $list$ to length $n$ by adding 0 on the right.
4940    <dt>'PadRight[$list$, $n$, $x$]'
4941        <dd>pads $list$ to length $n$ by adding $x$ on the right.
4942    <dt>'PadRight[$list$, {$n1$, $n2, ...}, $x$]'
4943        <dd>pads $list$ to lengths $n1$, $n2$ at levels 1, 2, ... respectively by adding $x$ on the right.
4944    <dt>'PadRight[$list$, $n$, $x$, $m$]'
4945        <dd>pads $list$ to length $n$ by adding $x$ on the left and adding a margin of $m$ on the left.
4946    <dt>'PadRight[$list$, $n$, $x$, {$m1$, $m2$, ...}]'
4947        <dd>pads $list$ to length $n$ by adding $x$ on the right and adding margins of $m1$, $m2$, ...
4948         on levels 1, 2, ... on the left.
4949    <dt>'PadRight[$list$]'
4950        <dd>turns the ragged list $list$ into a regular list by adding 0 on the right.
4951    </dl>
4952
4953    >> PadRight[{1, 2, 3}, 5]
4954     = {1, 2, 3, 0, 0}
4955    >> PadRight[x[a, b, c], 5]
4956     = x[a, b, c, 0, 0]
4957    >> PadRight[{1, 2, 3}, 2]
4958     = {1, 2}
4959    >> PadRight[{{}, {1, 2}, {1, 2, 3}}]
4960     = {{0, 0, 0}, {1, 2, 0}, {1, 2, 3}}
4961    >> PadRight[{1, 2, 3}, 10, {a, b, c}, 2]
4962     = {b, c, 1, 2, 3, a, b, c, a, b}
4963    >> PadRight[{{1, 2, 3}}, {5, 2}, x, 1]
4964     = {{x, x}, {x, 1}, {x, x}, {x, x}, {x, x}}
4965    """
4966
4967    _mode = 1
4968
4969
4970class _IllegalDistance(Exception):
4971    def __init__(self, distance):
4972        self.distance = distance
4973
4974
4975class _IllegalDataPoint(Exception):
4976    pass
4977
4978
4979def _to_real_distance(d):
4980    if not isinstance(d, (Real, Integer)):
4981        raise _IllegalDistance(d)
4982
4983    mpd = d.to_mpmath()
4984    if mpd is None or mpd < 0:
4985        raise _IllegalDistance(d)
4986
4987    return mpd
4988
4989
4990class _PrecomputedDistances(PrecomputedDistances):
4991    # computes all n^2 distances for n points with one big evaluation in the beginning.
4992
4993    def __init__(self, df, p, evaluation):
4994        distances_form = [df(p[i], p[j]) for i in range(len(p)) for j in range(i)]
4995        distances = Expression(
4996            SymbolN, Expression(SymbolList, *distances_form)
4997        ).evaluate(evaluation)
4998        mpmath_distances = [_to_real_distance(d) for d in distances.leaves]
4999        super(_PrecomputedDistances, self).__init__(mpmath_distances)
5000
5001
5002class _LazyDistances(LazyDistances):
5003    # computes single distances only as needed, caches already computed distances.
5004
5005    def __init__(self, df, p, evaluation):
5006        super(_LazyDistances, self).__init__()
5007        self._df = df
5008        self._p = p
5009        self._evaluation = evaluation
5010
5011    def _compute_distance(self, i, j):
5012        p = self._p
5013        d = Expression(SymbolN, self._df(p[i], p[j])).evaluate(self._evaluation)
5014        return _to_real_distance(d)
5015
5016
5017def _dist_repr(p):
5018    dist_p = repr_p = None
5019    if p.has_form("Rule", 2):
5020        if all(q.get_head_name() == "System`List" for q in p.leaves):
5021            dist_p, repr_p = (q.leaves for q in p.leaves)
5022        elif (
5023            p.leaves[0].get_head_name() == "System`List"
5024            and p.leaves[1].get_name() == "System`Automatic"
5025        ):
5026            dist_p = p.leaves[0].leaves
5027            repr_p = [Integer(i + 1) for i in range(len(dist_p))]
5028    elif p.get_head_name() == "System`List":
5029        if all(q.get_head_name() == "System`Rule" for q in p.leaves):
5030            dist_p, repr_p = ([q.leaves[i] for q in p.leaves] for i in range(2))
5031        else:
5032            dist_p = repr_p = p.leaves
5033    return dist_p, repr_p
5034
5035
5036class _Cluster(Builtin):
5037    options = {
5038        "Method": "Optimize",
5039        "DistanceFunction": "Automatic",
5040        "RandomSeed": "Automatic",
5041    }
5042
5043    messages = {
5044        "amtd": "`1` failed to pick a suitable distance function for `2`.",
5045        "bdmtd": 'Method in `` must be either "Optimize", "Agglomerate" or "KMeans".',
5046        "intpm": "Positive integer expected at position 2 in ``.",
5047        "list": "Expected a list or a rule with equally sized lists at position 1 in ``.",
5048        "nclst": "Cannot find more clusters than there are elements: `1` is larger than `2`.",
5049        "xnum": "The distance function returned ``, which is not a non-negative real value.",
5050        "rseed": "The random seed specified through `` must be an integer or Automatic.",
5051        "kmsud": "KMeans only supports SquaredEuclideanDistance as distance measure.",
5052    }
5053
5054    _criteria = {
5055        "Optimize": AutomaticSplitCriterion,
5056        "Agglomerate": AutomaticMergeCriterion,
5057        "KMeans": None,
5058    }
5059
5060    def _cluster(self, p, k, mode, evaluation, options, expr):
5061        method_string, method = self.get_option_string(options, "Method", evaluation)
5062        if method_string not in ("Optimize", "Agglomerate", "KMeans"):
5063            evaluation.message(
5064                self.get_name(), "bdmtd", Expression(SymbolRule, "Method", method)
5065            )
5066            return
5067
5068        dist_p, repr_p = _dist_repr(p)
5069
5070        if dist_p is None or len(dist_p) != len(repr_p):
5071            evaluation.message(self.get_name(), "list", expr)
5072            return
5073
5074        if not dist_p:
5075            return Expression(SymbolList)
5076
5077        if k is not None:  # the number of clusters k is specified as an integer.
5078            if not isinstance(k, Integer):
5079                evaluation.message(self.get_name(), "intpm", expr)
5080                return
5081            py_k = k.get_int_value()
5082            if py_k < 1:
5083                evaluation.message(self.get_name(), "intpm", expr)
5084                return
5085            if py_k > len(dist_p):
5086                evaluation.message(self.get_name(), "nclst", py_k, len(dist_p))
5087                return
5088            elif py_k == 1:
5089                return Expression(SymbolList, *repr_p)
5090            elif py_k == len(dist_p):
5091                return Expression(
5092                    SymbolList, [Expression(SymbolList, q) for q in repr_p]
5093                )
5094        else:  # automatic detection of k. choose a suitable method here.
5095            if len(dist_p) <= 2:
5096                return Expression(SymbolList, *repr_p)
5097            constructor = self._criteria.get(method_string)
5098            py_k = (constructor, {}) if constructor else None
5099
5100        seed_string, seed = self.get_option_string(options, "RandomSeed", evaluation)
5101        if seed_string == "Automatic":
5102            py_seed = 12345
5103        elif isinstance(seed, Integer):
5104            py_seed = seed.get_int_value()
5105        else:
5106            evaluation.message(
5107                self.get_name(), "rseed", Expression(SymbolRule, "RandomSeed", seed)
5108            )
5109            return
5110
5111        distance_function_string, distance_function = self.get_option_string(
5112            options, "DistanceFunction", evaluation
5113        )
5114        if distance_function_string == "Automatic":
5115            from mathics.builtin.tensors import get_default_distance
5116
5117            distance_function = get_default_distance(dist_p)
5118            if distance_function is None:
5119                name_of_builtin = strip_context(self.get_name())
5120                evaluation.message(
5121                    self.get_name(),
5122                    "amtd",
5123                    name_of_builtin,
5124                    Expression(SymbolList, *dist_p),
5125                )
5126                return
5127
5128        if (
5129            method_string == "KMeans"
5130            and distance_function != "SquaredEuclideanDistance"
5131        ):
5132            evaluation.message(self.get_name(), "kmsud")
5133            return
5134
5135        def df(i, j):
5136            return Expression(distance_function, i, j)
5137
5138        try:
5139            if method_string == "Agglomerate":
5140                clusters = self._agglomerate(mode, repr_p, dist_p, py_k, df, evaluation)
5141            elif method_string == "Optimize":
5142                clusters = optimize(
5143                    repr_p, py_k, _LazyDistances(df, dist_p, evaluation), mode, py_seed
5144                )
5145            elif method_string == "KMeans":
5146                clusters = self._kmeans(mode, repr_p, dist_p, py_k, py_seed, evaluation)
5147        except _IllegalDistance as e:
5148            evaluation.message(self.get_name(), "xnum", e.distance)
5149            return
5150        except _IllegalDataPoint:
5151            name_of_builtin = strip_context(self.get_name())
5152            evaluation.message(
5153                self.get_name(),
5154                "amtd",
5155                name_of_builtin,
5156                Expression(SymbolList, *dist_p),
5157            )
5158            return
5159
5160        if mode == "clusters":
5161            return Expression(
5162                SymbolList, *[Expression(SymbolList, *c) for c in clusters]
5163            )
5164        elif mode == "components":
5165            return Expression(SymbolList, *clusters)
5166        else:
5167            raise ValueError("illegal mode %s" % mode)
5168
5169    def _agglomerate(self, mode, repr_p, dist_p, py_k, df, evaluation):
5170        if mode == "clusters":
5171            clusters = agglomerate(
5172                repr_p, py_k, _PrecomputedDistances(df, dist_p, evaluation), mode
5173            )
5174        elif mode == "components":
5175            clusters = agglomerate(
5176                repr_p, py_k, _PrecomputedDistances(df, dist_p, evaluation), mode
5177            )
5178
5179        return clusters
5180
5181    def _kmeans(self, mode, repr_p, dist_p, py_k, py_seed, evaluation):
5182        items = []
5183
5184        def convert_scalars(p):
5185            for q in p:
5186                if not isinstance(q, (Real, Integer)):
5187                    raise _IllegalDataPoint
5188                mpq = q.to_mpmath()
5189                if mpq is None:
5190                    raise _IllegalDataPoint
5191                items.append(q)
5192                yield mpq
5193
5194        def convert_vectors(p):
5195            d = None
5196            for q in p:
5197                if q.get_head_name() != "System`List":
5198                    raise _IllegalDataPoint
5199                v = list(convert_scalars(q.leaves))
5200                if d is None:
5201                    d = len(v)
5202                elif len(v) != d:
5203                    raise _IllegalDataPoint
5204                yield v
5205
5206        if dist_p[0].is_numeric():
5207            numeric_p = [[x] for x in convert_scalars(dist_p)]
5208        else:
5209            numeric_p = list(convert_vectors(dist_p))
5210
5211        # compute epsilon similar to Real.__eq__, such that "numbers that differ in their last seven binary digits
5212        # are considered equal"
5213
5214        prec = min_prec(*items) or machine_precision
5215        eps = 0.5 ** (prec - 7)
5216
5217        return kmeans(numeric_p, repr_p, py_k, mode, py_seed, eps)
5218
5219
5220class FindClusters(_Cluster):
5221    """
5222    <dl>
5223    <dt>'FindClusters[$list$]'
5224        <dd>returns a list of clusters formed from the elements of $list$. The number of cluster is determined
5225        automatically.
5226    <dt>'FindClusters[$list$, $k$]'
5227        <dd>returns a list of $k$ clusters formed from the elements of $list$.
5228    </dl>
5229
5230    >> FindClusters[{1, 2, 20, 10, 11, 40, 19, 42}]
5231     = {{1, 2, 20, 10, 11, 19}, {40, 42}}
5232
5233    >> FindClusters[{25, 100, 17, 20}]
5234     = {{25, 17, 20}, {100}}
5235
5236    >> FindClusters[{3, 6, 1, 100, 20, 5, 25, 17, -10, 2}]
5237     = {{3, 6, 1, 5, -10, 2}, {100}, {20, 25, 17}}
5238
5239    >> FindClusters[{1, 2, 10, 11, 20, 21}]
5240     = {{1, 2}, {10, 11}, {20, 21}}
5241
5242    >> FindClusters[{1, 2, 10, 11, 20, 21}, 2]
5243     = {{1, 2, 10, 11}, {20, 21}}
5244
5245    >> FindClusters[{1 -> a, 2 -> b, 10 -> c}]
5246     = {{a, b}, {c}}
5247
5248    >> FindClusters[{1, 2, 5} -> {a, b, c}]
5249     = {{a, b}, {c}}
5250
5251    >> FindClusters[{1, 2, 3, 1, 2, 10, 100}, Method -> "Agglomerate"]
5252     = {{1, 2, 3, 1, 2, 10}, {100}}
5253
5254    >> FindClusters[{1, 2, 3, 10, 17, 18}, Method -> "Agglomerate"]
5255     = {{1, 2, 3}, {10}, {17, 18}}
5256
5257    >> FindClusters[{{1}, {5, 6}, {7}, {2, 4}}, DistanceFunction -> (Abs[Length[#1] - Length[#2]]&)]
5258     = {{{1}, {7}}, {{5, 6}, {2, 4}}}
5259
5260    >> FindClusters[{"meep", "heap", "deep", "weep", "sheep", "leap", "keep"}, 3]
5261     = {{meep, deep, weep, keep}, {heap, leap}, {sheep}}
5262
5263    FindClusters' automatic distance function detection supports scalars, numeric tensors, boolean vectors and
5264    strings.
5265
5266    The Method option must be either "Agglomerate" or "Optimize". If not specified, it defaults to "Optimize".
5267    Note that the Agglomerate and Optimize methods usually produce different clusterings.
5268
5269    The runtime of the Agglomerate method is quadratic in the number of clustered points n, builds the clustering
5270    from the bottom up, and is exact (no element of randomness). The Optimize method's runtime is linear in n,
5271    Optimize builds the clustering from top down, and uses random sampling.
5272    """
5273
5274    def apply(self, p, evaluation, options):
5275        "FindClusters[p_, OptionsPattern[%(name)s]]"
5276        return self._cluster(
5277            p,
5278            None,
5279            "clusters",
5280            evaluation,
5281            options,
5282            Expression("FindClusters", p, *options_to_rules(options)),
5283        )
5284
5285    def apply_manual_k(self, p, k, evaluation, options):
5286        "FindClusters[p_, k_Integer, OptionsPattern[%(name)s]]"
5287        return self._cluster(
5288            p,
5289            k,
5290            "clusters",
5291            evaluation,
5292            options,
5293            Expression("FindClusters", p, k, *options_to_rules(options)),
5294        )
5295
5296
5297class ClusteringComponents(_Cluster):
5298    """
5299    <dl>
5300    <dt>'ClusteringComponents[$list$]'
5301        <dd>forms clusters from $list$ and returns a list of cluster indices, in which each
5302        element shows the index of the cluster in which the corresponding element in $list$
5303        ended up.
5304    <dt>'ClusteringComponents[$list$, $k$]'
5305        <dd>forms $k$ clusters from $list$ and returns a list of cluster indices, in which
5306        each element shows the index of the cluster in which the corresponding element in
5307        $list$ ended up.
5308    </dl>
5309
5310    For more detailed documentation regarding options and behavior, see FindClusters[].
5311
5312    >> ClusteringComponents[{1, 2, 3, 1, 2, 10, 100}]
5313     = {1, 1, 1, 1, 1, 1, 2}
5314
5315    >> ClusteringComponents[{10, 100, 20}, Method -> "KMeans"]
5316     = {1, 0, 1}
5317    """
5318
5319    def apply(self, p, evaluation, options):
5320        "ClusteringComponents[p_, OptionsPattern[%(name)s]]"
5321        return self._cluster(
5322            p,
5323            None,
5324            "components",
5325            evaluation,
5326            options,
5327            Expression("ClusteringComponents", p, *options_to_rules(options)),
5328        )
5329
5330    def apply_manual_k(self, p, k, evaluation, options):
5331        "ClusteringComponents[p_, k_Integer, OptionsPattern[%(name)s]]"
5332        return self._cluster(
5333            p,
5334            k,
5335            "components",
5336            evaluation,
5337            options,
5338            Expression("ClusteringComponents", p, k, *options_to_rules(options)),
5339        )
5340
5341
5342class Nearest(Builtin):
5343    """
5344    <dl>
5345    <dt>'Nearest[$list$, $x$]'
5346        <dd>returns the one item in $list$ that is nearest to $x$.
5347    <dt>'Nearest[$list$, $x$, $n$]'
5348        <dd>returns the $n$ nearest items.
5349    <dt>'Nearest[$list$, $x$, {$n$, $r$}]'
5350        <dd>returns up to $n$ nearest items that are not farther from $x$ than $r$.
5351    <dt>'Nearest[{$p1$ -> $q1$, $p2$ -> $q2$, ...}, $x$]'
5352        <dd>returns $q1$, $q2$, ... but measures the distances using $p1$, $p2$, ...
5353    <dt>'Nearest[{$p1$, $p2$, ...} -> {$q1$, $q2$, ...}, $x$]'
5354        <dd>returns $q1$, $q2$, ... but measures the distances using $p1$, $p2$, ...
5355    </dl>
5356
5357    >> Nearest[{5, 2.5, 10, 11, 15, 8.5, 14}, 12]
5358     = {11}
5359
5360    Return all items within a distance of 5:
5361
5362    >> Nearest[{5, 2.5, 10, 11, 15, 8.5, 14}, 12, {All, 5}]
5363     = {11, 10, 14}
5364
5365    >> Nearest[{Blue -> "blue", White -> "white", Red -> "red", Green -> "green"}, {Orange, Gray}]
5366     = {{red}, {white}}
5367
5368    >> Nearest[{{0, 1}, {1, 2}, {2, 3}} -> {a, b, c}, {1.1, 2}]
5369     = {b}
5370    """
5371
5372    options = {
5373        "DistanceFunction": "Automatic",
5374        "Method": '"Scan"',
5375    }
5376
5377    messages = {
5378        "amtd": "`1` failed to pick a suitable distance function for `2`.",
5379        "list": "Expected a list or a rule with equally sized lists at position 1 in ``.",
5380        "nimp": "Method `1` is not implemented yet.",
5381    }
5382
5383    rules = {
5384        "Nearest[list_, pattern_]": "Nearest[list, pattern, 1]",
5385        "Nearest[pattern_][list_]": "Nearest[list, pattern]",
5386    }
5387
5388    def apply(self, items, pivot, limit, expression, evaluation, options):
5389        "Nearest[items_, pivot_, limit_, OptionsPattern[%(name)s]]"
5390
5391        method = self.get_option(options, "Method", evaluation)
5392        if not isinstance(method, String) or method.get_string_value() != "Scan":
5393            evaluation("Nearest", "nimp", method)
5394            return
5395
5396        dist_p, repr_p = _dist_repr(items)
5397
5398        if dist_p is None or len(dist_p) != len(repr_p):
5399            evaluation.message(self.get_name(), "list", expression)
5400            return
5401
5402        if limit.has_form("List", 2):
5403            up_to = limit.leaves[0]
5404            py_r = limit.leaves[1].to_mpmath()
5405        else:
5406            up_to = limit
5407            py_r = None
5408
5409        if isinstance(up_to, Integer):
5410            py_n = up_to.get_int_value()
5411        elif up_to.get_name() == "System`All":
5412            py_n = None
5413        else:
5414            return
5415
5416        if not dist_p or (py_n is not None and py_n < 1):
5417            return Expression(SymbolList)
5418
5419        multiple_x = False
5420
5421        distance_function_string, distance_function = self.get_option_string(
5422            options, "DistanceFunction", evaluation
5423        )
5424        if distance_function_string == "Automatic":
5425            from mathics.builtin.tensors import get_default_distance
5426
5427            distance_function = get_default_distance(dist_p)
5428            if distance_function is None:
5429                evaluation.message(
5430                    self.get_name(), "amtd", "Nearest", Expression(SymbolList, *dist_p)
5431                )
5432                return
5433
5434            if pivot.get_head_name() == "System`List":
5435                _, depth_x = walk_levels(pivot)
5436                _, depth_items = walk_levels(dist_p[0])
5437
5438                if depth_x > depth_items:
5439                    multiple_x = True
5440
5441        def nearest(x):
5442            calls = [Expression(distance_function, x, y) for y in dist_p]
5443            distances = Expression(SymbolList, *calls).evaluate(evaluation)
5444
5445            if not distances.has_form("List", len(dist_p)):
5446                raise ValueError()
5447
5448            py_distances = [
5449                (_to_real_distance(d), i) for i, d in enumerate(distances.leaves)
5450            ]
5451
5452            if py_r is not None:
5453                py_distances = [(d, i) for d, i in py_distances if d <= py_r]
5454
5455            def pick():
5456                if py_n is None:
5457                    candidates = sorted(py_distances)
5458                else:
5459                    candidates = heapq.nsmallest(py_n, py_distances)
5460
5461                for d, i in candidates:
5462                    yield repr_p[i]
5463
5464            return Expression(SymbolList, *list(pick()))
5465
5466        try:
5467            if not multiple_x:
5468                return nearest(pivot)
5469            else:
5470                return Expression(SymbolList, *[nearest(t) for t in pivot.leaves])
5471        except _IllegalDistance:
5472            return SymbolFailed
5473        except ValueError:
5474            return SymbolFailed
5475
5476
5477class Permutations(Builtin):
5478    """
5479    <dl>
5480    <dt>'Permutations[$list$]'
5481        <dd>gives all possible orderings of the items in $list$.
5482    <dt>'Permutations[$list$, $n$]'
5483        <dd>gives permutations up to length $n$.
5484    <dt>'Permutations[$list$, {$n$}]'
5485        <dd>gives permutations of length $n$.
5486    </dl>
5487
5488    >> Permutations[{y, 1, x}]
5489     = {{y, 1, x}, {y, x, 1}, {1, y, x}, {1, x, y}, {x, y, 1}, {x, 1, y}}
5490
5491    Elements are differentiated by their position in $list$, not their value.
5492
5493    >> Permutations[{a, b, b}]
5494     = {{a, b, b}, {a, b, b}, {b, a, b}, {b, b, a}, {b, a, b}, {b, b, a}}
5495
5496    >> Permutations[{1, 2, 3}, 2]
5497     = {{}, {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 1}, {2, 3}, {3, 1}, {3, 2}}
5498
5499    >> Permutations[{1, 2, 3}, {2}]
5500     = {{1, 2}, {1, 3}, {2, 1}, {2, 3}, {3, 1}, {3, 2}}
5501    """
5502
5503    messages = {
5504        "argt": "Permutation expects at least one argument.",
5505        "nninfseq": "The number specified at position 2 of `` must be a non-negative integer, All, or Infinity.",
5506    }
5507
5508    def apply_argt(self, evaluation):
5509        "Permutations[]"
5510        evaluation.message(self.get_name(), "argt")
5511
5512    def apply(self, l, evaluation):
5513        "Permutations[l_List]"
5514        return Expression(
5515            "List",
5516            *[
5517                Expression(SymbolList, *p)
5518                for p in permutations(l.leaves, len(l.leaves))
5519            ],
5520        )
5521
5522    def apply_n(self, l, n, evaluation):
5523        "Permutations[l_List, n_]"
5524
5525        rs = None
5526        if isinstance(n, Integer):
5527            py_n = min(n.get_int_value(), len(l.leaves))
5528        elif n.has_form("List", 1) and isinstance(n.leaves[0], Integer):
5529            py_n = n.leaves[0].get_int_value()
5530            rs = (py_n,)
5531        elif (
5532            n.has_form("DirectedInfinity", 1) and n.leaves[0].get_int_value() == 1
5533        ) or n.get_name() == "System`All":
5534            py_n = len(l.leaves)
5535        else:
5536            py_n = None
5537
5538        if py_n is None or py_n < 0:
5539            evaluation.message(
5540                self.get_name(), "nninfseq", Expression(self.get_name(), l, n)
5541            )
5542            return
5543
5544        if rs is None:
5545            rs = range(py_n + 1)
5546
5547        inner = structure("List", l, evaluation)
5548        outer = structure("List", inner, evaluation)
5549
5550        return outer([inner(p) for r in rs for p in permutations(l.leaves, r)])
5551
5552
5553class SubsetQ(Builtin):
5554    """
5555    <dl>
5556    <dt>'SubsetQ[$list1$, $list2$]'
5557        <dd>returns True if $list2$ is a subset of $list1$, and False otherwise.
5558    </dl>
5559
5560    >> SubsetQ[{1, 2, 3}, {3, 1}]
5561     = True
5562
5563    The empty list is a subset of every list:
5564    >> SubsetQ[{}, {}]
5565     = True
5566
5567    >> SubsetQ[{1, 2, 3}, {}]
5568     = True
5569
5570    Every list is a subset of itself:
5571    >> SubsetQ[{1, 2, 3}, {1, 2, 3}]
5572     = True
5573
5574    #> SubsetQ[{1, 2, 3}, {0, 1}]
5575     = False
5576
5577    #> SubsetQ[{1, 2, 3}, {1, 2, 3, 4}]
5578     = False
5579
5580    #> SubsetQ[{1, 2, 3}]
5581     : SubsetQ called with 1 argument; 2 arguments are expected.
5582     = SubsetQ[{1, 2, 3}]
5583
5584    #> SubsetQ[{1, 2, 3}, {1, 2}, {3}]
5585     : SubsetQ called with 3 arguments; 2 arguments are expected.
5586     = SubsetQ[{1, 2, 3}, {1, 2}, {3}]
5587
5588    #> SubsetQ[a + b + c, {1}]
5589     : Heads Plus and List at positions 1 and 2 are expected to be the same.
5590     = SubsetQ[a + b + c, {1}]
5591
5592    #> SubsetQ[{1, 2, 3}, n]
5593     : Nonatomic expression expected at position 2 in SubsetQ[{1, 2, 3}, n].
5594     = SubsetQ[{1, 2, 3}, n]
5595
5596    #> SubsetQ[f[a, b, c], f[a]]
5597     = True
5598    """
5599
5600    messages = {
5601        "argr": "SubsetQ called with 1 argument; 2 arguments are expected.",
5602        "argrx": "SubsetQ called with `1` arguments; 2 arguments are expected.",
5603        "heads": "Heads `1` and `2` at positions 1 and 2 are expected to be the same.",
5604        "normal": "Nonatomic expression expected at position `1` in `2`.",
5605    }
5606
5607    def apply(self, expr, subset, evaluation):
5608        "SubsetQ[expr_, subset___]"
5609
5610        if expr.is_atom():
5611            return evaluation.message(
5612                "SubsetQ", "normal", Integer(1), Expression("SubsetQ", expr, subset)
5613            )
5614
5615        subset = subset.get_sequence()
5616        if len(subset) > 1:
5617            return evaluation.message("SubsetQ", "argrx", Integer(len(subset) + 1))
5618        elif len(subset) == 0:
5619            return evaluation.message("SubsetQ", "argr")
5620
5621        subset = subset[0]
5622        if subset.is_atom():
5623            return evaluation.message(
5624                "SubsetQ", "normal", Integer(2), Expression("SubsetQ", expr, subset)
5625            )
5626        if expr.get_head_name() != subset.get_head_name():
5627            return evaluation.message(
5628                "SubsetQ", "heads", expr.get_head(), subset.get_head()
5629            )
5630
5631        if set(subset.leaves).issubset(set(expr.leaves)):
5632            return Symbol("True")
5633        else:
5634            return Symbol("False")
5635
5636
5637def delete_one(expr, pos):
5638    if expr.is_atom():
5639        raise PartDepthError(pos)
5640    leaves = expr.leaves
5641    if pos == 0:
5642        return Expression(Symbol("System`Sequence"), *leaves)
5643    l = len(leaves)
5644    truepos = pos
5645    if truepos < 0:
5646        truepos = l + truepos
5647    else:
5648        truepos = truepos - 1
5649    if truepos < 0 or truepos >= l:
5650        raise PartRangeError
5651    leaves = leaves[:truepos] + (Expression("System`Sequence"),) + leaves[truepos + 1 :]
5652    return Expression(expr.get_head(), *leaves)
5653
5654
5655def delete_rec(expr, pos):
5656    if len(pos) == 1:
5657        return delete_one(expr, pos[0])
5658    truepos = pos[0]
5659    if truepos == 0 or expr.is_atom():
5660        raise PartDepthError(pos[0])
5661    leaves = expr.leaves
5662    l = len(leaves)
5663    if truepos < 0:
5664        truepos = truepos + l
5665        if truepos < 0:
5666            raise PartRangeError
5667        newleaf = delete_rec(leaves[truepos], pos[1:])
5668        leaves = leaves[:truepos] + (newleaf,) + leaves[truepos + 1 :]
5669    else:
5670        if truepos > l:
5671            raise PartRangeError
5672        newleaf = delete_rec(leaves[truepos - 1], pos[1:])
5673        leaves = leaves[: truepos - 1] + (newleaf,) + leaves[truepos:]
5674    return Expression(expr.get_head(), *leaves)
5675
5676
5677class Delete(Builtin):
5678    """
5679    <dl>
5680    <dt>'Delete[$expr$, $i$]'
5681        <dd>deletes the element at position $i$ in $expr$. The position is counted from the end if $i$ is negative.
5682    <dt>'Delete[$expr$, {$m$, $n$, ...}]'
5683        <dd>deletes the element at position {$m$, $n$, ...}.
5684    <dt>'Delete[$expr$, {{$m1$, $n1$, ...}, {$m2$, $n2$, ...}, ...}]'
5685        <dd>deletes the elements at several positions.
5686    </dl>
5687
5688    Delete the element at position 3:
5689    >> Delete[{a, b, c, d}, 3]
5690     = {a, b, d}
5691
5692    Delete at position 2 from the end:
5693    >> Delete[{a, b, c, d}, -2]
5694     = {a, b, d}
5695
5696    Delete at positions 1 and 3:
5697    >> Delete[{a, b, c, d}, {{1}, {3}}]
5698     = {b, d}
5699
5700    Delete in a 2D array:
5701    >> Delete[{{a, b}, {c, d}}, {2, 1}]
5702     = {{a, b}, {d}}
5703
5704    Deleting the head of a whole expression gives a Sequence object:
5705    >> Delete[{a, b, c}, 0]
5706     = Sequence[a, b, c]
5707
5708    Delete in an expression with any head:
5709    >> Delete[f[a, b, c, d], 3]
5710     = f[a, b, d]
5711
5712    Delete a head to splice in its arguments:
5713    >> Delete[f[a, b, u + v, c], {3, 0}]
5714     = f[a, b, u, v, c]
5715
5716    >> Delete[{a, b, c}, 0]
5717     = Sequence[a, b, c]
5718
5719    #> Delete[1 + x ^ (a + b + c), {2, 2, 3}]
5720     = 1 + x ^ (a + b)
5721
5722    #> Delete[f[a, g[b, c], d], {{2}, {2, 1}}]
5723     = f[a, d]
5724
5725    #> Delete[f[a, g[b, c], d], m + n]
5726     : The expression m + n cannot be used as a part specification. Use Key[m + n] instead.
5727     = Delete[f[a, g[b, c], d], m + n]
5728
5729    Delete without the position:
5730    >> Delete[{a, b, c, d}]
5731     : Delete called with 1 argument; 2 arguments are expected.
5732     = Delete[{a, b, c, d}]
5733
5734    Delete with many arguments:
5735    >> Delete[{a, b, c, d}, 1, 2]
5736     : Delete called with 3 arguments; 2 arguments are expected.
5737     = Delete[{a, b, c, d}, 1, 2]
5738
5739    Delete the element out of range:
5740    >> Delete[{a, b, c, d}, 5]
5741     : Part {5} of {a, b, c, d} does not exist.
5742     = Delete[{a, b, c, d}, 5]
5743
5744    #> Delete[{a, b, c, d}, {1, 2}]
5745     : Part 2 of {a, b, c, d} does not exist.
5746     = Delete[{a, b, c, d}, {1, 2}]
5747
5748    Delete the position not integer:
5749    >> Delete[{a, b, c, d}, {1, n}]
5750     : Position specification n in {a, b, c, d} is not a machine-sized integer or a list of machine-sized integers.
5751     = Delete[{a, b, c, d}, {1, n}]
5752
5753    #> Delete[{a, b, c, d}, {{1}, n}]
5754     : Position specification {n, {1}} in {a, b, c, d} is not a machine-sized integer or a list of machine-sized integers.
5755     = Delete[{a, b, c, d}, {{1}, n}]
5756
5757    #> Delete[{a, b, c, d}, {{1}, {n}}]
5758     : Position specification n in {a, b, c, d} is not a machine-sized integer or a list of machine-sized integers.
5759     = Delete[{a, b, c, d}, {{1}, {n}}]
5760    """
5761
5762    messages = {
5763        "argr": "Delete called with 1 argument; 2 arguments are expected.",
5764        "argt": "Delete called with `1` arguments; 2 arguments are expected.",
5765        "psl": "Position specification `1` in `2` is not a machine-sized integer or a list of machine-sized integers.",
5766        "pkspec": "The expression `1` cannot be used as a part specification. Use `2` instead.",
5767    }
5768
5769    def apply_one(self, expr, position, evaluation):
5770        "Delete[expr_, position_Integer]"
5771        pos = position.get_int_value()
5772        try:
5773            return delete_one(expr, pos)
5774        except PartRangeError:
5775            evaluation.message("Part", "partw", Expression(SymbolList, pos), expr)
5776
5777    def apply(self, expr, positions, evaluation):
5778        "Delete[expr_, positions___]"
5779        positions = positions.get_sequence()
5780        if len(positions) > 1:
5781            return evaluation.message("Delete", "argt", Integer(len(positions) + 1))
5782        elif len(positions) == 0:
5783            return evaluation.message("Delete", "argr")
5784
5785        positions = positions[0]
5786        if not positions.has_form("List", None):
5787            return evaluation.message(
5788                "Delete", "pkspec", positions, Expression("Key", positions)
5789            )
5790
5791        # Create new python list of the positions and sort it
5792        positions = (
5793            [l for l in positions.leaves]
5794            if positions.leaves[0].has_form("List", None)
5795            else [positions]
5796        )
5797        positions.sort(key=lambda e: e.get_sort_key(pattern_sort=True))
5798        leaves = expr.leaves
5799        newexpr = expr
5800        for position in positions:
5801            pos = [p.get_int_value() for p in position.get_leaves()]
5802            if None in pos:
5803                return evaluation.message(
5804                    "Delete", "psl", position.leaves[pos.index(None)], expr
5805                )
5806            if len(pos) == 0:
5807                return evaluation.message(
5808                    "Delete", "psl", Expression(SymbolList, *positions), expr
5809                )
5810            try:
5811                newexpr = delete_rec(newexpr, pos)
5812            except PartDepthError as exc:
5813                return evaluation.message("Part", "partw", Integer(exc.index), expr)
5814            except PartError:
5815                return evaluation.message(
5816                    "Part", "partw", Expression(SymbolList, *pos), expr
5817                )
5818        return newexpr
5819
5820
5821class Association(Builtin):
5822    """
5823    <dl>
5824    <dt>'Association[$key1$ -> $val1$, $key2$ -> $val2$, ...]'
5825    <dt>'<|$key1$ -> $val1$, $key2$ -> $val2$, ...|>'
5826        <dd> represents an association between keys and values.
5827    </dl>
5828
5829    'Association' is the head of associations:
5830    >> Head[<|a -> x, b -> y, c -> z|>]
5831     = Association
5832
5833    >> <|a -> x, b -> y|>
5834     = <|a -> x, b -> y|>
5835
5836    >> Association[{a -> x, b -> y}]
5837     = <|a -> x, b -> y|>
5838
5839    Associations can be nested:
5840    >> <|a -> x, b -> y, <|a -> z, d -> t|>|>
5841     = <|a -> z, b -> y, d -> t|>
5842
5843    #> <|a -> x, b -> y, c -> <|d -> t|>|>
5844     = <|a -> x, b -> y, c -> <|d -> t|>|>
5845    #> %["s"]
5846     = Missing[KeyAbsent, s]
5847
5848    #> <|a -> x, b + c -> y, {<|{}|>, a -> {z}}|>
5849     = <|a -> {z}, b + c -> y|>
5850    #> %[a]
5851     = {z}
5852
5853    #> <|"x" -> 1, {y} -> 1|>
5854     = <|x -> 1, {y} -> 1|>
5855    #> %["x"]
5856     = 1
5857
5858    #> <|<|a -> v|> -> x, <|b -> y, a -> <|c -> z|>, {}, <||>|>, {d}|>[c]
5859     =  Association[Association[a -> v] -> x, Association[b -> y, a -> Association[c -> z], {}, Association[]], {d}][c]
5860
5861    #> <|<|a -> v|> -> x, <|b -> y, a -> <|c -> z|>, {d}|>, {}, <||>|>[a]
5862     = Association[Association[a -> v] -> x, Association[b -> y, a -> Association[c -> z], {d}], {}, Association[]][a]
5863
5864    #> <|<|a -> v|> -> x, <|b -> y, a -> <|c -> z, {d}|>, {}, <||>|>, {}, <||>|>
5865     = <|<|a -> v|> -> x, b -> y, a -> Association[c -> z, {d}]|>
5866    #> %[a]
5867     = Association[c -> z, {d}]
5868
5869    #> <|a -> x, b -> y, c -> <|d -> t|>|> // ToBoxes
5870     = RowBox[{<|, RowBox[{RowBox[{a, ->, x}], ,, RowBox[{b, ->, y}], ,, RowBox[{c, ->, RowBox[{<|, RowBox[{d, ->, t}], |>}]}]}], |>}]
5871
5872    #> Association[a -> x, b -> y, c -> Association[d -> t, Association[e -> u]]] // ToBoxes
5873     = RowBox[{<|, RowBox[{RowBox[{a, ->, x}], ,, RowBox[{b, ->, y}], ,, RowBox[{c, ->, RowBox[{<|, RowBox[{RowBox[{d, ->, t}], ,, RowBox[{e, ->, u}]}], |>}]}]}], |>}]
5874    """
5875
5876    error_idx = 0
5877
5878    attributes = (
5879        "HoldAllComplete",
5880        "Protected",
5881    )
5882
5883    def apply_makeboxes(self, rules, f, evaluation):
5884        """MakeBoxes[<|rules___|>,
5885        f:StandardForm|TraditionalForm|OutputForm|InputForm]"""
5886
5887        def validate(exprs):
5888            for expr in exprs:
5889                if expr.has_form(("Rule", "RuleDelayed"), 2):
5890                    pass
5891                elif expr.has_form("List", None) or expr.has_form("Association", None):
5892                    if validate(expr.leaves) is not True:
5893                        return False
5894                else:
5895                    return False
5896            return True
5897
5898        rules = rules.get_sequence()
5899        if self.error_idx == 0 and validate(rules) is True:
5900            expr = Expression(
5901                "RowBox", Expression(SymbolList, *list_boxes(rules, f, "<|", "|>"))
5902            )
5903        else:
5904            self.error_idx += 1
5905            symbol = Expression(SymbolMakeBoxes, SymbolAssociation, f)
5906            expr = Expression(
5907                "RowBox",
5908                Expression(SymbolList, symbol, *list_boxes(rules, f, "[", "]")),
5909            )
5910
5911        expr = expr.evaluate(evaluation)
5912        if self.error_idx > 0:
5913            self.error_idx -= 1
5914        return expr
5915
5916    def apply(self, rules, evaluation):
5917        "Association[rules__]"
5918
5919        def make_flatten(exprs, dic={}, keys=[]):
5920            for expr in exprs:
5921                if expr.has_form(("Rule", "RuleDelayed"), 2):
5922                    key = expr.leaves[0].evaluate(evaluation)
5923                    value = expr.leaves[1].evaluate(evaluation)
5924                    dic[key] = Expression(expr.get_head(), key, value)
5925                    if key not in keys:
5926                        keys.append(key)
5927                elif expr.has_form("List", None) or expr.has_form("Association", None):
5928                    make_flatten(expr.leaves, dic, keys)
5929                else:
5930                    raise
5931            return [dic[key] for key in keys]
5932
5933        try:
5934            return Expression(SymbolAssociation, *make_flatten(rules.get_sequence()))
5935        except:
5936            return None
5937
5938    def apply_key(self, rules, key, evaluation):
5939        "Association[rules__][key_]"
5940
5941        def find_key(exprs, dic={}):
5942            for expr in exprs:
5943                if expr.has_form(("Rule", "RuleDelayed"), 2):
5944                    if expr.leaves[0] == key:
5945                        dic[key] = expr.leaves[1]
5946                elif expr.has_form("List", None) or expr.has_form("Association", None):
5947                    find_key(expr.leaves)
5948                else:
5949                    raise
5950            return dic
5951
5952        try:
5953            result = find_key(rules.get_sequence())
5954        except:
5955            return None
5956
5957        return (
5958            result[key] if result else Expression("Missing", Symbol("KeyAbsent"), key)
5959        )
5960
5961
5962class AssociationQ(Test):
5963    """
5964    <dl>
5965    <dt>'AssociationQ[$expr$]'
5966        <dd>return True if $expr$ is a valid Association object, and False otherwise.
5967    </dl>
5968
5969    >> AssociationQ[<|a -> 1, b :> 2|>]
5970     = True
5971
5972    >> AssociationQ[<|a, b|>]
5973     = False
5974    """
5975
5976    def test(self, expr):
5977        def validate(leaves):
5978            for leaf in leaves:
5979                if leaf.has_form(("Rule", "RuleDelayed"), 2):
5980                    pass
5981                elif leaf.has_form("List", None) or leaf.has_form("Association", None):
5982                    if validate(leaf.leaves) is not True:
5983                        return False
5984                else:
5985                    return False
5986            return True
5987
5988        return expr.get_head_name() == "System`Association" and validate(expr.leaves)
5989
5990
5991class Keys(Builtin):
5992    """
5993    <dl>
5994    <dt>'Keys[<|$key1$ -> $val1$, $key2$ -> $val2$, ...|>]'
5995        <dd>return a list of the keys $keyi$ in an association.
5996    <dt>'Keys[{$key1$ -> $val1$, $key2$ -> $val2$, ...}]'
5997        <dd>return a list of the $keyi$ in a list of rules.
5998    </dl>
5999
6000    >> Keys[<|a -> x, b -> y|>]
6001     = {a, b}
6002
6003    >> Keys[{a -> x, b -> y}]
6004     = {a, b}
6005
6006    Keys automatically threads over lists:
6007    >> Keys[{<|a -> x, b -> y|>, {w -> z, {}}}]
6008     = {{a, b}, {w, {}}}
6009
6010    Keys are listed in the order of their appearance:
6011    >> Keys[{c -> z, b -> y, a -> x}]
6012     = {c, b, a}
6013
6014    #> Keys[a -> x]
6015     = a
6016
6017    #> Keys[{a -> x, a -> y, {a -> z, <|b -> t|>, <||>, {}}}]
6018     = {a, a, {a, {b}, {}, {}}}
6019
6020    #> Keys[{a -> x, a -> y, <|a -> z, {b -> t}, <||>, {}|>}]
6021     = {a, a, {a, b}}
6022
6023    #> Keys[<|a -> x, a -> y, <|a -> z, <|b -> t|>, <||>, {}|>|>]
6024     = {a, b}
6025
6026    #> Keys[<|a -> x, a -> y, {a -> z, {b -> t}, <||>, {}}|>]
6027     = {a, b}
6028
6029    #> Keys[<|a -> x, <|a -> y, b|>|>]
6030     : The argument Association[a -> x, Association[a -> y, b]] is not a valid Association or a list of rules.
6031     = Keys[Association[a -> x, Association[a -> y, b]]]
6032
6033    #> Keys[<|a -> x, {a -> y, b}|>]
6034     : The argument Association[a -> x, {a -> y, b}] is not a valid Association or a list of rules.
6035     = Keys[Association[a -> x, {a -> y, b}]]
6036
6037    #> Keys[{a -> x, <|a -> y, b|>}]
6038     : The argument Association[a -> y, b] is not a valid Association or a list of rules.
6039     = Keys[{a -> x, Association[a -> y, b]}]
6040
6041    #> Keys[{a -> x, {a -> y, b}}]
6042     : The argument b is not a valid Association or a list of rules.
6043     = Keys[{a -> x, {a -> y, b}}]
6044
6045    #> Keys[a -> x, b -> y]
6046     : Keys called with 2 arguments; 1 argument is expected.
6047     = Keys[a -> x, b -> y]
6048    """
6049
6050    attributes = ("Protected",)
6051
6052    messages = {
6053        "argx": "Keys called with `1` arguments; 1 argument is expected.",
6054        "invrl": "The argument `1` is not a valid Association or a list of rules.",
6055    }
6056
6057    def apply(self, rules, evaluation):
6058        "Keys[rules___]"
6059
6060        def get_keys(expr):
6061            if expr.has_form(("Rule", "RuleDelayed"), 2):
6062                return expr.leaves[0]
6063            elif expr.has_form("List", None) or (
6064                expr.has_form("Association", None)
6065                and AssociationQ(expr).evaluate(evaluation) == Symbol("True")
6066            ):
6067                return Expression(SymbolList, *[get_keys(leaf) for leaf in expr.leaves])
6068            else:
6069                evaluation.message("Keys", "invrl", expr)
6070                raise
6071
6072        rules = rules.get_sequence()
6073        if len(rules) != 1:
6074            return evaluation.message("Keys", "argx", Integer(len(rules)))
6075
6076        try:
6077            return get_keys(rules[0])
6078        except:
6079            return None
6080
6081
6082class Values(Builtin):
6083    """
6084    <dl>
6085    <dt>'Values[<|$key1$ -> $val1$, $key2$ -> $val2$, ...|>]'
6086        <dd>return a list of the values $vali$ in an association.
6087    <dt>'Values[{$key1$ -> $val1$, $key2$ -> $val2$, ...}]'
6088        <dd>return a list of the $vali$ in a list of rules.
6089    </dl>
6090
6091    >> Values[<|a -> x, b -> y|>]
6092     = {x, y}
6093
6094    >> Values[{a -> x, b -> y}]
6095     = {x, y}
6096
6097    Values automatically threads over lists:
6098    >> Values[{<|a -> x, b -> y|>, {c -> z, {}}}]
6099     = {{x, y}, {z, {}}}
6100
6101    Values are listed in the order of their appearance:
6102    >> Values[{c -> z, b -> y, a -> x}]
6103     = {z, y, x}
6104
6105    #> Values[a -> x]
6106     = x
6107
6108    #> Values[{a -> x, a -> y, {a -> z, <|b -> t|>, <||>, {}}}]
6109     = {x, y, {z, {t}, {}, {}}}
6110
6111    #> Values[{a -> x, a -> y, <|a -> z, {b -> t}, <||>, {}|>}]
6112     = {x, y, {z, t}}
6113
6114    #> Values[<|a -> x, a -> y, <|a -> z, <|b -> t|>, <||>, {}|>|>]
6115     = {z, t}
6116
6117    #> Values[<|a -> x, a -> y, {a -> z, {b -> t}, <||>, {}}|>]
6118     = {z, t}
6119
6120    #> Values[<|a -> x, <|a -> y, b|>|>]
6121     : The argument Association[a -> x, Association[a -> y, b]] is not a valid Association or a list of rules.
6122     = Values[Association[a -> x, Association[a -> y, b]]]
6123
6124    #> Values[<|a -> x, {a -> y, b}|>]
6125     : The argument Association[a -> x, {a -> y, b}] is not a valid Association or a list of rules.
6126     = Values[Association[a -> x, {a -> y, b}]]
6127
6128    #> Values[{a -> x, <|a -> y, b|>}]
6129     : The argument {a -> x, Association[a -> y, b]} is not a valid Association or a list of rules.
6130     = Values[{a -> x, Association[a -> y, b]}]
6131
6132    #> Values[{a -> x, {a -> y, b}}]
6133     : The argument {a -> x, {a -> y, b}} is not a valid Association or a list of rules.
6134     = Values[{a -> x, {a -> y, b}}]
6135
6136    #> Values[a -> x, b -> y]
6137     : Values called with 2 arguments; 1 argument is expected.
6138     = Values[a -> x, b -> y]
6139    """
6140
6141    attributes = ("Protected",)
6142
6143    messages = {
6144        "argx": "Values called with `1` arguments; 1 argument is expected.",
6145        "invrl": "The argument `1` is not a valid Association or a list of rules.",
6146    }
6147
6148    def apply(self, rules, evaluation):
6149        "Values[rules___]"
6150
6151        def get_values(expr):
6152            if expr.has_form(("Rule", "RuleDelayed"), 2):
6153                return expr.leaves[1]
6154            elif expr.has_form("List", None) or (
6155                expr.has_form("Association", None)
6156                and AssociationQ(expr).evaluate(evaluation) == Symbol("True")
6157            ):
6158                return Expression(
6159                    SymbolList, *[get_values(leaf) for leaf in expr.leaves]
6160                )
6161            else:
6162                raise
6163
6164        rules = rules.get_sequence()
6165        if len(rules) != 1:
6166            return evaluation.message("Values", "argx", Integer(len(rules)))
6167
6168        try:
6169            return get_values(rules[0])
6170        except:
6171            return evaluation.message("Values", "invrl", rules[0])
6172
6173
6174class ContainsOnly(Builtin):
6175    """
6176    <dl>
6177    <dt>'ContainsOnly[$list1$, $list2$]'
6178        <dd>yields True if $list1$ contains only elements that appear in $list2$.
6179    </dl>
6180
6181    >> ContainsOnly[{b, a, a}, {a, b, c}]
6182     = True
6183
6184    The first list contains elements not present in the second list:
6185    >> ContainsOnly[{b, a, d}, {a, b, c}]
6186     = False
6187
6188    >> ContainsOnly[{}, {a, b, c}]
6189     = True
6190
6191    #> ContainsOnly[1, {1, 2, 3}]
6192     : List or association expected instead of 1.
6193     = ContainsOnly[1, {1, 2, 3}]
6194
6195    #> ContainsOnly[{1, 2, 3}, 4]
6196     : List or association expected instead of 4.
6197     = ContainsOnly[{1, 2, 3}, 4]
6198
6199    Use Equal as the comparison function to have numerical tolerance:
6200    >> ContainsOnly[{a, 1.0}, {1, a, b}, {SameTest -> Equal}]
6201     = True
6202
6203    #> ContainsOnly[{c, a}, {a, b, c}, IgnoreCase -> True]
6204     : Unknown option IgnoreCase -> True in ContainsOnly.
6205     : Unknown option IgnoreCase in .
6206     = True
6207    """
6208
6209    attributes = ("ReadProtected",)
6210
6211    messages = {
6212        "lsa": "List or association expected instead of `1`.",
6213        "nodef": "Unknown option `1` for ContainsOnly.",
6214        "optx": "Unknown option `1` in `2`.",
6215    }
6216
6217    options = {
6218        "SameTest": "SameQ",
6219    }
6220
6221    def check_options(self, expr, evaluation, options):
6222        for key in options:
6223            if key != "System`SameTest":
6224                if expr is None:
6225                    evaluation.message("ContainsOnly", "optx", Symbol(key))
6226                else:
6227                    return evaluation.message("ContainsOnly", "optx", Symbol(key), expr)
6228        return None
6229
6230    def apply(self, list1, list2, evaluation, options={}):
6231        "ContainsOnly[list1_?ListQ, list2_?ListQ, OptionsPattern[ContainsOnly]]"
6232
6233        same_test = self.get_option(options, "SameTest", evaluation)
6234
6235        def sameQ(a, b) -> bool:
6236            """Mathics SameQ"""
6237            result = Expression(same_test, a, b).evaluate(evaluation)
6238            return result.is_true()
6239
6240        self.check_options(None, evaluation, options)
6241        for a in list1.leaves:
6242            if not any(sameQ(a, b) for b in list2.leaves):
6243                return Symbol("False")
6244        return Symbol("True")
6245
6246    def apply_msg(self, e1, e2, evaluation, options={}):
6247        "ContainsOnly[e1_, e2_, OptionsPattern[ContainsOnly]]"
6248
6249        opts = (
6250            options_to_rules(options)
6251            if len(options) <= 1
6252            else [Expression(SymbolList, *options_to_rules(options))]
6253        )
6254        expr = Expression("ContainsOnly", e1, e2, *opts)
6255
6256        if not isinstance(e1, Symbol) and not e1.has_form("List", None):
6257            evaluation.message("ContainsOnly", "lsa", e1)
6258            return self.check_options(expr, evaluation, options)
6259
6260        if not isinstance(e2, Symbol) and not e2.has_form("List", None):
6261            evaluation.message("ContainsOnly", "lsa", e2)
6262            return self.check_options(expr, evaluation, options)
6263
6264        return self.check_options(expr, evaluation, options)
6265
6266
6267## From backports in CellsToTeX. This functions provides compatibility to WMA 10.
6268##  TODO:
6269##  * Add doctests
6270##  * Translate to python the more complex rules
6271##  * Complete the support.
6272
6273
6274class Key(Builtin):
6275    """
6276    <dl>
6277    <dt>Key[$key$]
6278        <dd> represents a key used to access a value in an association.
6279    <dt>Key[$key$][$assoc$]
6280        <dd>
6281    </dl>
6282    """
6283
6284    rules = {
6285        "Key[key_][assoc_Association]": "assoc[key]",
6286    }
6287
6288
6289class Lookup(Builtin):
6290    """
6291    <dl>
6292    <dt>Lookup[$assoc$, $key$]
6293        <dd> looks up the value associated with $key$ in the association $assoc$, or Missing[$KeyAbsent$].
6294    </dl>
6295    """
6296
6297    attributes = "HoldAllComplete"
6298    rules = {
6299        "Lookup[assoc_?AssociationQ, key_, default_]": "FirstCase[assoc, _[Verbatim[key], val_] :> val, default]",
6300        "Lookup[assoc_?AssociationQ, key_]": 'Lookup[assoc, key, Missing["KeyAbsent", key]]',
6301    }
6302
6303
6304class Failure(Builtin):
6305    """
6306    <dl>
6307    <dt>Failure[$tag$, $assoc$]
6308        <dd> represents a failure of a type indicated by $tag$, with details given by the association $assoc$.
6309    </dl>
6310    """
6311
6312    pass
6313
6314
6315#    rules = {'Failure /: MakeBoxes[Failure[tag_, assoc_Association], StandardForm]' :
6316# 		'With[{msg = assoc["MessageTemplate"], msgParam = assoc["MessageParameters"], type = assoc["Type"]}, ToBoxes @ Interpretation["Failure" @ Panel @ Grid[{{Style["\[WarningSign]", "Message", FontSize -> 35], Style["Message:", FontColor->GrayLevel[0.5]], ToString[StringForm[msg, Sequence @@ msgParam], StandardForm]}, {SpanFromAbove, Style["Tag:", FontColor->GrayLevel[0.5]], ToString[tag, StandardForm]},{SpanFromAbove,Style["Type:", FontColor->GrayLevel[0.5]],ToString[type, StandardForm]}},Alignment -> {Left, Top}], Failure[tag, assoc]] /; msg =!= Missing["KeyAbsent", "MessageTemplate"] && msgParam =!= Missing["KeyAbsent", "MessageParameters"] && msgParam =!= Missing["KeyAbsent", "Type"]]',
6317#     }
6318
6319
6320class FirstCase(Builtin):
6321    """
6322    <dl>
6323    <dt> FirstCase[{$e1$, $e2$, ...}, $pattern$]
6324        <dd>gives the first $ei$ to match $pattern$, or $Missing[\"NotFound\"]$ if none matching pattern is found.
6325
6326    <dt> FirstCase[{$e1$,$e2$, ...}, $pattern$ -> $rhs$]
6327        <dd> gives the value of $rhs$ corresponding to the first $ei$ to match pattern.
6328    <dt> FirstCase[$expr$, $pattern$, $default$]
6329         <dd> gives $default$ if no element matching $pattern$ is found.
6330
6331    <dt>FirstCase[$expr$, $pattern$, $default$, $levelspec$] \
6332         <dd>finds only objects that appear on levels specified by $levelspec$.
6333
6334    <dt>FirstCase[$pattern$]
6335        <dd>represents an operator form of FirstCase that can be applied to an expression.
6336    </dl>
6337
6338
6339    """
6340
6341    attributes = "HoldRest"
6342    options = Cases.options
6343    rules = {
6344        'FirstCase[expr_, pattOrRule_, Shortest[default_:Missing["NotFound"], 1],Shortest[levelspec_:{1}, 2], opts:OptionsPattern[]]': "Replace[Cases[expr, pattOrRule, levelspec, 1, opts],{{} :> default, {match_} :> match}]",
6345        "FirstCase[pattOrRule_][expr_]": "FirstCase[expr, pattOrRule]",
6346    }
6347