# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %% [markdown]
# # Effect of Heuristic on Search Efficiency
#
# This project will show you how to solve planning problems -- Think logistics and scheduling -- using search and logic to find tractable solutions.  Classic planning problems can solve a problem where the rules are known and clear-cut.  In other words, they are _deterministic_.  Where to effects of an action are known to all, in AI, this is known as being _observable_.  In controlled environments, where change only occurs through the actions taken by a planning agent, these systems work well.
#
# Acting in these worlds can be modeled logically.  Actions can take place by first meeting all their preconditions.  And once completed, its effects known.  Imagine you are in charge of logistics at Fed Ex and need to find solutions for getting packages from place to place.  How could that be done?  To model the action of flying a package, we can use a [Planning Domain Definition Language](https://planning.wiki/) (PDDL).
#
# ```
# Action(Fly, (plane, from, to)
#     Precondition: At(plane, from) and Plane(from) and Airport(from) and Airport(to)
#     Effect: not At(plane, from) and At(plane, to)
# )
# ```
#
# Propositional logic describes each action briefly.   Just add actions to represent new problems or new business domains.  Knowing that we have a starting state and goal state, a set of actions and costs associated with each action, allows us to apply search algorithms to find solutions.
#

# %%
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from dataclasses import make_dataclass

sns.set_theme(style="whitegrid")

# %%
Record = make_dataclass(
    "Record",
    [
        ("Problem", int),
        ("Alg_Num", int),
        ("Alg_Name", str),
        ("Heuristic", str),
        ("Actions", int),
        ("Expansions", int),
        ("Goal_Tests", int),
        ("New_Nodes", int),
        ("Length", int),
        ("Time", float),
    ],
)

# %% [markdown]
# ## Experimental Results
#
# I ran pypy to generate my experiments results using my mac desktop computer.

# %%
dataset = [
    Record(1, 1, "BFS", "", 20, 43, 56, 178, 6, 0.0284),
    Record(1, 2, "DFS", "", 20, 21, 22, 84, 20, 0.0073),
    Record(2, 1, "BFS", "", 72, 3343, 4609, 30503, 9, 0.3459),
    Record(2, 2, "DFS", "", 72, 624, 625, 5602, 619, 0.4810),
    Record(4, 4, "Greedy", "h_unmet_goals", 104, 29, 31, 280, 18, 0.0612),
    Record(
        4, 11, "A*", "h_pg_setlevel", 104, 79863, 79865, 768641, 14, 30733.3346,
    ),
    Record(
        3, 11, "A*", "h_pg_setlevel", 88, 12872, 12874, 115220, 12, 3046.3672
    ),
    Record(1, 3, "UCS", "", 20, 60, 62, 240, 6, 0.0379),
    Record(1, 4, "Greedy", "h_unmet_goals", 20, 7, 9, 29, 6, 0.0032),
    Record(1, 5, "Greedy", "h_pg_levelsum", 20, 6, 8, 28, 6, 0.4590),
    Record(1, 6, "Greedy", "h_pg_maxlevel", 20, 6, 8, 24, 6, 0.1218),
    Record(1, 7, "Greedy", "h_pg_setlevel", 20, 13, 15, 53, 6, 1.2660),
    Record(1, 8, "A*", "h_unmet_goals", 20, 50, 52, 206, 6, 0.0216),
    Record(1, 9, "A*", "h_pg_levelsum", 20, 28, 30, 122, 6, 0.2154),
    Record(1, 10, "A*", "h_pg_maxlevel", 20, 43, 45, 180, 6, 0.1616),
    Record(1, 11, "A*", "h_pg_setlevel", 20, 46, 48, 192, 6, 0.7168),
    Record(2, 3, "UCS", "", 72, 5154, 5156, 46618, 9, 0.7450),
    Record(2, 4, "Greedy", "h_unmet_goals", 72, 17, 19, 170, 9, 0.0205),
    Record(2, 5, "Greedy", "h_pg_levelsum", 72, 9, 11, 86, 9, 0.4228),
    Record(2, 6, "Greedy", "h_pg_maxlevel", 72, 27, 29, 249, 9, 0.6970),
    Record(2, 7, "Greedy", "h_pg_setlevel", 72, 304, 306, 2846, 10, 54.8001),
    Record(2, 8, "A*", "h_unmet_goals", 72, 2467, 2469, 22522, 9, 0.7026),
    Record(2, 9, "A*", "h_pg_levelsum", 72, 357, 359, 3426, 9, 10.3350),
    Record(2, 10, "A*", "h_pg_maxlevel", 72, 2887, 2889, 26594, 9, 59.0657),
    Record(2, 11, "A*", "h_pg_setlevel", 72, 2879, 2881, 26622, 9, 577.1144),
    Record(3, 4, "Greedy", "h_unmet_goals", 88, 25, 27, 230, 15, 0.0519),
    Record(3, 6, "Greedy", "h_pg_maxlevel", 88, 21, 23, 195, 13, 1.8522),
    Record(4, 4, "Greedy", "h_unmet_goals", 104, 29, 31, 280, 18, 0.0615),
    Record(4, 6, "Greedy", "h_pg_maxlevel", 104, 56, 58, 580, 17, 3.6040),
    Record(3, 8, "A*", "h_unmet_goals", 88, 7388, 7390, 65711, 12, 1.2993),
    Record(3, 10, "A*", "h_pg_maxlevel", 88, 9580, 9582, 86312, 12, 331.0158),
    Record(4, 8, "A*", "h_unmet_goals", 104, 34330, 34332, 328509, 14, 4.0136),
    Record(
        4, 10, "A*", "h_pg_maxlevel", 104, 62077, 62079, 599376, 14, 3268.4160
    ),
    Record(3, 2, "DFS", "", 88, 408, 409, 3364, 392, 0.2780),
    Record(4, 2, "DFS", "", 104, 25174, 25175, 228849, 24132, 771.7955),
]
df = pd.DataFrame(dataset)

# %% [markdown]
# ---
#
# ## Supporting Functions

# %%
def create_catplot(df, y_var, title, is_log=True):
    """Creates a categorical plot using Seaborn's catplot.

    Parameters:
        df: Dataframe contianing data to be plotted
        y_var: Name of column in dataframe used for y-axis values
        title: Title for the plot
    """

    g = sns.catplot(
        data=df,
        kind="bar",
        x="Problem",
        y=y_var,
        hue="Alg_Title",
        height=7,
        aspect=1.0,
    )
    g.despine(left=True, bottom=True)
    if is_log:
        g.set(yscale="log")
    g.ax.xaxis.grid(True, linewidth=0.25)
    g.ax.yaxis.grid(True, "minor", linewidth=0.25)
    g.despine(left=True, bottom=False)
    plt.title(title)

# %%
def create_relplot(df, y_var, title):
    """Creates a relational plot using Seaborn's catplot.

    Parameters:
        df: Dataframe contianing data to be plotted
        y_var: Name of column in dataframe used for y-axis values
        title: Title for the plot
    """

    g = sns.relplot(
        data=df,
        x="Actions",
        y=y_var,
        hue="Alg_Title",
        size="Problem",
        height=7,
        aspect=1.0,
        sizes=(40, 120),
    )
    g.set(yscale="log")
    g.ax.xaxis.grid(True, linewidth=0.25)
    g.ax.yaxis.grid(True, "minor", linewidth=0.25)
    g.despine(left=True, bottom=False)
    plt.title(title)

# %%
df["Alg_Title"] = (
    df.Alg_Num.astype(str) + " " + df["Alg_Name"] + " " + df["Heuristic"]
)
df = df.sort_values(by=["Problem", "Alg_Num"])

# %% [markdown]
# ---
#
# ## Analysis: Nodes vs Actions

# %%
create_relplot(df, "New_Nodes", "New Nodes vs Actions")

# %%
create_catplot(df, "New_Nodes", "New Nodes by Problem")

# %%
df[["Problem", "Alg_Title", "New_Nodes", "Actions"]].set_index(
    "Problem"
).sort_values(by=["Problem", "New_Nodes"])

# %% [markdown]
# ### Discussion
#
# Solving more complex problems requires more actions.  To solve the problems here, the number of actions grows linearly.   Finding the sequence of actions needed to reach the goal requires exponentially more new nodes as problem complexity increases.  When plotted logarithmically, exponential growth appears linearly, just as we see here.  The linear pattern occurs as logarithms and exponents are inverse functions.  One notable exception is the performance of depth-first search on problem 3 is an order of magnitude better than expected.  Repeated runs are consistent, so I imagine it has to do with the structure of the problem that allows DFS to find a solution with minimal backtracking.
#
# For all problems and heuristics, greedy algorithms outperform A* by needing fewer nodes to arrive at a solution.  A* has exponential growth in space complexity and puts A* at a disadvantage.  This especially true when the number of actions grows with the increased complexity of the problem.  Breadth-first search algorithms performed the worst -- as expected -- since its frontier size grows exponentially.
# %% [markdown]
# ---
#
# ## Analysis: Time vs Actions
#

# %%
create_relplot(df, "Time", "Runtime vs Actions")

# %%
create_catplot(df, "Time", "Runtime vs Actions")

# %%
df[["Problem", "Alg_Title", "Time", "Actions"]].set_index(
    "Problem"
).sort_values(by=["Problem", "Time"])

# %%

# %% [markdown]
# ### Discussion
# Greedy search -- especially when using the h_unmet_goals heuristic -- finds solutions the fastest. The h_pg_setlevel heuristic is the slowest for all problems. For problem 1, which only requires 20 actions to solve, greedy search using the h_pg_setlevel is the slowest.  For all other problems, it is A* using this same heuristic. The h_pg_setlevel heuristic estimates:
#
# > The set level of a planning graph is the first level where all goals appear such that no pair of goal literals are mutex in the last layer of the planning graph.
#
# The h_pg_setlevel heuristic is known to perform poorly on complex problems, and here it lives up to its billing.
#
#
#
# %% [markdown]
# ---
#
# ## Analysis: Length of plan by algorithm
#
#

# %%
create_catplot(df, "Actions", "Actions by Algorithm", is_log=False)

# %%
df[["Problem", "Alg_Title", "Actions"]].set_index("Problem").sort_values(
    by=["Problem", "Actions"]
)

# %% [markdown]
# ### Discussion
#
# The length of the plan required to arrive at the goal is consistent for all algorithms for each problem. The work to derive the plan -- measured in time expended of new nodes created -- varies considerably across algorithms. Why is this? Conflicting actions have been eliminated by only allowing actions that are not mutually exclusive. Mutually exclusive or mutex actions happen in three different ways.
#
# 1. _Inconsistent_ actions negate the effects of one another. For example, `Load(C, p)` and `Unload(C, p)`.
#
# 1. _Interference_ s when the effects of one action are mutually exclusive of the precondition of another.  For example, the actions `Fly(a, b)` and `At(a)` are _mutex_ since once flown to airport b.  The plane is no longer at airport a.
#
# 1. _Competing needs_ occur when the preconditions of one action interfere with the preconditions of the other.  One plane can only fly to one airport, so the actions `Fly(p, a, b)` and `Fly(P, a, c)` compete for the having the plane `At(a)`.
#
# By eliminating mutex actions, the size of the resulting plan shrinks dramatically.
# %% [markdown]
# # Questions
#
# 1. Which algorithm or algorithms would be most appropriate for planning in a very restricted domain (i.e., one that has only a few actions) and needs to operate in real time?
#
# > Based on my analysis of time vs. actions, I would select Greedy h_unmet_goals or DFS because they performed the best for small problems (i.e., problem 1). Both completed their search in less than one-hundredth of a second, an order of magnitude faster than A* h_unmet_goals.
#
# 1. _Which algorithm or algorithms would be most appropriate for planning in very large domains (e.g., planning delivery routes for all UPS drivers in the U.S. on a given day)_
#
# > Greedy h_unmet_goals is by far the best choice for large problem domains (i.e., problem 4). Other choices are Greedy h_pg_maxlevel or A* h_unmet_goals, but they are three orders of magnitude slower.
#
#
# 1. _Which algorithm or algorithms would be most appropriate for planning problems where it is important to find only optimal plans?
#
# > A planning graph is a tree with its root at the initial state. Actions are the branches of the tree and whose effects result in a new state. A* is guaranteed to return an optimal plan when searching a tree using an admissible heuristic. Heuristics are estimates of the cost measured in the number of new states from the initial state. An admissible heuristic will always be less than or equal to the actual cost. Of the heuristics considered in this analysis, here are those that are admissible:
#
# * The h_pg_maxlevel. heuristic is the largest level-cost needed to achieve any one of the goal's conditions.
#
# * The h_pg_setlevel heuristic is the level where all the goal's conditions are achieved.
#
# > __So, to find optimal plans use either A*_h_pg_setlevel heuristic or A*_h_pg_maxlevel. In my results, A_h_pg_setlevel is about 10 faster than A_h_pg_maxlevel for all problems analyzed in this exercise. So, choose A*_h_pg_setlevel. It will return an optimal plan in the least amount of time.__
#
# BFS, too, will return an optimal plan but expands an exponentially expanding number of nodes in the process, so it consumes a lot of memory and time.__
#
# ### Cite
#
# Russell, S. J., & Norvig, P. (2010). Artificial intelligence: A modern approach (3rd ed.). Upper Saddle River: Prentice-Hall.  Chapter 3 - Solving Problems by Searching and Chapter 10 - Classical Planning.

# %%


r"""
Interface to Axiom

TODO:

- Evaluation using a file is not done. Any input line with more than a
  few thousand characters would hang the system, so currently it
  automatically raises an exception.

- All completions of a given command.

- Interactive help.

Axiom is a free GPL-compatible (modified BSD license) general
purpose computer algebra system whose development started in 1973
at IBM. It contains symbolic manipulation algorithms, as well as
implementations of special functions, including elliptic functions
and generalized hypergeometric functions. Moreover, Axiom has
implementations of many functions relating to the invariant theory
of the symmetric group `S_n.` For many links to Axiom
documentation see http://wiki.axiom-developer.org.

AUTHORS:

- Bill Page (2006-10): Created this (based on Maxima interface)

  .. note::

     Bill Page put a huge amount of effort into the Sage Axiom
     interface over several days during the Sage Days 2 coding
     sprint. This is contribution is greatly appreciated.

- William Stein (2006-10): misc touchup.

- Bill Page (2007-08): Minor modifications to support axiom4sage-0.3

.. note::

   The axiom4sage-0.3.spkg is based on an experimental version of the
   FriCAS fork of the Axiom project by Waldek Hebisch that uses
   pre-compiled cached Lisp code to build Axiom very quickly with
   clisp.

If the string "error" (case insensitive) occurs in the output of
anything from axiom, a RuntimeError exception is raised.

EXAMPLES: We evaluate a very simple expression in axiom.

::

    sage: axiom('3 * 5')                     #optional - axiom
    15
    sage: a = axiom(3) * axiom(5); a         #optional - axiom
    15

The type of a is AxiomElement, i.e., an element of the axiom
interpreter.

::

    sage: type(a)                            #optional - axiom
    <class 'sage.interfaces.axiom.AxiomElement'>
    sage: parent(a)                          #optional - axiom
    Axiom

The underlying Axiom type of a is also available, via the type
method::

    sage: a.type()                           #optional - axiom
    PositiveInteger

We factor `x^5 - y^5` in Axiom in several different ways.
The first way yields a Axiom object.

::

    sage: F = axiom.factor('x^5 - y^5'); F      #optional - axiom
               4      3    2 2    3     4
    - (y - x)(y  + x y  + x y  + x y + x )
    sage: type(F)                               #optional - axiom
    <class 'sage.interfaces.axiom.AxiomElement'>
    sage: F.type()                              #optional - axiom
    Factored Polynomial Integer

Note that Axiom objects are normally displayed using "ASCII art".

::

    sage: a = axiom(2/3); a          #optional - axiom
      2
      -
      3
    sage: a = axiom('x^2 + 3/7'); a      #optional - axiom
       2   3
      x  + -
           7

The ``axiom.eval`` command evaluates an expression in
axiom and returns the result as a string. This is exact as if we
typed in the given line of code to axiom; the return value is what
Axiom would print out.

::

    sage: print axiom.eval('factor(x^5 - y^5)')   #optional - axiom
               4      3    2 2    3     4
    - (y - x)(y  + x y  + x y  + x y + x )
    Type: Factored Polynomial Integer

We can create the polynomial `f` as a Axiom polynomial,
then call the factor method on it. Notice that the notation
``f.factor()`` is consistent with how the rest of Sage
works.

::

    sage: f = axiom('x^5 - y^5')                  #optional - axiom
    sage: f^2                                     #optional - axiom
       10     5 5    10
      y   - 2x y  + x
    sage: f.factor()                              #optional - axiom
               4      3    2 2    3     4
    - (y - x)(y  + x y  + x y  + x y + x )

Control-C interruption works well with the axiom interface, because
of the excellent implementation of axiom. For example, try the
following sum but with a much bigger range, and hit control-C.

::

    sage:  f = axiom('(x^5 - y^5)^10000')       # not tested
    Interrupting Axiom...
    ...
    <type 'exceptions.TypeError'>: Ctrl-c pressed while running Axiom

::

    sage: axiom('1/100 + 1/101')                  #optional - axiom
       201
      -----
      10100
    sage: a = axiom('(1 + sqrt(2))^5'); a         #optional - axiom
         +-+
      29\|2  + 41

TESTS: We check to make sure the subst method works with keyword
arguments.

::

    sage: a = axiom(x+2); a  #optional - axiom
    x + 2
    sage: a.subst(x=3)       #optional - axiom
    5

We verify that Axiom floating point numbers can be converted to
Python floats.

::

    sage: float(axiom(2))     #optional - axiom
    2.0
"""

###########################################################################
#       Copyright (C) 2008 Mike Hansen <mhansen@gmail.com>
#                     2007 Bill Page
#                     2006 William Stein <wstein@gmail.com>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
###########################################################################

import os, re

from expect import Expect, ExpectElement, FunctionElement, ExpectFunction
from sage.misc.misc import verbose, DOT_SAGE
from pexpect import EOF
from sage.misc.multireplace import multiple_replace

# The Axiom commands ")what thing det" ")show Matrix" and ")display
# op det" commands, gives a list of all identifiers that begin in
# a certain way.  This could maybe be useful somehow... (?)  Also
# axiom has a lot a lot of ways for getting documentation from the
# system -- this could also be useful.

class PanAxiom(Expect):
    """
    Interface to a PanAxiom interpreter.
    """
    def __init__(self, name='axiom', command='axiom -nox -noclef',
                 script_subdirectory=None, logfile=None,
                 server=None, server_tmpdir=None,
                 init_code=[')lisp (si::readline-off)']):
        """
        Create an instance of the Axiom interpreter.

        TESTS::

            sage: axiom == loads(dumps(axiom))
            True
        """
        eval_using_file_cutoff = 200
        self.__eval_using_file_cutoff = eval_using_file_cutoff
        self._COMMANDS_CACHE = '%s/%s_commandlist_cache.sobj'%(DOT_SAGE, name)
        Expect.__init__(self,
                        name = name,
                        prompt = '\([0-9]+\) -> ',
                        command = command,
                        maxread = 10,
                        script_subdirectory = script_subdirectory,
                        server=server,
                        server_tmpdir=server_tmpdir,
                        restart_on_ctrlc = False,
                        verbose_start = False,
                        init_code = init_code,
                        logfile = logfile,
                        eval_using_file_cutoff=eval_using_file_cutoff)
        self._prompt_wait = self._prompt

    def _start(self):
        """
        Start the Axiom interpreter.

        EXAMPLES::

            sage: a = Axiom()
            sage: a.is_running()
            False
            sage: a._start()     #optional - axiom
            sage: a.is_running() #optional - axiom
            True
            sage: a.quit()       #optional - axiom
        """
        Expect._start(self)
        out = self._eval_line(')set functions compile on', reformat=False)
        out = self._eval_line(')set output length 245', reformat=False)
        out = self._eval_line(')set message autoload off', reformat=False)

    def _read_in_file_command(self, filename):
        r"""
        EXAMPLES::

            sage: axiom._read_in_file_command('test.input')
            ')read test.input \n'
            sage: axiom._read_in_file_command('test')
            Traceback (most recent call last):
            ...
            ValueError: the filename must end with .input

        ::

            sage: filename = tmp_filename(ext='.input')
            sage: f = open(filename, 'w')
            sage: f.write('xx := 22;\n')
            sage: f.close()
            sage: axiom.read(filename)    # optional - axiom
            sage: axiom.get('xx')         #optional
            '22'
        """
        if not filename.endswith('.input'):
            raise ValueError, "the filename must end with .input"

        # For some reason this trivial comp
        # keeps certain random freezes from occurring.  Do not remove this.
        # The space before the \n is also important.
        return ')read %s \n'%filename

    def _quit_string(self):
        """
        Returns the string used to quit Axiom.

        EXAMPLES::

            sage: axiom._quit_string()
            ')lisp (quit)'
            sage: a = Axiom()
            sage: a.is_running()
            False
            sage: a._start()     #optional - axiom
            sage: a.is_running() #optional - axiom
            True
            sage: a.quit()       #optional - axiom
            sage: a.is_running() #optional - axiom
            False
        """
        return ')lisp (quit)'

    def _commands(self):
        """
        Returns a list of commands available. This is done by parsing the
        result of the first section of the output of ')what things'.

        EXAMPLES::

            sage: cmds = axiom._commands() #optional - axiom
            sage: len(cmds) > 100  #optional - axiom
            True
            sage: '<' in cmds      #optional - axiom
            True
            sage: 'factor' in cmds #optional - axiom
            True
        """
        s = self.eval(")what things")
        start = '\r\n\r\n#'
        i = s.find(start)
        end = "To get more information about"
        j = s.find(end)
        s = s[i+len(start):j].split()
        return s

    def trait_names(self, verbose=True, use_disk_cache=True):
        """
        Returns a list of all the commands defined in Axiom and optionally
        (per default) store them to disk.

        EXAMPLES::

            sage: c = axiom.trait_names(use_disk_cache=False, verbose=False) #optional - axiom
            sage: len(c) > 100  #optional - axiom
            True
            sage: 'factor' in c  #optional - axiom
            True
            sage: '**' in c     #optional - axiom
            False
            sage: 'upperCase?' in c  #optional - axiom
            False
            sage: 'upperCase_q' in c #optional - axiom
            True
            sage: 'upperCase_e' in c #optional - axiom
            True
        """
        try:
            return self.__trait_names
        except AttributeError:
            import sage.misc.persist
            if use_disk_cache:
                try:
                    self.__trait_names = sage.misc.persist.load(self._COMMANDS_CACHE)
                    return self.__trait_names
                except IOError:
                    pass
            if verbose:
                print "\nBuilding %s command completion list (this takes"%(self)
                print "a few seconds only the first time you do it)."
                print "To force rebuild later, delete %s."%self._COMMANDS_CACHE
            v = self._commands()

            #Process we now need process the commands to strip out things which
            #are not valid Python identifiers.
            import re
            valid = re.compile('[^a-zA-Z0-9_]+')
            names = [x for x in v if valid.search(x) is None]

            #Change everything that ends with ? to _q and
            #everything that ends with ! to _e
            names += [x[:-1]+"_q" for x in v if x.endswith("?")]
            names += [x[:-1]+"_e" for x in v if x.endswith("!")]

            self.__trait_names = names
            if len(v) > 200:
                # Axiom is actually installed.
                sage.misc.persist.save(v, self._COMMANDS_CACHE)
            return names

    def set(self, var, value):
        """
        Set the variable var to the given value.

        EXAMPLES::

            sage: axiom.set('xx', '2')    #optional - axiom
            sage: axiom.get('xx')         #optional - axiom
            '2'

            sage: fricas.set('xx', '2')    #optional - fricas
            sage: fricas.get('xx')         #optional - fricas
            '2'

        """
        cmd = '%s := %s'%(var, value)
        out = self._eval_line(cmd, reformat=False)

        if out.find("error") != -1:
            raise TypeError, "Error executing code in Axiom\nCODE:\n\t%s\nAxiom ERROR:\n\t%s"%(cmd, out)

    def get(self, var):
        r"""
        Get the string value of the Axiom variable var.

        EXAMPLES::

            sage: axiom.set('xx', '2')    #optional - axiom
            sage: axiom.get('xx')         #optional - axiom
            '2'
            sage: a = axiom('(1 + sqrt(2))^5') #optional - axiom
            sage: axiom.get(a.name())          #optional - axiom
            '     +-+\r\r\n  29\\|2  + 41'
        """
        s = self._eval_line(str(var))
        i = s.rfind('Type:')
        s = s[:i].rstrip().lstrip("\n")
        if '\n' not in s:
            s = s.strip()
        return s

    def _eval_line(self, line, reformat=True, allow_use_file=False,
                   wait_for_prompt=True, restart_if_needed=False):
        """
        EXAMPLES::

            sage: print axiom._eval_line('2+2')  #optional - axiom
              4
                                                       Type: PositiveInteger
        """
        if not wait_for_prompt:
            return Expect._eval_line(self, line)
        line = line.rstrip().rstrip(';')
        if line == '':
            return ''
        if len(line) > 3000:
            raise NotImplementedError, "evaluation of long input lines (>3000 characters) in Axiom not yet implemented."
        if self._expect is None:
            self._start()
        if allow_use_file and self.__eval_using_file_cutoff and \
                            len(line) > self.__eval_using_file_cutoff:
            return self._eval_line_using_file(line)
        try:
            E = self._expect
            # debug
            self._synchronize(cmd='1+%s\n')
            verbose("in = '%s'"%line,level=3)
            E.sendline(line)
            self._expect.expect(self._prompt)
            out = self._expect.before
            # debug
            verbose("out = '%s'"%out,level=3)
        except EOF:
          if self._quit_string() in line:
             return ''
        except KeyboardInterrupt:
            self._keyboard_interrupt()

        if '>> Error detected within library code:' in out or \
           'Cannot find a definition or applicable library operation named' in out:
            raise RuntimeError, out

        if not reformat:
            return out
        if 'error' in out:
            return out
        #out = out.lstrip()
        i = out.find('\n')
        out = out[i+1:]
        outs = out.split("\n")
        i = 0
        outline = ''
        for line in outs:
            line = line.rstrip()
            # print "'%s'"%line
            if line[:4] == '   (':
                i = line.find('(')
                i += line[i:].find(')')
                if line[i+1:] == "":
                    i = 0
                    outs = outs[1:]
                break;
        out = "\n".join(line[i+1:] for line in outs[1:])
        return out

    # define relational operators
    def _equality_symbol(self):
        """equality symbol

        EXAMPLES::

            sage: a = axiom(x==6); a    #optional axiom
            x= 6
            sage: a = fricas(x==6); a   #optional fricas
            x= 6
        """
        return "="

class Axiom(PanAxiom):
    def __reduce__(self):
        """
        EXAMPLES::

            sage: axiom.__reduce__()
            (<function reduce_load_Axiom at 0x...>, ())
            sage: f, args = _
            sage: f(*args)
            Axiom
        """
        return reduce_load_Axiom, tuple([])

    def _function_class(self):
        """
        Return the AxiomExpectFunction class.

        EXAMPLES::

            sage: axiom._function_class()
            <class 'sage.interfaces.axiom.AxiomExpectFunction'>
            sage: type(axiom.gcd)
            <class 'sage.interfaces.axiom.AxiomExpectFunction'>
        """
        return AxiomExpectFunction

    def _object_class(self):
        """
        EXAMPLES::

            sage: axiom._object_class()
            <class 'sage.interfaces.axiom.AxiomElement'>
            sage: type(axiom(2)) #optional - axiom
            <class 'sage.interfaces.axiom.AxiomElement'>
        """
        return AxiomElement

    def _function_element_class(self):
        """
        Returns the Axiom function element class.

        EXAMPLES::

            sage: axiom._function_element_class()
            <class 'sage.interfaces.axiom.AxiomFunctionElement'>
            sage: type(axiom(2).gcd) #optional - axiom
            <class 'sage.interfaces.axiom.AxiomFunctionElement'>
        """
        return AxiomFunctionElement

    def console(self):
        """
        Spawn a new Axiom command-line session.

        EXAMPLES::

            sage: axiom.console() #not tested
                                    AXIOM Computer Algebra System
                                    Version: Axiom (January 2009)
                           Timestamp: Sunday January 25, 2009 at 07:08:54
            -----------------------------------------------------------------------------
               Issue )copyright to view copyright notices.
               Issue )summary for a summary of useful system commands.
               Issue )quit to leave AXIOM and return to shell.
            -----------------------------------------------------------------------------
        """
        axiom_console()

class PanAxiomElement(ExpectElement):
    def __call__(self, x):
        """
        EXAMPLES::

            sage: f = axiom(x+2) #optional - axiom
            sage: f(2)           #optional - axiom
            4
        """
        self._check_valid()
        P = self.parent()
        return P('%s(%s)'%(self.name(), x))

    def __cmp__(self, other):
        """
        EXAMPLES::

            sage: two = axiom(2)  #optional - axiom
            sage: two == 2        #optional - axiom
            True
            sage: two == 3        #optional - axiom
            False
            sage: two < 3         #optional - axiom
            True
            sage: two > 1         #optional - axiom
            True

            sage: a = axiom(1); b = axiom(2)  #optional - axiom
            sage: a == b                      #optional - axiom
            False
            sage: a < b                       #optional - axiom
            True
            sage: a > b                       #optional - axiom
            False
            sage: b < a                       #optional - axiom
            False
            sage: b > a                       #optional - axiom
            True

        We can also compare more complicated object such as functions::

            sage: f = axiom('sin(x)'); g = axiom('cos(x)')    #optional - axiom
            sage: f == g                                      #optional - axiom
            False

        """
        P = self.parent()
        if 'true' in P.eval("(%s = %s) :: Boolean"%(self.name(),other.name())):
            return 0
        elif 'true' in P.eval("(%s < %s) :: Boolean"%(self.name(), other.name())):
            return -1
        elif 'true' in P.eval("(%s > %s) :: Boolean"%(self.name(),other.name())):
            return 1

        # everything is supposed to be comparable in Python, so we define
        # the comparison thus when no comparable in interfaced system.
        if (hash(self) < hash(other)):
            return -1
        else:
            return 1

    def type(self):
        """
        Returns the type of an AxiomElement.

        EXAMPLES::

            sage: axiom(x+2).type()  #optional - axiom
            Polynomial Integer
        """
        P = self._check_valid()
        s = P._eval_line(self.name())
        i = s.rfind('Type:')
        return P(s[i+5:].strip())

    def __len__(self):
        """
        Return the length of a list.

        EXAMPLES::

            sage: v = axiom('[x^i for i in 0..5]')            # optional - axiom
            sage: len(v)                                      # optional - axiom
            6
        """
        P = self._check_valid()
        s = P.eval('# %s '%self.name())
        i = s.rfind('Type')
        return int(s[:i-1])

    def __getitem__(self, n):
        r"""
        Return the n-th element of this list.

        .. note::

           Lists are 1-based.

        EXAMPLES::

            sage: v = axiom('[i*x^i for i in 0..5]'); v          # optional - axiom
                     2   3   4   5
              [0,x,2x ,3x ,4x ,5x ]
            sage: v[4]                                           # optional - axiom
                3
              3x
            sage: v[1]                                           # optional - axiom
            0
            sage: v[10]                                          # optional - axiom
            Traceback (most recent call last):
            ...
            IndexError: index out of range
        """
        n = int(n)
        if n <= 0 or n > len(self):
            raise IndexError, "index out of range"
        P = self._check_valid()
        if not isinstance(n, tuple):
            return P.new('%s(%s)'%(self._name, n))
        else:
            return P.new('%s(%s)'%(self._name, str(n)[1:-1]))

    def comma(self, *args):
        """
        Returns a Axiom tuple from self and args.

        EXAMPLES::

            sage: two = axiom(2)  #optional - axiom
            sage: two.comma(3)    #optional - axiom
            [2,3]
            sage: two.comma(3,4)  #optional - axiom
            [2,3,4]
            sage: _.type()        #optional - axiom
            Tuple PositiveInteger

            sage: two = fricas(2)  #optional - fricas
            sage: two.comma(3)     #optional - fricas
            [2,3]
            sage: two.comma(3,4)   #optional - fricas
            [2,3,4]
            sage: _.type()         #optional - fricas
            Tuple(PositiveInteger)

        """
        P = self._check_valid()
        args = list(args)
        for i, arg in enumerate(args):
            if not isinstance(arg, AxiomElement) or arg.parent() is not P:
                args[i] = P(arg)
        cmd = "(" + ",".join([x.name() for x in [self]+args]) + ")"
        return P(cmd)

    def _latex_(self):
        r"""
        EXAMPLES::

            sage: a = axiom(1/2) #optional - axiom
            sage: latex(a)       #optional - axiom
            \frac{1}{2}

            sage: a = fricas(1/2) #optional - fricas
            sage: latex(a)        #optional - fricas
            1 \over 2
        """
        self._check_valid()
        P = self.parent()
        s = P._eval_line('outputAsTex(%s)'%self.name(), reformat=False)
        if not '$$' in s:
            raise RuntimeError, "Error texing axiom object."
        i = s.find('$$')
        j = s.rfind('$$')
        s = s[i+2:j]
        s = multiple_replace({'\r':'', '\n':' ',
                              ' \\sp ':'^',
                              '\\arcsin ':'\\sin^{-1} ',
                              '\\arccos ':'\\cos^{-1} ',
                              '\\arctan ':'\\tan^{-1} '},
            re.sub(r'\\leqno\(.*?\)','',s)) # no eq number!
        return s

    def as_type(self, type):
        """
        Returns self as type.

        EXAMPLES::

            sage: a = axiom(1.2); a            #optional - axiom
            1.2
            sage: a.as_type(axiom.DoubleFloat) #optional - axiom
            1.2
            sage: _.type()                     #optional - axiom
            DoubleFloat

        ::

            sage: a = fricas(1.2); a            #optional - fricas
            1.2
            sage: a.as_type(fricas.DoubleFloat) #optional - fricas
            1.2
            sage: _.type()                      #optional - fricas
            DoubleFloat

        """
        P = self._check_valid()
        type = P(type)
        return P.new("%s :: %s"%(self.name(), type.name()))

    def unparsed_input_form(self):
        """
        Get the linear string representation of this object, if possible
        (often it isn't).

        EXAMPLES::

            sage: a = axiom(x^2+1); a     #optional - axiom
               2
              x  + 1
            sage: a.unparsed_input_form() #optional - axiom
            'x*x+1'

            sage: a = fricas(x^2+1)       #optional - fricas
            sage: a.unparsed_input_form() #optional - fricas
            'x^2+1'
        """
        P = self._check_valid()
        s = P.eval('unparse(%s::InputForm)'%self._name)
        if 'translation error' in s or 'Cannot convert' in s:
            raise NotImplementedError
        s = multiple_replace({'\r\n':'', # fix stupid Fortran-ish
                              'DSIN(':'sin(',
                              'DCOS(':'cos(',
                              'DTAN(':'tan(',
                              'DSINH(':'sinh('}, s)
        r = re.search(r'"(.*)"',s)
        if r:
            return r.groups(0)[0]
        else:
            return s

    def _sage_(self):
        """
        Convert self to a Sage object.

        EXAMPLES::

            sage: a = axiom(1/2); a #optional - axiom
              1
              -
              2
            sage: a.sage()          #optional - axiom
            1/2
            sage: _.parent()        #optional - axiom
            Rational Field

            sage: gp(axiom(1/2))    #optional - axiom
            1/2

            sage: fricas(1/2).sage() #optional - fricas
            1/2

        DoubleFloat's in Axiom are converted to be in RDF in Sage.

        ::

            sage: axiom(2.0).as_type('DoubleFloat').sage()  #optional - axiom
            2.0
            sage: _.parent() #optional - axiom
            Real Double Field

            sage: axiom(2.1234)._sage_() #optional - axiom
            2.12340000000000
            sage: _.parent()             #optional - axiom
            Real Field with 53 bits of precision
            sage: a = RealField(100)(pi)
            sage: axiom(a)._sage_()      #optional - axiom
            3.1415926535897932384626433833
            sage: _.parent()             #optional - axiom
            Real Field with 100 bits of precision
            sage: axiom(a)._sage_() == a #optional - axiom
            True
            sage: axiom(2.0)._sage_() #optional - axiom
            2.00000000000000
            sage: _.parent() #optional  - axiom
            Real Field with 53 bits of precision

        We can also convert Axiom's polynomials to Sage polynomials.
            sage: a = axiom(x^2 + 1)   #optional - axiom
            sage: a.type()             #optional - axiom
            Polynomial Integer
            sage: a.sage()             #optional - axiom
            x^2 + 1
            sage: _.parent()           #optional - axiom
            Univariate Polynomial Ring in x over Integer Ring
            sage: axiom('x^2 + y^2 + 1/2').sage()    #optional - axiom
            y^2 + x^2 + 1/2
            sage: _.parent()                         #optional - axiom
            Multivariate Polynomial Ring in y, x over Rational Field

        """
        P = self._check_valid()
        type = str(self.type())

        if type in ["Type", "Domain"]:
            return self._sage_domain()

        if type == "Float":
            from sage.rings.all import RealField, ZZ
            prec = max(self.mantissa().length()._sage_(), 53)
            R = RealField(prec)
            x,e,b = self.unparsed_input_form().lstrip('float(').rstrip(')').split(',')
            return R(ZZ(x)*ZZ(b)**ZZ(e))
        elif type == "DoubleFloat":
            from sage.rings.all import RDF
            return RDF(repr(self))
        elif type.startswith('Polynomial'):
            from sage.rings.all import PolynomialRing
            base_ring = P(type.lstrip('Polynomial '))._sage_domain()
            vars = str(self.variables())[1:-1]
            R = PolynomialRing(base_ring, vars)
            return R(self.unparsed_input_form())

        #If all else fails, try using the unparsed input form
        try:
            import sage.misc.sage_eval
            return sage.misc.sage_eval.sage_eval(self.unparsed_input_form())
        except StandardError:
            raise NotImplementedError

    def _sage_domain(self):
        """
        A helper function for converting Axiom domains to the corresponding
        Sage object.

        EXAMPLES::

            sage: axiom('Integer').sage()  #optional - axiom
            Integer Ring
            sage: fricas('Integer').sage() #optional - fricas
            Integer Ring

            sage: axiom('Fraction Integer').sage()  #optional - axiom
            Rational Field
            sage: fricas('Fraction Integer').sage() #optional - fricas
            Rational Field

            sage: axiom('DoubleFloat').sage()  #optional - axiom
            Real Double Field
            sage: fricas('DoubleFloat').sage() #optional - fricas
            Real Double Field

        """
        P = self._check_valid()
        name = str(self)
        if name == 'Integer':
            from sage.rings.all import ZZ
            return ZZ
        elif name == 'DoubleFloat':
            from sage.rings.all import RDF
            return RDF
        elif name.startswith('Fraction '):
            return P(name.lstrip('Fraction '))._sage_domain().fraction_field()

        raise NotImplementedError

class AxiomElement(PanAxiomElement):
    pass

class PanAxiomFunctionElement(FunctionElement):
    def __init__(self, object, name):
        """
        TESTS::

            sage: a = axiom('"Hello"') #optional - axiom
            sage: a.upperCase_q        #optional - axiom
            upperCase?
            sage: a.upperCase_e        #optional - axiom
            upperCase!
            sage: a.upperCase_e()      #optional - axiom
            "HELLO"
        """
        if name.endswith("_q"):
            name = name[:-2] + "?"
        elif name.endswith("_e"):
            name = name[:-2] + "!"
        FunctionElement.__init__(self, object, name)

class AxiomFunctionElement(PanAxiomFunctionElement):
    pass

class PanAxiomExpectFunction(ExpectFunction):
    def __init__(self, parent, name):
        """
        TESTS::

            sage: axiom.upperCase_q
            upperCase?
            sage: axiom.upperCase_e
            upperCase!
        """
        if name.endswith("_q"):
            name = name[:-2] + "?"
        elif name.endswith("_e"):
            name = name[:-2] + "!"
        ExpectFunction.__init__(self, parent, name)

class AxiomExpectFunction(PanAxiomExpectFunction):
    pass

def is_AxiomElement(x):
    """
    Returns True of x is of type AxiomElement.

    EXAMPLES::

        sage: from sage.interfaces.axiom import is_AxiomElement
        sage: is_AxiomElement(axiom(2)) #optional - axiom
        True
        sage: is_AxiomElement(2)
        False
    """
    return isinstance(x, AxiomElement)

#Instances
axiom = Axiom(name='axiom')

def reduce_load_Axiom():
    """
    Returns the Axiom interface object defined in
    sage.interfaces.axiom.

    EXAMPLES::

        sage: from sage.interfaces.axiom import reduce_load_Axiom
        sage: reduce_load_Axiom()
        Axiom
    """
    return axiom

import os
def axiom_console():
    """
    Spawn a new Axiom command-line session.

    EXAMPLES::

        sage: axiom_console() #not tested
                                AXIOM Computer Algebra System
                                Version: Axiom (January 2009)
                       Timestamp: Sunday January 25, 2009 at 07:08:54
        -----------------------------------------------------------------------------
           Issue )copyright to view copyright notices.
           Issue )summary for a summary of useful system commands.
           Issue )quit to leave AXIOM and return to shell.
        -----------------------------------------------------------------------------

    """
    os.system('axiom -nox')


from sympy.core.basic import Basic
from sympy.printing import pprint

import random

def interactive_traversal(expr):
    """Traverse a tree asking a user which branch to choose. """

    RED, BRED = '\033[0;31m', '\033[1;31m'
    GREEN, BGREEN = '\033[0;32m', '\033[1;32m'
    YELLOW, BYELLOW = '\033[0;33m', '\033[1;33m'  # noqa
    BLUE, BBLUE = '\033[0;34m', '\033[1;34m'      # noqa
    MAGENTA, BMAGENTA = '\033[0;35m', '\033[1;35m'# noqa
    CYAN, BCYAN = '\033[0;36m', '\033[1;36m'      # noqa
    END = '\033[0m'

    def cprint(*args):
        print("".join(map(str, args)) + END)

    def _interactive_traversal(expr, stage):
        if stage > 0:
            print()

        cprint("Current expression (stage ", BYELLOW, stage, END, "):")
        print(BCYAN)
        pprint(expr)
        print(END)

        if isinstance(expr, Basic):
            if expr.is_Add:
                args = expr.as_ordered_terms()
            elif expr.is_Mul:
                args = expr.as_ordered_factors()
            else:
                args = expr.args
        elif hasattr(expr, "__iter__"):
            args = list(expr)
        else:
            return expr

        n_args = len(args)

        if not n_args:
            return expr

        for i, arg in enumerate(args):
            cprint(GREEN, "[", BGREEN, i, GREEN, "] ", BLUE, type(arg), END)
            pprint(arg)
            print()

        if n_args == 1:
            choices = '0'
        else:
            choices = '0-%d' % (n_args - 1)

        try:
            choice = input("Your choice [%s,f,l,r,d,?]: " % choices)
        except EOFError:
            result = expr
            print()
        else:
            if choice == '?':
                cprint(RED, "%s - select subexpression with the given index" %
                       choices)
                cprint(RED, "f - select the first subexpression")
                cprint(RED, "l - select the last subexpression")
                cprint(RED, "r - select a random subexpression")
                cprint(RED, "d - done\n")

                result = _interactive_traversal(expr, stage)
            elif choice in ('d', ''):
                result = expr
            elif choice == 'f':
                result = _interactive_traversal(args[0], stage + 1)
            elif choice == 'l':
                result = _interactive_traversal(args[-1], stage + 1)
            elif choice == 'r':
                result = _interactive_traversal(random.choice(args), stage + 1)
            else:
                try:
                    choice = int(choice)
                except ValueError:
                    cprint(BRED,
                           "Choice must be a number in %s range\n" % choices)
                    result = _interactive_traversal(expr, stage)
                else:
                    if choice < 0 or choice >= n_args:
                        cprint(BRED, "Choice must be in %s range\n" % choices)
                        result = _interactive_traversal(expr, stage)
                    else:
                        result = _interactive_traversal(args[choice], stage + 1)

        return result

    return _interactive_traversal(expr, 0)


import numpy as np

a = np.arange(10) * 10
print(a)
# [ 0 10 20 30 40 50 60 70 80 90]

print(a[5])
# 50

print(a[8])
# 80

print(a[[5, 8]])
# [50 80]

print(a[[5, 4, 8, 0]])
# [50 40 80  0]

print(a[[5, 5, 5, 5]])
# [50 50 50 50]

idx = np.array([[5, 4], [8, 0]])
print(idx)
# [[5 4]
#  [8 0]]

print(a[idx])
# [[50 40]
#  [80  0]]

# print(a[[[5, 4], [8, 0]]])
# IndexError: too many indices for array

print(a[[[[5, 4], [8, 0]]]])
# [[50 40]
#  [80  0]]

a_2d = np.arange(12).reshape((3, 4))
print(a_2d)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

print(a_2d[0])
# [0 1 2 3]

print(a_2d[2])
# [ 8  9 10 11]

print(a_2d[[2, 0]])
# [[ 8  9 10 11]
#  [ 0  1  2  3]]

print(a_2d[[2, 2, 2]])
# [[ 8  9 10 11]
#  [ 8  9 10 11]
#  [ 8  9 10 11]]

print(a_2d[:, 1])
# [1 5 9]

print(a_2d[:, 3])
# [ 3  7 11]

print(a_2d[:, 1:2])
# [[1]
#  [5]
#  [9]]

print(a_2d[:, [3, 1]])
# [[ 3  1]
#  [ 7  5]
#  [11  9]]

print(a_2d[:, [3, 3, 3]])
# [[ 3  3  3]
#  [ 7  7  7]
#  [11 11 11]]

print(a_2d[0, 1])
# 1

print(a_2d[2, 3])
# 11

print(a_2d[[0, 2], [1, 3]])
# [ 1 11]

# index
# [[0, 1] [2, 3]]

# print(a_2d[[0, 2, 1], [1, 3]])
# IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (3,) (2,)

print(a_2d[[[0, 0], [2, 2]], [[1, 3], [1, 3]]])
# [[ 1  3]
#  [ 9 11]]

# index
# [[0, 1] [0, 3]
#  [2, 1] [2, 3]]

print(a_2d[[[0], [2]], [1, 3]])
# [[ 1  3]
#  [ 9 11]]

idxs = np.ix_([0, 2], [1, 3])
print(idxs)
# (array([[0],
#        [2]]), array([[1, 3]]))

print(type(idxs))
# <class 'tuple'>

print(type(idxs[0]))
# <class 'numpy.ndarray'>

print(idxs[0])
# [[0]
#  [2]]

print(idxs[1])
# [[1 3]]

print(a_2d[np.ix_([0, 2], [1, 3])])
# [[ 1  3]
#  [ 9 11]]

print(a_2d[np.ix_([2, 0], [3, 3, 3])])
# [[11 11 11]
#  [ 3  3  3]]

print(a_2d[[0, 2]][:, [1, 3]])
# [[ 1  3]
#  [ 9 11]]

a_2d = np.arange(12).reshape((3, 4))
print(a_2d)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

a_2d[np.ix_([0, 2], [1, 3])] = 100
print(a_2d)
# [[  0 100   2 100]
#  [  4   5   6   7]
#  [  8 100  10 100]]

a_2d[np.ix_([0, 2], [1, 3])] = [100, 200]
print(a_2d)
# [[  0 100   2 200]
#  [  4   5   6   7]
#  [  8 100  10 200]]

a_2d[np.ix_([0, 2], [1, 3])] = [[100, 200], [300, 400]]
print(a_2d)
# [[  0 100   2 200]
#  [  4   5   6   7]
#  [  8 300  10 400]]

print(a_2d[[0, 2]][:, [1, 3]])
# [[100 200]
#  [300 400]]

a_2d[[0, 2]][:, [1, 3]] = 0
print(a_2d)
# [[  0 100   2 200]
#  [  4   5   6   7]
#  [  8 300  10 400]]

a_2d = np.arange(12).reshape((3, 4))
print(a_2d)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

a_2d[[2, 0]] = [[100, 200, 300, 400], [500, 600, 700, 800]]
print(a_2d)
# [[500 600 700 800]
#  [  4   5   6   7]
#  [100 200 300 400]]

a_2d[[2, 2]] = [[-1, -2, -3, -4], [-5, -6, -7, -8]]
print(a_2d)
# [[500 600 700 800]
#  [  4   5   6   7]
#  [ -5  -6  -7  -8]]

a_2d = np.arange(12).reshape((3, 4))
print(a_2d)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

a_fancy = a_2d[np.ix_([0, 2], [1, 3])]
print(a_fancy)
# [[ 1  3]
#  [ 9 11]]

a_fancy[0, 0] = 100
print(a_fancy)
# [[100   3]
#  [  9  11]]

print(a_2d)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

a_2d = np.arange(12).reshape((3, 4))
print(a_2d)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

print(a_2d[[2, 0], ::-1])
# [[11 10  9  8]
#  [ 3  2  1  0]]

print(a_2d[::2, [3, 0, 1]])
# [[ 3  0  1]
#  [11  8  9]]


# encoding: UTF-8

'''
本文件中包含的是遗传算法参数优化器的实现
华富资产 李来佳
'''

from typing import Callable, List
from itertools import product
from functools import lru_cache
from time import time
from datetime import  datetime
import multiprocessing
import random
import traceback
from copy import copy, deepcopy
from uuid import uuid1
import logging
import os
import numpy as np

# deap是遗传算法的实现工具
from deap import creator, base, tools, algorithms

from vnpy.app.cta_strategy_pro.portfolio_testing import PortfolioTestingEngine
from vnpy.trader.util_logger import setup_logger

class OptimizationSetting:
    """
    Setting for runnning optimization.
    设置参数优化
    """

    def __init__(self):
        """"""
        self.params = {}
        self.target_name = ""

    def add_parameter(
            self, name: str,
            start: float,
            end: float = None,
            step: float = None
    ):
        """添加int/float类型参数优化"""
        if not end and not step:
            self.params[name] = [start]
            return

        if start >= end:
            print("参数优化起始点必须小于终止点")
            return

        if step <= 0:
            print("参数优化步进必须大于0")
            return

        value = start
        value_list = []

        while value <= end:
            value_list.append(value)
            value += step

        # win: [20,22,24]
        self.params[name] = value_list

    def add_parameters(self,
                       name: str,
                       values: List = []):
        """添加参数清单"""
        self.params[name] = values

    def set_target(self, target_name: str):
        """"""
        self.target_name = target_name

    def generate_setting(self) -> List:
        """"""
        keys = self.params.keys()
        values = self.params.values()
        products = list(product(*values))

        settings = []
        for p in products:
            setting = dict(zip(keys, p))
            settings.append(setting)

        return settings

    def generate_setting_ga(self):
        """"""
        settings_ga = []  # [ [(key1:value1), (key2:value2)], [],,,]
        settings = self.generate_setting()
        for d in settings:
            # dict => [(key1:value1), (key2:value2)]
            param = [tuple(i) for i in d.items()]
            settings_ga.append(param)
        return settings_ga

class GeneticOptimize(object):
    """
    遗传算法优化器
    """
    def __init__(self):

        self.s = OptimizationSetting()  # 参数生成器
        self.be = PortfolioTestingEngine()  # (组合)回测引擎

        self.settings_for_ga = []  # 供遗传算法进行优化的参数列表
        self.settings_for_env = {}  # 供回测使用的参数配置，如合约、资金账号、回测时间等
        self.settings_for_strategies = {}  # cta_strategy_settings， dict格式：{ "strategy_instance_name": {策略配置}}

        self.target_names = []  # 回测结果目标值名称列表，如最大化收益回撤比，最大化夏普比率
        self.logger = None

    creator.create("FitnessMulti", base.Fitness, weights=(1.0, 1.0))
    creator.create("Individual", list, fitness=creator.FitnessMulti)

    def write_log(self,msg: str, level: int = logging.DEBUG):

        if self.logger:
            self.logger.log(msg=msg, level=level)

    def init_setting(self,
                     test_settings,
                     strategy_settings,
                     target_names):

        # 回测引擎所需的所有基本配置
        self.settings_for_env = deepcopy(test_settings)

        if not self.logger:
            logs_folder = os.path.abspath(os.path.join(os.getcwd(), 'log'))
            filename = os.path.abspath(os.path.join(logs_folder, '{}'.format(test_settings['name'])))
            self.logger = setup_logger(file_name=filename,
                                       name="go_{}".format(datetime.now().strftime("%Y%m%d_%H%M%S")),
                                       log_level=logging.DEBUG,
                                       backtesing=True)

        # 回测实例及配置
        self.settings_for_strategies = deepcopy(strategy_settings)

        # 优化的目标值清单
        self.target_names = deepcopy(target_names)

        return self.settings_for_env

    def add_parameter(self, name, start, end, step):
        """
        添加参数
        :param name: 参数名称
        :param start: 开始数字
        :param end: 结束数字
        :param step: 步进
        :return:
        """
        self.s.add_parameter(name, start, end, step)

    def add_parameters(self, name, values):
        """
        添加参数
        :param name: 参数名
        :param values: 参数值列表
        :return:
        """
        self.s.add_parameters(name, deepcopy(values))

    def generate_setting_for_ga(self):
        """
        产生遗传算法的参数列表
        :return:
        """
        settings = self.s.generate_setting()
        for d in settings:
            param = [tuple(i) for i in d.items()]
            self.settings_for_ga.append(param)
        return self.settings_for_ga

    def generate_parameter(self):
        """"""
        return random.choice(self.settings_for_ga)

    def mutArrayGroup(self, individual, indpb):
        size = len(individual)
        paralist = self.generate_parameter()
        for i in range(size):
            if random.random() < indpb:
                individual[i] = paralist[i]
        return individual,

    def object_func(self, strategy_avg):
        """
        遗传优化的目标执行函数
        :param strategy_avg:
        :return:
        """
        return self._object_func(tuple(strategy_avg))

    def optimize(self):
        """
        运行优化
        :return:
        """
        start = time()
        toolbox = base.Toolbox()

        # 使用多进程
        pool = multiprocessing.Pool(multiprocessing.cpu_count())
        toolbox.register("map", pool.map)
        # 初始化
        toolbox.register("individual", tools.initIterate, creator.Individual, self.generate_parameter)
        toolbox.register("population", tools.initRepeat, list, toolbox.individual)
        toolbox.register("mate", tools.cxTwoPoint)
        toolbox.register("mutate", self.mutArrayGroup, indpb=1)
        toolbox.register("evaluate", self.object_func)
        toolbox.register("select", tools.selNSGA2)

        MU = 16
        LAMBDA = 20
        POP = 20
        pop = toolbox.population(POP)
        CXPB, MUTPB, NGEN = 0.95, 0.05, 4
        hof = tools.ParetoFront()

        stats = tools.Statistics(lambda ind: ind.fitness.values)
        np.set_printoptions(suppress=True)
        stats.register("mean", np.mean, axis=0)
        stats.register("std", np.std, axis=0)
        stats.register("min", np.min, axis=0)
        stats.register("max", np.max, axis=0)

        self.write_log("开始运行遗传算法，每代族群总数：%s, 优良品种筛选个数：%s，迭代次数：%s，交叉概率：%s，突变概率：%s" % (POP, MU, NGEN, CXPB, MUTPB))
        algorithms.eaMuPlusLambda(pop, toolbox, MU, LAMBDA, CXPB, MUTPB, NGEN, stats, halloffame=hof, verbose=True)
        end = time()
        cost = int((end - start))

        self.write_log("遗传算法优化完成，耗时%s秒" % (cost))
        self.write_log("----------输出帕累托前沿解集,解集数量%s----------" % (len(hof)))
        # return hof
        for i in range(len(hof)):
            solution = hof[i]
            self.write_log(solution)

    @lru_cache(maxsize=1000000)
    def _object_func(self, strategy_avg):
        """
        使用了缓存，可以减少重复参数的运行
        :param strategy_avg: 选取的策略参数
        :return:
        """
        engine = self.be
        self.settings_for_env['name'] = self.settings_for_env['name'] + '_' + str(uuid1())
        engine.prepare_env(self.settings_for_env)
        # 选取参数 => dict 结构
        ga_setting = dict(strategy_avg)
        settings_for_strategies = deepcopy(self.settings_for_strategies)

        # 策略实例名
        for ins_name in list(settings_for_strategies.keys()):
            # 策略实例配置
            ins_config = settings_for_strategies[ins_name]
            # 策略实例的策略参数
            strategy_setting = ins_config['setting']
            # 候选参数值 => 更新 => 策略参数值
            for k, v in ga_setting.items():
                if k in strategy_setting:
                    strategy_setting.update({k: v})
            # 更新回测实例配置
            settings_for_strategies.update({ins_name: ins_config})

        try:
            # 运行（组合）回测
            engine.run_portfolio_test(settings_for_strategies)
            # 回测结果，保存
            result = engine.show_backtesting_result()

            # 保存策略得内部数据
            engine.save_strategy_data()

            # 根据target_names => (value1, value2)
            return tuple([round(result.get(k, 0), 2) for k in self.target_names])

        except Exception as ex:
            self.write_log('组合回测异常{}'.format(str(ex)))
            traceback.print_exc()
            engine.save_fail_to_mongo(f'回测异常{str(ex)}')
            return tuple([0 for k in self.target_names])

def run_go(test_setting: dict, strategy_setting: dict, ga_setting: dict, target_names: List):
    """
    遗传算法优化+回测
    : test_setting, 组合回测所需的配置，包括合约信息，数据bar信息，回测时间，资金等。
    ：strategy_setting, dict, 一个或多个策略配置
    : ga_setting,dict, 代优化参数的设定： name: (开始值，结束值，步进）, name: [value1, value2,value3,,,]

    : return 定义的
    """
    # 创建遗传优化器
    GO = GeneticOptimize()

    # 初始化 环境参数、策略参数，优化目标清单
    GO.init_setting(
        test_settings=test_setting,
        strategy_settings=strategy_setting,
        target_names=target_names
    )

    # 添加代优化的参数
    for k, v in ga_setting.items():
        if isinstance(v, tuple) and len(v) == 3:
            GO.add_parameter(k, v[0], v[1], v[2])

        if isinstance(v, list):
            GO.add_parameters(k, v)

    GO.generate_setting_for_ga()

    GO.generate_parameter()

    GO.optimize()


from typing import Optional, Sequence, Union

import numpy as np

from mygrad.tensor_base import Tensor, _resolve_constant, implements_numpy_override
from mygrad.typing import ArrayLike, DTypeLikeReals, Real

Shape = Union[Sequence[int], int]

def _anything_but_tensor(x):
    if isinstance(x, Tensor):
        x = x.data
    return x

__all__ = [
    "arange",
    "empty",
    "empty_like",
    "eye",
    "geomspace",
    "identity",
    "linspace",
    "logspace",
    "ones",
    "ones_like",
    "full",
    "full_like",
    "zeros",
    "zeros_like",
]

def empty(
    shape: Shape, dtype: DTypeLikeReals = np.float32, *, constant: Optional[bool] = None
) -> Tensor:
    """Return a new Tensor of the given shape and type, without initializing entries.

    This docstring was adapted from ``numpy.empty`` [1]_

    Parameters
    ----------
    shape : Union[int, Tuple[int]]
        The shape of the empty array.

    dtype : data-type, optional (default=numpy.float32)
        The data type of the output Tensor.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor
        A tensor of uninitialized data of the given shape and dtype.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.empty.html

    See Also
    --------
    empty_like : Return an empty tensor with shape and type of input.
    ones : Return a new tensor setting values to one.
    zeros : Return a new tensor setting values to zero.
    full : Return a new tensor of given shape filled with value.

    Notes
    -----
    `empty`, unlike `zeros`, does not set the array values to zero,
    and may therefore be marginally faster.  On the other hand, it requires
    the user to manually set all the values in the array, and should be
    used with caution.

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.empty([2, 2], constant=True)
    Tensor([[ -9.74499359e+001,   6.69583040e-309],
            [  2.13182611e-314,   3.06959433e-309]])         #random

    >>> mg.empty([2, 2], dtype=int)
    Tensor([[-1073741821, -1067949133],
            [  496041986,    19249760]])                     #random
    """
    return Tensor(np.empty(shape=shape, dtype=dtype), constant=constant, copy=False)

@implements_numpy_override()
def empty_like(
    other: ArrayLike,
    dtype: Optional[DTypeLikeReals] = None,
    shape: Optional[Union[int, Sequence[int]]] = None,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """Return a new Tensor of the same shape and type as the given array.

    This docstring was adapted from ``numpy.empty_like`` [1]_

    Parameters
    ----------
    other : ArrayLike
        The Tensor or array whose shape and datatype should be mirrored.

    dtype : Optional[DTypeLikeReals]
        Override the data type of the returned Tensor with this value, or None to not override.

    shape : Optional[Union[int, Sequence[int]]]
        If specified, overrides the shape of the result

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation. If ``None`` then:

        Inferred from ``other``, if other is a tensor
        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

    Returns
    -------
    Tensor
        A tensor of uninitialized data whose shape and type match `other`.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.empty_like.html

    See Also
    --------
    empty : Return a new Tensor of the given shape and type, without initializing entries.
    ones : Return a new tensor setting values to one.
    zeros : Return a new tensor setting values to zero.
    full : Return a new tensor of given shape filled with value.

    Examples
    --------
    >>> import mygrad as mg
    >>> x = mg.arange(4).reshape(2, 2)
    >>> mg.empty_like(x, constant=True)
    Tensor([[ -9.74499359e+001,   6.69583040e-309],
            [  2.13182611e-314,   3.06959433e-309]])         #random

    >>> mg.empty_like(x, dtype=int)
    Tensor([[-1073741821, -1067949133],
            [  496041986,    19249760]])                     #random
    """
    constant = _resolve_constant(other, constant=constant)
    return Tensor(
        np.empty_like(_anything_but_tensor(other), dtype=dtype, shape=shape),
        constant=constant,
        copy=False,
    )

def eye(
    N: int,
    M: Optional[int] = None,
    k: int = 0,
    dtype: DTypeLikeReals = float,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """Return a 2D Tensor with ones on the diagonal and zeros elsewhere.

    This docstring was adapted from ``numpy.eye`` [1]_

    Parameters
    ----------
    N : int
        The number of rows in the output Tensor.

    M : int, optional (default=None)
        The number of columns in the output, or None to match `rows`.

    k : int, optional (default=0)
        The index of the diagonal. 0 is the main diagonal; a positive value is the upper
        diagonal, while a negative value refers to the lower diagonal.

    dtype : data-type, optional (default=numpy.float32)
        The data type of the output Tensor.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.eye.html

    Returns
    -------
    Tensor
        A tensor whose elements are 0, except for the :math:`k`-th diagonal, whose values are 1.

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.eye(2, dtype=int)
    Tensor([[1, 0],
            [0, 1]])
    >>> mg.eye(3, k=1)
    Tensor([[ 0.,  1.,  0.],
            [ 0.,  0.,  1.],
            [ 0.,  0.,  0.]])
    """
    return Tensor(
        np.eye(N, M=M, k=k, dtype=dtype),
        constant=constant,
        copy=False,
    )

def identity(
    n: int, dtype: DTypeLikeReals = float, *, constant: Optional[bool] = None
) -> Tensor:
    """Return the identity Tensor; a square Tensor with 1s on the main diagonal and 0s elsewhere.

    This docstring was adapted from ``numpy.identity`` [1]_

    Parameters
    ----------
    n : int
        The number of rows and columns in the output Tensor.

    dtype : data-type, optional (default=numpy.float32)
        The data type of the output Tensor.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor
        A square Tensor whose main diagonal is 1 and all other elements are 0.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.identity.html

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.identity(3)
    Tensor([[ 1.,  0.,  0.],
            [ 0.,  1.,  0.],
            [ 0.,  0.,  1.]])
    """
    return Tensor(np.identity(n, dtype=dtype), constant=constant, copy=False)

def ones(
    shape: Shape, dtype: DTypeLikeReals = np.float32, *, constant: Optional[bool] = None
) -> Tensor:
    """
    Return a Tensor of the given shape and type, filled with ones.

    This docstring was adapted from ``numpy.ones`` [1]_

    Parameters
    ----------
    shape : Union[int, Tuple[int]]
        The shape of the output Tensor.

    dtype : data-type, optional (default=numpy.float32)
        The data type of the output Tensor.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor
        A Tensor of ones with the given shape and data type.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.ones.html

    See Also
    --------
    ones_like : Return an tensor of ones with shape and type of input.
    empty : Return a new uninitialized tensor.
    zeros : Return a new tensor setting values to zero.
    full : Return a new tensor of given shape filled with value.

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.ones(5)
    Tensor([ 1.,  1.,  1.,  1.,  1.])

    >>> mg.ones((5,), dtype=int)
    Tensor([1, 1, 1, 1, 1])

    >>> mg.ones((2, 1))
    Tensor([[ 1.],
           [ 1.]])

    >>> mg.ones((2, 2))
    Tensor([[ 1.,  1.],
            [ 1.,  1.]])
    """
    return Tensor(np.ones(shape, dtype=dtype), constant=constant, copy=False)

@implements_numpy_override()
def ones_like(
    other: ArrayLike,
    dtype: Optional[DTypeLikeReals] = None,
    shape: Optional[Union[int, Sequence[int]]] = None,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """
    Return a Tensor of the same shape and type as the given, filled with ones.

    This docstring was adapted from ``numpy.ones_like`` [1]_

    Parameters
    ----------
    other : array_like
        The Tensor or array whose shape and datatype should be mirrored.

    dtype : Optional[DTypeLikeReals]
        Override the data type of the returned Tensor with this value, or None to not override.

    shape : Optional[Union[int, Sequence[int]]]
        If specified, overrides the shape of the result

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation. If ``None`` then:

        Inferred from ``other``, if other is a tensor
        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

    Returns
    -------
    Tensor
        A Tensor of ones whose shape and data type match `other`.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.ones_like.html

    Examples
    --------
    >>> import mygrad as mg
    >>> x = mg.arange(6).reshape((2, 3))
    >>> x
    Tensor([[0, 1, 2],
            [3, 4, 5]])

    >>> mg.ones_like(x)
    Tensor([[1, 1, 1],
            [1, 1, 1]])

    >>> y = mg.arange(3, dtype=float)
    >>> y
    Tensor([ 0.,  1.,  2.])

    >>> mg.ones_like(y)
    Tensor([ 1.,  1.,  1.])
    """
    constant = _resolve_constant(other, constant=constant)

    return Tensor(
        np.ones_like(_anything_but_tensor(other), dtype=dtype, shape=shape),
        constant=constant,
        copy=False,
    )

def zeros(
    shape: Shape, dtype: DTypeLikeReals = np.float32, *, constant: Optional[bool] = None
) -> Tensor:
    """
    Return a Tensor of the given shape and type, filled with zeros.

    This docstring was adapted from ``numpy.zeros`` [1]_

    Parameters
    ----------
    shape : Union[int, Tuple[int]]
        The shape of the output Tensor.

    dtype : data-type, optional (default=numpy.float32)
        The data type of the output Tensor.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor
        A Tensor of zeros with the given shape and data type.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.zeros.html

    See Also
    --------
    ones_like : Return an tensor of ones with shape and type of input.
    empty : Return a new uninitialized tensor.
    ones : Return a new tensor setting values to one.
    full : Return a new tensor of given shape filled with value.

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.zeros(5)
    Tensor([ 0.,  0.,  0.,  0.,  0.])

    >>> mg.zeros((5,), dtype=int, constant=True) # tensor will not back-propagate a gradient
    Tensor([0, 0, 0, 0, 0])

    >>> mg.zeros((2, 1))
    Tensor([[ 0.],
            [ 0.]])

    >>> mg.zeros((2, 2))
    Tensor([[ 0.,  0.],
            [ 0.,  0.]])
    """
    return Tensor(np.zeros(shape, dtype), constant=constant, copy=False)

@implements_numpy_override()
def zeros_like(
    other: ArrayLike,
    dtype: Optional[DTypeLikeReals] = None,
    shape: Optional[Union[int, Shape]] = None,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """
    Return a Tensor of the same shape and type as the given, filled with zeros.

    This docstring was adapted from ``numpy.zeros_like`` [1]_

    Parameters
    ----------
    other : ArrayLike
        The Tensor or array whose shape and datatype should be mirrored.

    dtype : Optional[DTypeLikeReals]
        Override the data type of the returned Tensor with this value, or None to not override.

    shape : Optional[int, Sequence[int]]
        If specified, overrides the shape of the result

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation. If ``None`` then:

        Inferred from ``other``, if other is a tensor
        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor
        A Tensor of zeros whose shape and data type match `other`.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.zeros_like.html

    See Also
    --------
    empty_like : Return an empty tensor with shape and type of input.
    ones_like : Return an tensor of ones with shape and type of input.
    full_like : Return a new tensor with shape of input filled with value.
    zeros : Return a new tensor setting values to zero.

    Examples
    --------
    >>> import mygrad as mg
    >>> x = mg.arange(6).reshape((2, 3))
    >>> x
    Tensor([[0, 1, 2],
            [3, 4, 5]])

    >>> mg.zeros_like(x, constant=True)  # tensor will not back-propagate a gradient
    Tensor([[0, 0, 0],
            [0, 0, 0]])

    >>> y = mg.arange(3, dtype=float)
    >>> y
    Tensor([ 0.,  1.,  2.])

    >>> mg.zeros_like(y)
    Tensor([ 0.,  0.,  0.])
    """
    constant = _resolve_constant(other, constant=constant)
    return Tensor(
        np.zeros_like(_anything_but_tensor(other), dtype=dtype, shape=shape),
        constant=constant,
        copy=False,
    )

def full(
    shape: Shape,
    fill_value: ArrayLike,
    dtype: Optional[DTypeLikeReals] = None,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """
    Return a Tensor of the given shape and type, filled with `fill_value`.

    This docstring was adapted from ``numpy.full`` [1]_

    Parameters
    ----------
    shape : Union[int, Iterable[int]]
        The shape of the output Tensor.

    fill_value : ArrayLike
        The value with which to fill the output Tensor. Note that this function
        is not differentiable – the resulting tensor will not backprop through
        `fill_value`.

        The value with which to fill the output Tensor.

    dtype : Optional[DTypeLikeReals]
        The data type of the output Tensor, or None to match `fill_value`..

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor
        A Tensor of `fill_value` with the given shape and dtype.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.full.html

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.full((2, 2), 33)
    Tensor([[ 33,  33],
            [ 33,  33]])

    >>> mg.full((2, 2), 10)
    Tensor([[10, 10],
            [10, 10]])
    """
    return Tensor(
        np.full(shape, fill_value=fill_value, dtype=dtype),
        constant=constant,
        copy=False,
    )

@implements_numpy_override()
def full_like(
    other: ArrayLike,
    fill_value: Real,
    dtype: Optional[DTypeLikeReals] = None,
    shape: Optional[Union[int, Shape]] = None,
    constant: Optional[bool] = None,
) -> Tensor:
    """Return a Tensor of the same shape and type as the given, filled with `fill_value`.

    This docstring was adapted from ``numpy.full_like`` [1]_

    Parameters
    ----------
    other : ArrayLike
        The tensor or array whose shape and datatype should be mirrored.

    fill_value : Real
        The value with which to fill the output Tensor.

    dtype : Optional[DTypeLikeReals]
        Override the data type of the returned Tensor with this value, or None to not override.

    shape : Optional[int, Sequence[int]]
        If specified, overrides the shape of the result

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation. If ``None`` then:

        Inferred from ``other``, if other is a tensor
        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

    Returns
    -------
    Tensor
        A Tensor of `fill_value` whose shape and data type match `other`.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.full_like.html

    Examples
    --------
    >>> import mygrad as mg
    >>> x = mg.arange(6, dtype=int)
    >>> mg.full_like(x, 1)
    Tensor([1, 1, 1, 1, 1, 1])
    >>> mg.full_like(x, 0.1)
    Tensor([0, 0, 0, 0, 0, 0])
    >>> mg.full_like(x, 0.1, dtype=np.double)
    Tensor([ 0.1,  0.1,  0.1,  0.1,  0.1,  0.1])
    >>> mg.full_like(x, np.nan, dtype=np.double)
    Tensor([ nan,  nan,  nan,  nan,  nan,  nan])

    >>> y = mg.arange(6, dtype=np.double)
    >>> mg.full_like(y, 0.1)
    Tensor([ 0.1,  0.1,  0.1,  0.1,  0.1,  0.1])
    """
    constant = _resolve_constant(other, constant=constant)

    return Tensor(
        np.full_like(
            _anything_but_tensor(other),
            fill_value=_anything_but_tensor(fill_value),
            dtype=dtype,
            shape=shape,
        ),
        constant=constant,
        copy=False,
    )

def arange(
    *args,
    constant: Optional[bool] = None,
    **kwargs,
) -> Tensor:
    """
    arange([start,] stop[, step,], dtype=None, *, constant=None)

    Return a Tensor with evenly-spaced values within a given interval.

    Values are generated within [start, stop). Note that for non-integer steps, results may be
    inconsistent; you are better off using `linspace` instead.

    This docstring was adapted from ``numpy.arange`` [1]_

    Parameters
    ----------
    start : Real, optional, default=0
        The start of the interval, inclusive.

    stop : Real
        The end of the interval, exclusive.

    step : int, optional (default=1)
        The spacing between successive values.

    dtype : Optional[DTypeLikeReals]
        The data type of the output Tensor, or None to infer from the inputs.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor
        A Tensor of evenly-spaced values in [start, end).

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.arange.html

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.arange(3)
    Tensor([0, 1, 2])
    >>> mg.arange(3.0, constant=True)  # resulting tensor will not back-propagate a gradient
    Tensor([ 0.,  1.,  2.])
    >>> mg.arange(3,7)
    Tensor([3, 4, 5, 6])
    >>> mg.arange(3,7,2)
    Tensor([3, 5])
    """
    return Tensor(np.arange(*args, **kwargs), constant=constant, copy=False)

def linspace(
    start: ArrayLike,
    stop: ArrayLike,
    num: int = 50,
    endpoint: bool = True,
    dtype: Optional[DTypeLikeReals] = None,
    axis: int = 0,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """Return a Tensor with evenly-spaced numbers over a specified interval.

    Values are generated within [start, stop], with the endpoint optionally excluded.

    This docstring was adapted from ``numpy.linspace`` [1]_

    Parameters
    ----------
    start : ArrayLike
        The starting value of the sequence, inclusive.

    stop : ArrayLike
        The ending value of the sequence, inclusive unless `include_endpoint` is False.

    num : int, optional (default=50)
        The number of values to generate. Must be non-negative.

    endpoint : bool, optional (default=True)
        Whether to include the endpoint in the Tensor. Note that if False, the step size changes
        to accommodate the sequence excluding the endpoint.

    dtype : Optional[DTypeLikeReals]
        The data type of the output Tensor, or None to infer from the inputs.

    axis : int, optional (default=0)
        The axis in the result to store the samples - for array-like start/stop.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.linspace.html

    See Also
    --------
    arange : Similar to `linspace`, but uses a step size (instead of the
             number of samples).
    logspace : Samples uniformly distributed in log space.

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.linspace(2.0, 3.0, num=5)
    Tensor([ 2.  ,  2.25,  2.5 ,  2.75,  3.  ])
    >>> mg.linspace(2.0, 3.0, num=5, endpoint=False)
    Tensor([ 2. ,  2.2,  2.4,  2.6,  2.8])
    """
    return Tensor(
        np.linspace(
            start,
            stop,
            num,
            endpoint=endpoint,
            dtype=dtype,
            axis=axis,
        ),
        constant=constant,
        copy=False,
    )

def logspace(
    start: ArrayLike,
    stop: ArrayLike,
    num: int = 50,
    endpoint: bool = True,
    base: Real = 10,
    dtype: Optional[DTypeLikeReals] = None,
    axis: int = 0,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """Return a Tensor with evenly-spaced numbers over a specified interval on a log scale.
    This is not a differentiable function - it does not propagate gradients to its inputs.

    In linear space, values are generated within [base**start, base**stop], with the endpoint
    optionally excluded.

    This docstring was adapted from ``numpy.logspace`` [1]_

    Parameters
    ----------
    start : ArrayLike
        The starting value of the sequence, inclusive; start at `base ** start`.

    stop : ArrayLike
        The ending value of the sequence, inclusive unless `include_endpoint` is False; end at
        `base ** stop`.

    num : int, optional (default=50)
        The number of values to generate. Must be non-negative.

    endpoint : bool, optional (default=True)
        Whether to include the endpoint in the Tensor. Note that if False, the step size changes
        to accommodate the sequence excluding the endpoint.

    base : Real, optional (default=10)
        The base of the log space.

    dtype : Optional[DTypeLikeReals]
        The data type of the output Tensor, or None to infer from the inputs.

    axis : int, optional (default=0)
        The axis in the result to store the samples - for array-like start/stop.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor

    See Also
    --------
    arange : Similar to linspace, with the step size specified instead of the
             number of samples. Note that, when used with a float endpoint, the
             endpoint may or may not be included.
    linspace : Similar to logspace, but with the samples uniformly distributed
               in linear space, instead of log space.
    geomspace : Similar to logspace, but with endpoints specified directly.

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.logspace.html

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.logspace(2.0, 3.0, num=4)
    Tensor([  100.        ,   215.443469  ,   464.15888336,  1000.        ])
    >>> mg.logspace(2.0, 3.0, num=4, endpoint=False)
    Tensor([ 100.        ,  177.827941  ,  316.22776602,  562.34132519])
    >>> mg.logspace(2.0, 3.0, num=4, base=2.0)
    Tensor([ 4.        ,  5.0396842 ,  6.34960421,  8.        ])

    """
    return Tensor(
        np.logspace(
            start=start,
            stop=stop,
            num=num,
            endpoint=endpoint,
            base=base,
            dtype=dtype,
            axis=axis,
        ),
        constant=constant,
        copy=False,
    )

def geomspace(
    start: ArrayLike,
    stop: ArrayLike,
    num=50,
    endpoint=True,
    dtype=None,
    axis=0,
    *,
    constant: Optional[bool] = None,
) -> Tensor:
    """Return a Tensor with evenly-spaced values in a geometric progression.

    Each output sample is a constant multiple of the previous output.

    This docstring was adapted from ``numpy.geomspace`` [1]_

    Parameters
    ----------
    start : ArrayLike
        The starting value of the output.

    stop : ArrayLike
        The ending value of the sequence, inclusive unless `endpoint` is false.

    num : int, optional (default=50)
        The number of values to generate. Must be non-negative.

    endpoint : bool, optional (default=True)
        Whether to include the endpoint in the Tensor. Note that if False, the step size changes
        to accommodate the sequence excluding the endpoint.

    dtype : Optional[DTypeLikeReals]
        The data type of the output Tensor, or None to infer from the inputs.

    axis : int, optional (default=0)
        The axis in the result to store the samples - for array-like start/stop.

    constant : Optional[bool]
        If ``True``, this tensor is a constant, and thus does not facilitate
        back propagation.

        Defaults to ``False`` for float-type data.
        Defaults to ``True`` for integer-type data.

        Integer-type tensors must be constant.

    Returns
    -------
    Tensor

    References
    ----------
    .. [1] Retrieved from https://numpy.org/doc/stable/reference/generated/numpy.geomspace.html

    See Also
    --------
    logspace : Similar to geomspace, but with endpoints specified using log
               and base.
    linspace : Similar to geomspace, but with arithmetic instead of geometric
               progression.
    arange : Similar to linspace, with the step size specified instead of the
             number of samples.

    Examples
    --------
    >>> import mygrad as mg
    >>> mg.geomspace(1, 1000, num=4)
    Tensor([    1.,    10.,   100.,  1000.])
    >>> mg.geomspace(1, 1000, num=3, endpoint=False)
    Tensor([   1.,   10.,  100.])
    >>> mg.geomspace(1, 1000, num=4, endpoint=False)
    Tensor([   1.        ,    5.62341325,   31.6227766 ,  177.827941  ])
    >>> mg.geomspace(1, 256, num=9)
    Tensor([   1.,    2.,    4.,    8.,   16.,   32.,   64.,  128.,  256.])

    Note that the above may not produce exact integers:

    >>> mg.geomspace(1, 256, num=9, dtype=int)
    Tensor([  1,   2,   4,   7,  16,  32,  63, 127, 256])
    >>> np.around(mg.geomspace(1, 256, num=9).data).astype(int)
    array([  1,   2,   4,   8,  16,  32,  64, 128, 256])

    Negative, and decreasing inputs are allowed:

    >>> mg.geomspace(1000, 1, num=4)
    Tensor([ 1000.,   100.,    10.,     1.])
    >>> mg.geomspace(-1000, -1, num=4)
    Tensor([-1000.,  -100.,   -10.,    -1.])
    """
    return Tensor(
        np.geomspace(
            start=start,
            stop=stop,
            num=num,
            endpoint=endpoint,
            dtype=dtype,
            axis=axis,
        ),
        constant=constant,
        copy=False,
    )