""" Test grader module
"""

from os.path import join as pjoin, dirname
from io import StringIO
import re
from hashlib import sha1
from glob import glob
from copy import deepcopy

import numpy as np

from rnbgrader import JupyterKernel
from rnbgrader.grader import (OPTIONAL_PROMPT, MARK_MARKUP_RE, NBRunner,
                              report, duplicates, Grader, CanvasGrader,
                              NotebookError, CachedBuiltNotebook)
from rnbgrader.answers import RegexAnswer, ImgAnswer, raw2regex, RawRegexAnswer

import pytest

from gradools.canvastools import CanvasError

DATA = pjoin(dirname(__file__), 'data')
MB_NB_FN = 'brettmatthew_139741_6519327_some_name.Rmd'
VR2_NB_FN = 'rodriguezvalia_140801_6518299_notebook.rmd'

def test_optional_prompt():
    assert re.search(OPTIONAL_PROMPT, '[1] ') is not None
    assert re.search(OPTIONAL_PROMPT, '[100] ') is not None
    assert re.search(OPTIONAL_PROMPT, ' [100] ') is not None
    assert re.search(OPTIONAL_PROMPT + 'here', '[100] here') is not None
    assert re.search(OPTIONAL_PROMPT + 'here', '[100] here') is not None

def test_report():
    runner = NBRunner()
    nb_fileobj0 = StringIO("""
Text

```{r}
first_var <- 1
```

More text

```{r}
first_var
```
""")
    nb_fileobj1 = StringIO("""
Text

```{r}
```

More text

```{r}
first_var <- 2
```
""")
    with JupyterKernel('ir') as rk:
        results0 = runner.run(nb_fileobj0, rk)
        results1 = runner.run(nb_fileobj1, rk)
    assert (report(results0) ==
            ' 0: first_var <- 1 - None\n 1: first_var - [1] 1')
    assert (report(results1) ==
            ' 0: (no code) - None\n 1: first_var <- 2 - None')

def test_duplicates():
    fnames = glob(pjoin(DATA, 'test_submissions', '*.Rmd'))
    with open(fnames[0], 'rb') as fobj:
        hash = sha1(fobj.read()).hexdigest()
    hashes = duplicates(glob(pjoin(DATA, 'test_submissions', '*')))
    assert list(hashes) == [hash]
    assert sorted(hashes[hash]) == sorted(fnames)

def test_get_submissions():
    g = CanvasGrader()
    pth = pjoin(DATA, 'test_submissions2')
    fnames = sorted(glob(pjoin(pth, '*')))
    assert g.get_submissions(pth) == fnames

def test_get_submissions_same_id():
    g = CanvasGrader()
    with pytest.raises(CanvasError):
        g.get_submissions(pjoin(DATA, 'test_submissions'))

class Skip:
    """ Flag to indicate we should skip this answer in compiling answers
    """
    pass

class CarsGrader(CanvasGrader):

    solution_rmds = (pjoin(DATA, 'solution.Rmd'),)
    standard_box = (44, 81, 800, 770)
    total = 50

    # Solution chunk positions, used in make_answers.
    # This allows testing different chunk position specifications.
    # Here we use simple positions.  Use Skip class to indicate we should
    # skip this question in the specification.
    _positions = tuple(range(1, 8))

    def make_answers(self):
        solution_dir = self.solution_dirs[0]
        # Positions are class variable to allow for testing different position
        # specifications.
        ps = self._positions

        if ps[0] is not Skip:
            self._chk_answer(RegexAnswer(
                5,
                OPTIONAL_PROMPT + r'50  2'),
                ps[0])

        raw = """
            'data.frame':	50 obs. of  2 variables:
            $ speed: num  4 4 7 7 8 9 10 10 10 11 ...
            $ dist : num  2 10 4 22 16 10 18 26 34 17 ..."""

        if ps[1] is not Skip:
            self._chk_answer(RegexAnswer(5, raw2regex(raw)), ps[1])

        raw = """
            speed dist
            1 4      2
            2 4     10
            3 7      4
            4 7     22
            5 8     16
            6 9     10"""
        if ps[2] is not Skip:
            self._chk_answer(RegexAnswer(5, raw2regex(raw)), ps[2])

        if ps[3] is not Skip:
            self._chk_answer(ImgAnswer(10,
                pjoin(solution_dir, 'chunk-4_item-0.png'),
                self.standard_box), ps[3])

        raw = """
            4  7  8  9 10 11 12 13 14 15 16 17 18 19 20 22 23 24 25
            2  2  1  1  3  2  4  4  4  3  2  3  4  3  5  1  1  4  1"""
        if ps[4] is not Skip:
            self._chk_answer(RawRegexAnswer(5, raw), ps[4])

        raw = """
        speed dist
        27    16   32
        28    16   40
        29    17   32
        30    17   40
        31    17   50
        32    18   42"""
        if ps[5] is not Skip:
            self._chk_answer(RegexAnswer(10, raw2regex(raw)), ps[5])

        if ps[6] is not Skip:
            self._chk_img_answer(10, ps[6])

CARS_GRADER = CarsGrader()

def test_solutions_with_offsets():
    soln_fname = pjoin(DATA, 'solution.Rmd')
    # Basic position specification.
    assert sum(CARS_GRADER.grade_notebook(soln_fname)) == 50

    class G2(CarsGrader):
        # Strings for integer positions are OK.
        _positions = [str(i) for i in range(1, 8)]

    assert sum(G2().grade_notebook(soln_fname)) == 50

    class G3(CarsGrader):
        # Offsets OK, as long as there is a base.
        _positions = [1] + ['+1'] * 6

    assert sum(G3().grade_notebook(soln_fname)) == 50

    class G4(CarsGrader):
        # There must be a base.
        _positions = ['+1'] * 7

    with pytest.raises(ValueError):
        G4().grade_notebook(soln_fname)

    class G5(CarsGrader):
        # Can mix positions and offsets.
        _positions = [1, '+1', 3, 4, '+1', 6, '+1']

    assert sum(G5().grade_notebook(soln_fname)) == 50

    class G6(CarsGrader):
        # Can have offsets > 1
        _positions = [1, Skip, '+2', Skip, 5, Skip, '+2']
        total = 25

    assert sum(G6().grade_notebook(soln_fname)) == 25

def test_bit_bad():
    # This one has a couple of wrong answers
    assert sum(CARS_GRADER.grade_notebook(
        pjoin(DATA, 'not_solution.Rmd'))) == 35

def test_grade_all_error():
    with pytest.raises(CanvasError):
        CARS_GRADER.grade_all_notebooks(pjoin(DATA, 'test_submissions'))

def test_main():
    args = ["foo"]
    assert CARS_GRADER.main(args) == 1

def test_check_names():
    args = ["check-names", pjoin(DATA, "test_submissions")]
    with pytest.raises(CanvasError):
        CARS_GRADER.main(args)

def test_error_report():
    nb = StringIO("""

Some text.

```{r}
a <- 1
a
```

More text.

```{r}
b
```
""")
    runner = NBRunner()
    with pytest.raises(NotebookError):
        with JupyterKernel('ir') as rk:
            runner.run(nb, rk)

def test_mark_markup():
    assert MARK_MARKUP_RE.match('#M: -2.5').groups() == ('-2.5',)
    assert MARK_MARKUP_RE.match('#M:-2.5').groups() == ('-2.5',)
    assert MARK_MARKUP_RE.match('# M : -2.5').groups() == ('-2.5',)
    assert MARK_MARKUP_RE.search('foo\n# M : -2.5  \nbar').groups() == ('-2.5',)
    assert MARK_MARKUP_RE.search('foo  # M : -2.5  \nbar') is None
    assert MARK_MARKUP_RE.match('#M : -2.5  ').groups() == ('-2.5',)
    assert MARK_MARKUP_RE.match('#M : +2.5  ').groups() == ('+2.5',)
    assert MARK_MARKUP_RE.match('#M : +22. ').groups() == ('+22.',)
    assert MARK_MARKUP_RE.match('#M : 11.999 ').groups() == ('11.999',)
    assert MARK_MARKUP_RE.match('\t#M : -2.5  ').groups() == ('-2.5',)
    assert MARK_MARKUP_RE.match('#M: --2.5') is None
    assert MARK_MARKUP_RE.match('#M: +-2.5') is None
    assert MARK_MARKUP_RE.match('#M: ++2.5') is None
    assert MARK_MARKUP_RE.match('#M: 2.5 ish') is None

def test_initial_check():
    g = CanvasGrader()
    with pytest.raises(CanvasError):
        g.initial_check(pjoin(DATA, 'test_submissions'))
    with pytest.raises(NotebookError):
        g.initial_check(pjoin(DATA, 'test_submissions_markup'))
    pth = pjoin(DATA, 'test_submissions2')
    res = g.initial_check(pth)
    mb = pjoin(pth, MB_NB_FN)
    with open(mb, 'rb') as fobj:
        mb_sha = sha1(fobj.read()).hexdigest()
    assert res == {mb_sha: [mb, pjoin(pth, VR2_NB_FN)]}

def test_markup_in_nb():
    bare_nb = StringIO("""

Some text.

```{r}
a <- 1
a
```

More text.

#M: 10

```{r}
b <- 2
```
""")
    assert CARS_GRADER.mark_markups(bare_nb) == ()

    annotated_nb = StringIO("""

Some text.

```{r}
a <- 1
a
# M: 2.0
```

More text.

#M: 10

```{r}
#M : -2.5
b <- 2
```
""")

    assert CARS_GRADER.mark_markups(annotated_nb) == (2., -2.5)

def test_raise_for_markup():
    g = Grader()
    for sdir in ('test_submissions', 'test_submissions2'):
        pth = pjoin(DATA, sdir)
        g.raise_for_markup(g.get_submissions(pth))

def test_markup_used():
    g = CARS_GRADER
    pth = pjoin(DATA, 'test_submissions_markup')
    mb = pjoin(pth, MB_NB_FN)
    vr2 = pjoin(pth, VR2_NB_FN)
    assert g.mark_markups(mb) == (-2, 42)
    assert g.mark_markups(vr2) == ()
    mb_marks = g.grade_notebook(mb)
    assert list(mb_marks.index) == ['unnamed'] * 7 + ['adjustments', 'markups']
    assert sum(mb_marks) == 80
    assert sum(g.grade_notebook(vr2)) == 40

def test_cached_nb_file_like():
    # Test we can used file-likes for notebook caching
    nb_text = """
Text

```{r}
#- A question
first_var <- 99
```

More text

```{r}
first_var
```

```{r}
#- Actual answer
first_var
```
"""
    nb_fobj = StringIO(nb_text)
    cbn = CachedBuiltNotebook(nb_fobj, NBRunner())
    assert len(cbn.solution) == 3
    assert cbn.solution[-1].chunk.code == '#- Actual answer\nfirst_var\n'

    class MyG(Grader):

        solution_rmds = (nb_fobj,)

    g = MyG()
    solns = g.solutions
    assert len(solns) == 1
    assert len(solns[0]) == 3

def assert_seq_equal(s1, s2):
    assert tuple(s1) == tuple(s2)

def test_chunk_is_answer():
    runner = NBRunner()
    nb_text = """
Text

```{r}
#- A question
first_var <- 99
```

More text

```{r}
first_var
```

```{r}
#- Actual answer
first_var
```
"""
    nb_fobj = StringIO(nb_text)
    with JupyterKernel('ir') as rk:
        chunks = runner.run(nb_fobj, rk)
    g = Grader()
    assert_seq_equal(chunks, g.clear_not_answers(chunks))
    # Second chunk has results.
    assert len(chunks[1].results) == 1

    # First check case where not-answer present raises error
    nb_fobj.seek(0)

    class MyG(Grader):

        solution_rmds = (nb_fobj,)
        total = 5

        def make_answers(self):
            self._chk_answer(RegexAnswer(
                5,
                OPTIONAL_PROMPT + r'99'),
                2)

    g = MyG()
    # Doesn't clear any chunks.
    assert_seq_equal(chunks, g.clear_not_answers(chunks))
    assert len(chunks[1].results) == 1
    # Therefore errors for duplicate outputs
    with pytest.raises(NotebookError):
        g.grade_notebook(StringIO(nb_text))

    # Check case with answer removed.
    class MyG2(MyG):

        solution_rmds = (nb_fobj,)

        def chunk_is_answer(self, chunk):
            return chunk.chunk.code != 'first_var\n'

    g2 = MyG2()
    # Clears output for not-answer chunk.
    cleared = deepcopy(chunks[1])
    cleared.results = []
    assert_seq_equal([chunks[0], cleared, chunks[2]], g2.clear_not_answers(chunks))
    # Answers now not duplicated.
    assert np.all(np.array(g2.grade_notebook(StringIO(nb_text))) ==
                  [5, 0, 0])


"""
Accurate timing information for Sage commands

This is an implementation of nice timeit functionality, like the
``%timeit`` magic command in IPython.  To use it, use the ``timeit``
command. This command then calls :func:`sage_timeit`, which you can
find below.

EXAMPLES::

    sage: timeit('1+1')    # random output
    625 loops, best of 3: 314 ns per loop

AUTHOR:

    -- William Stein, based on code by Fernando Perez included in IPython
"""

class SageTimeitResult():
    r"""
    Represent the statistics of a timeit() command.

    Prints as a string so that it can be easily returned to a user.

    INPUT:

    - ``stats`` -- tuple of length 5 containing the following information:

        - integer, number of loops
        - integer, repeat number
        - Python integer, number of digits to print
        - number, best timing result
        - str, time unit

    EXAMPLES::

        sage: from sage.misc.sage_timeit import SageTimeitResult
        sage: SageTimeitResult( (3, 5, int(8), pi, 'ms') )
        3 loops, best of 5: 3.1415927 ms per loop

    ::

        sage: units = ["s", "ms", "\xc2\xb5s", "ns"]
        sage: scaling = [1, 1e3, 1e6, 1e9]
        sage: number = 7
        sage: repeat = 13
        sage: precision = int(5)
        sage: best = pi / 10 ^ 9
        sage: order = 3
        sage: stats = (number, repeat, precision, best * scaling[order], units[order])
        sage: SageTimeitResult(stats)
        7 loops, best of 13: 3.1416 ns per loop

    If the third argument is not a Python integer, a ``TypeError`` is raised::

        sage: SageTimeitResult( (1, 2, 3, 4, 's') )
        Traceback (most recent call last):
        ...
        TypeError: * wants int

    """
    def __init__(self, stats, series=None):
        r"""
        Construction of a timing result.

        See documentation of ``SageTimeitResult`` for more details and
        examples.

        EXAMPLES::

            sage: from sage.misc.sage_timeit import SageTimeitResult
            sage: SageTimeitResult( (3, 5, int(8), pi, 'ms') )
            3 loops, best of 5: 3.1415927 ms per loop
            sage: s = SageTimeitResult( (3, 5, int(8), pi, 'ms'), [1.0,1.1,0.5])
            sage: s.series
            [1.00000000000000, 1.10000000000000, 0.500000000000000]
        """
        self.stats = stats
        self.series = series if not None else []

    def __repr__(self):
        r"""
        String representation.

        EXAMPLES::

            sage: from sage.misc.sage_timeit import SageTimeitResult
            sage: stats = (1, 2, int(3), pi, 'ns')
            sage: SageTimeitResult(stats)           #indirect doctest
            1 loops, best of 2: 3.14 ns per loop
        """
        return "%d loops, best of %d: %.*g %s per loop" % self.stats

def sage_timeit(stmt, globals_dict=None, preparse=None, number=0, repeat=3, precision=3, seconds=False):
    """
    Accurately measure the wall time required to execute ``stmt``.

    INPUT:

    - ``stmt`` -- a text string.

    - ``globals_dict`` -- a dictionary or ``None`` (default). Evaluate
      ``stmt`` in the context of the globals dictionary. If not set,
      the current ``globals()`` dictionary is used.

    - ``preparse`` -- (default: use globals preparser default) if
      ``True`` preparse ``stmt`` using the Sage preparser.

    - ``number`` -- integer, (optional, default: 0), number of loops.

    - ``repeat`` -- integer, (optional, default: 3), number of
      repetition.

    - ``precision`` -- integer, (optional, default: 3), precision of
      output time.

    - ``seconds`` -- boolean (default: ``False``). Whether to just
      return time in seconds.

    OUTPUT:

    An instance of ``SageTimeitResult`` unless the optional parameter
    ``seconds=True`` is passed. In that case, the elapsed time in
    seconds is returned as a floating-point number.

    EXAMPLES::

        sage: from sage.misc.sage_timeit import sage_timeit
        sage: sage_timeit('3^100000', globals(), preparse=True, number=50)      # random output
        '50 loops, best of 3: 1.97 ms per loop'
        sage: sage_timeit('3^100000', globals(), preparse=False, number=50)     # random output
        '50 loops, best of 3: 67.1 ns per loop'
        sage: a = 10
        sage: sage_timeit('a^2', globals(), number=50)                            # random output
        '50 loops, best of 3: 4.26 us per loop'

    If you only want to see the timing and not have access to additional
    information, just use the ``timeit`` object::

        sage: timeit('10^2', number=50)
        50 loops, best of 3: ... per loop

    Using sage_timeit gives you more information though::

        sage: s = sage_timeit('10^2', globals(), repeat=1000)
        sage: len(s.series)
        1000
        sage: mean(s.series)   # random output
        3.1298141479492283e-07
        sage: min(s.series)    # random output
        2.9258728027343752e-07
        sage: t = stats.TimeSeries(s.series)
        sage: t.scale(10^6).plot_histogram(bins=20,figsize=[12,6], ymax=2)

    The input expression can contain newlines (but doctests cannot, so
    we use ``os.linesep`` here)::

        sage: from sage.misc.sage_timeit import sage_timeit
        sage: from os import linesep as CR
        sage: # sage_timeit(r'a = 2\\nb=131\\nfactor(a^b-1)')
        sage: sage_timeit('a = 2' + CR + 'b=131' + CR + 'factor(a^b-1)',
        ...               globals(), number=10)
        10 loops, best of 3: ... per loop

    Test to make sure that ``timeit`` behaves well with output::

        sage: timeit("print 'Hi'", number=50)
        50 loops, best of 3: ... per loop

    If you want a machine-readable output, use the ``seconds=True`` option::

        sage: timeit("print 'Hi'", seconds=True)   # random output
        1.42555236816e-06
        sage: t = timeit("print 'Hi'", seconds=True)
        sage: t     #r random output
        3.6010742187499999e-07

    TESTS:

    Make sure that garbage collection is re-enabled after an exception
    occurs in timeit::

        sage: def f(): raise ValueError
        sage: import gc
        sage: gc.isenabled()
        True
        sage: timeit("f()")
        Traceback (most recent call last):
        ...
        ValueError
        sage: gc.isenabled()
        True
    """
    import time, math
    import timeit as timeit_

    import preparser, sage.repl.interpreter as interpreter

    number=int(number)
    repeat=int(repeat)
    precision=int(precision)
    if preparse is None:
        preparse = interpreter._do_preparse
    if preparse:
        stmt = preparser.preparse(stmt)
    if stmt == "":
        return ''

    units = ["s", "ms", "\xc2\xb5s", "ns"]
    scaling = [1, 1e3, 1e6, 1e9]

    timer = timeit_.Timer()

    # this code has tight coupling to the inner workings of timeit.Timer,
    # but is there a better way to achieve that the code stmt has access
    # to the shell namespace?

    src = timeit_.template % {'stmt': timeit_.reindent(stmt, 8),
                             'setup': "pass"}
    code = compile(src, "<magic-timeit>", "exec")
    ns = {}
    if not globals_dict:
        globals_dict = globals()
    exec code in globals_dict, ns
    timer.inner = ns["inner"]

    try:
        import sys
        f = sys.stdout
        sys.stdout = open('/dev/null', 'w')

        if number == 0:
            # determine number so that 0.2 <= total time < 2.0
            number = 1
            for i in range(1, 5):
                number *= 5
                if timer.timeit(number) >= 0.2:
                    break

        series = [s/number for s in timer.repeat(repeat, number)]
        best = min(series)

    finally:
        sys.stdout.close()
        sys.stdout = f
        import gc
        gc.enable()

    if seconds:
        return best

    if best > 0.0:
        order = min(-int(math.floor(math.log10(best)) // 3), 3)
    else:
        order = 3
    stats = (number, repeat, precision, best * scaling[order], units[order])
    return SageTimeitResult(stats,series=series)


import matplotlib.pyplot as plt
import numpy as np

def plot_data(X, y):
    plt.figure()

    # ===================== Your Code Here =====================
    # Instructions : Plot the positive and negative examples on a
    #                2D plot, using the marker="+" for the positive
    #                examples and marker="o" for the negative examples
    #


## Brian Blaylock
## October 1, 2021

"""
I have no idea what I'm doing yet. First time making a test.
"""
import numpy as np

def test_sqrt():
   num = 25
   assert np.sqrt(num) == 5

def test_str():
   a = "Herbie"
   b = "Herbie"
   assert a==b


import time
import pandas as pd
import numpy as np

CITY_DATA = { 'chicago': 'chicago.csv',
              'new york city': 'new_york_city.csv',
              'washington': 'washington.csv' }
months = ['january', 'february', 'march', 'april', 'may', 'june','all']
days_of_week = ['Sunday','Monday','Tuesday','Wednesday','Thursday','Friday','Saturday','All']

def get_filters():
    """
    Asks user to specify a city, month, and day to analyze.

    Returns:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    """
    print('Hello! Let\'s explore some US bikeshare data!')
    # user input for city (chicago, new york city, washington).
    citylist = list(CITY_DATA.keys())
    print("would you like to see data for",citylist)
    while True:
        city = input().lower()
        if city not in CITY_DATA:
            print("please try again")
            continue
        else:
            break

    # user input for month (all, january, february, ... , june)
    print("which month?",months)
    while True:
        month = input().lower()
        if month not in months:
            print("please try again")
            continue
        else:
            break

    # user input for day of week (all, monday, tuesday, ... sunday)
    print("which day?",days_of_week)
    while True:
        day = input().title()
        if day not in days_of_week:
            print("please try again")
            continue
        else:
            break

    print('-'*40)
    return city, month, day

def load_data(city, month, day):
    """
    Loads data for the specified city and filters by month and day if applicable.

    Args:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    Returns:
        df - Pandas DataFrame containing city data filtered by month and day
    """
    df = pd.read_csv(CITY_DATA[city])
    # convert the Start Time column to datetime
    df['Start Time'] = pd.to_datetime(df['Start Time'])

    # extract month and day of week from Start Time to create new columns
    df['month'] = df['Start Time'].dt.month
    df['day_of_week'] = df['Start Time'].dt.weekday_name

    # filter by month if applicable
    if month != 'all':
        # use the index of the months list to get the corresponding int
        month = months.index(month) + 1

        # filter by month to create the new dataframe
        df = df[df['month'] == month]

    # filter by day of week if applicable
    if day != 'All':
        # filter by day of week to create the new dataframe
        df = df[df['day_of_week'] == day.title()]

    return df

def time_stats(df):
    """Displays statistics on the most frequent times of travel."""

    print('\nCalculating The Most Frequent Times of Travel...\n')
    start_time = time.time()

    #  the most common month
    popular_month = df['month'].mode()[0]

    # the most common day of week
    popular_day = df['day_of_week'].mode()[0]

    # the most common start hour

    # convert the Start Time column to datetime
    df['Start Time'] = pd.to_datetime(df['Start Time'])

    # extract hour from the Start Time column to create an hour column
    df['hour'] = df['Start Time'].dt.hour

    # find the most popular hour
    popular_hour = df['hour'].mode()[0]

    print("Most Popular month:",popular_month)
    print("Most Popular day:", popular_day)
    print("Most Popular Start Hour:", popular_hour)
    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def station_stats(df):
    """Displays statistics on the most popular stations and trip."""

    print('\nCalculating The Most Popular Stations and Trip...\n')
    start_time = time.time()

    #  display most commonly used start station
    popular_start_station = df['Start Station'].mode()[0]

    #  display most commonly used end station
    popular_end_station = df['End Station'].mode()[0]

    #  display most frequent combination of start station and end station trip
    popular_combination = df.groupby(['Start Station','End Station']).size().idxmax()
    print("Most Popular Start Station:",popular_start_station)
    print("Most Popular End Station:",popular_end_station)
    print("Most Popular trip:",popular_combination)

    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def trip_duration_stats(df):
    """Displays statistics on the total and average trip duration."""

    print('\nCalculating Trip Duration...\n')
    start_time = time.time()

    #  display total travel time
    total_travel_time = df['Trip Duration'].sum()

    # display mean travel time
    mean_travel_time = df['Trip Duration'].mean()

    print("total travel time:",total_travel_time)
    print("the avreage travel time:",mean_travel_time)
    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def user_stats(df):
    """Displays statistics on bikeshare users."""
    print('\nCalculating User Stats...\n')
    start_time = time.time()
    # Display counts of user types
    user_types = df['User Type'].value_counts()
    # Display counts of gender
    if 'Gender' in df:
        gender_types = df['Gender'].value_counts()
        print('count of males and females is :' , gender_types)
    else:
        print('there is no gender data for this city')
    # Display earliest, most recent, and most common year of birth
    if 'Birth Year' in df:
        earliest = int(df['Birth Year'].min())
        recent = int(df['Birth Year'].max())
        popular = int(df['Birth Year'].mode()[0])
        print("the earlies birth was in",earliest ,"and the most recent birth was in",recent,"the most common birth was in",popular)
    else:
        print('there is no Birth Year data for this city')

    print("\nThis took %s seconds." % (time.time() - start_time))
    print('-'*40)

def data(df):
    raw_input = 0

    while True:
        # to ask the user if he wants to see the raw data or not
        choice = input('do you want to see five rows of raw data ? please type yes if you want to and no If you don\'t: ').lower()
        if choice == 'no':
            break
        elif choice == 'yes':
            raw_input += 5
            print(df.iloc[raw_input : raw_input + 5])
            print(choice)

        else:
            print('Wrong input')

def main():
    city = ()
    month = ()
    day = ()
    while True:
        city, month, day = get_filters()
        df = load_data(city, month, day)

        time_stats(df)
        station_stats(df)
        trip_duration_stats(df)
        user_stats(df)
        data(df)

        restart = input('\nWould you like to restart? Enter yes or no.\n')
        if restart.lower() != 'yes':
            break

if __name__ == "__main__":
	main()


"Slicing unit tests"

"""
NOTES

Indexing with nd boolean array fails:

>>> lar[lar>0]
Traceback (most recent call last):
  File "la\deflarry.py", line 722, in __getitem__
    raise IndexError, msg
IndexError: Only slice, integer, and seq (list, tuple, 1d array) indexing supported
lar.x[lar.x>0]  # this works but creates 1d array, only useful for changes

Comment: Note that the numpy example (lar.x[lar.x>0]) creates a 1d array. If
larry would allow this kind of indexing what would it do with the 2d labels
when the output is 1d? Should each label element become a tuple?

Indexing with 1d boolean array fails:

>>> la1_3d[lar.x[:,0,0]>0,:,:]
Traceback (most recent call last):
  File "la\deflarry.py", line 725, in __getitem__
    return type(self)(x, label)
  File "la\deflarry.py", line 60, in __init__
    raise ValueError, msg2 % (i, value, key)
ValueError: Elements of label not unique along dimension 0. There are 2 labels named `1`.
>>> lar.x[lar.x[:,0,0]>0,:,:] # this works

Comment: Yes, you can't create a larry that doesn't have unique label elements
along the axis. It works if the resulting labels are unique.

List slicing fails:

>>> lar.x[[0,1], [0,1], [0,1]]
array([ 2.,  4.])
>>> lar[[0,1], [0,1], [0,1]]
Traceback (most recent call last):
  File "la\deflarry.py", line 725, in __getitem__
    return type(self)(x, label)
  File "la\deflarry.py", line 50, in __init__
    if x.shape[i] != nlabel:
IndexError: tuple index out of range

Broadcasting of slices fails in larry:

>>> lar.x[(np.array([0,1])[:,None], [0,1], [0,1])]
array([[ 2.,  2.],
       [ 4.,  4.]])
>>> lar[(np.array([0,1])[:,None], [0,1], [0,1])]
Traceback (most recent call last):
  File "la\deflarry.py", line 725, in __getitem__
    return type(self)(x, label)
  File "la\deflarry.py", line 50, in __init__
    if x.shape[i] != nlabel:
IndexError: tuple index out of range

"""

import numpy as np
from numpy.testing import assert_
import la

nan = np.nan

def getnplabel3d(idx, slices, reduced=False): #True):
    '''constructs list of labels for comparison with generic larry

    does not preserve ordering of labels in each axis
    '''
    ndim = idx[0].ndim
    labs = [np.unique(idx[ii][slices]).tolist() for ii in range(ndim)]
    if reduced:
        labs = [ii for ii in labs if len(ii) > 1]
    return labs

def getnplabel3dm(idx, slices, reduced=True):
    '''constructs list of labels for comparison with generic larry

    preserves ordering of labels in each axis, usable for morph comparison
    '''
    #not sure anymore if it doesn't work easier
    ndim = idx[0].ndim
    labs = []
    for ii in range(ndim):
        sliceidx = list(np.zeros(ndim,int))
        sliceidx[ii] = slices[ii]
        labs.append((idx[ii][sliceidx]).tolist())
    if reduced:
        labs = [ii for ii in labs if len(ii) > 1]
    return labs

# larry definition on module level or in class setup, for now only 1 larry used

x1 = np.array([[ 2.0, 2.0, 3.0, 1.0],
               [ 3.0, 2.0, 2.0, 1.0],
               [ 1.0, 1.0, 1.0, 1.0]])

# slice tests for 3d

la1_2d0 = la.larry(x1)
la1_2d1 = la.larry(x1)
x1_3d = np.rollaxis(np.dstack([x1,2*x1]),2)
la1_3d = la.larry(x1_3d)
n0,n1,n2 = la1_3d.shape
idx = np.mgrid[0:n0,0:n1,0:n2]

def test_slicing():
    "larry slicing nose tests"
    slices = (slice(None), slice(None), slice(None,2))
    larli = [la1_3d, la1_2d0]
    sliceli = [(slice(None), slice(None), slice(None,2)),
               (slice(None), slice(None,1), slice(None,2)),
               (slice(None,1), slice(None), slice(None,2)),
               (slice(None), slice(None), slice(2)),
               (slice(None), slice(None), 0),
               (slice(None), 0, 0),
               (slice(None), slice(2)),
               #([0,1], [0,1], [0,1]),  # fails in larry
               #(np.array([0,1])[:,None], [0,1], [0,1]), #broadcasting
               (0,0,0)
              ]

    for lar in larli:
        for slices in sliceli:
            ndim = lar.ndim
            # create arrays corresponding to labels for slicing
            idxslice = [slice(0,nn) for nn in lar.shape]
            idx = np.mgrid[idxslice]
            slices = slices[:ndim]  # to run same loop on 2d and 3d
            lasliced = lar[slices]
            if not np.isscalar(lasliced):
                lasliced_x = lasliced.x
                lasliced_label = lasliced.label
            else:
                lasliced_x = lasliced
                lasliced_label = []
            reduced = np.ndim(lar.x) > np.ndim(lasliced_x)

            newlabels = getnplabel3d(idx, slices, reduced=reduced)
            yield assert_, np.all(lasliced_x == lar.x[slices]), \
                    'slicing\n%s\n%s' % (repr(lasliced_x), repr(lar.x[slices]))
            yield assert_, lasliced_label == newlabels,\
                    'slicing\n%s\n%s' % (repr(lasliced_label),
                    repr(newlabels))

def test_morph():
    "larry.morph nose tests"
    larli = [la1_3d, la1_2d0]
    slicesmorph = [([0,2,1,3], -1),
                   ([0,2,1], 2),
                   ([1,0], 0),
                   ([0,2,1], 1)
                  ]

    # Some of the things in the following loop are not necessary because
    # morph doesn't reduce dimension
    for lar in larli:
        for newlab, axis in slicesmorph:
            ndim = lar.ndim
            if axis > (ndim - 1):
                continue  # skip infeasible cases

            slices = [slice(None)]*ndim
            slices[axis] = newlab

            # Create arrays corresponding to labels for slicing
            idxslice = [slice(0,nn) for nn in lar.shape]
            idx = np.mgrid[idxslice]
            slices = slices[:ndim]  # to run same loop on 2d and 3d
            lasliced = lar.morph(newlab, axis)
            if not np.isscalar(lasliced):
                lasliced_x = lasliced.x
                lasliced_label = lasliced.label
            else:
                lasliced_x = lasliced
                lasliced_label = []
            reduced = np.ndim(lar.x) > np.ndim(lasliced_x)

            newlabels = getnplabel3dm(idx, slices, reduced=reduced)

            yield assert_, np.all(lasliced_x == lar.x[slices]), \
                    'slicing\n%s\n%s' % (repr(lasliced_x), repr(lar.x[slices]))
            yield assert_, lasliced_label == newlabels,\
                    'slicing\n%s\n%s' % (repr(lasliced_label),
                    repr(newlabels))


#!/usr/bin/python3
# -*- coding: utf-8 -*-

DEFAULT_OUT = "out_alg_sort_selection.txt"
DEFAULT_SEED = None

DEFAULT_N_START = 1
DEFAULT_N_STOP = 10
DEFAULT_N_STEP = 1
DEFAULT_TRIALS = 3

from subprocess import Popen, PIPE
from time import sleep, time
from multiprocessing import Process
import shlex
import json

import sys
import os
import argparse
import logging
import subprocess

import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as opt
import matplotlib.colors as colors
import matplotlib.cm as cmx

import timeit

class Node:

    def __init__(self, data):
        self.data = data
        self.proximo = None

class ListaEncadeada:

    def __init__(self):
        self.cabeca = None

    def reverter(self):
        if self.cabeca is None or self.cabeca.proximo is None:
            return

        anterior = None
        atual = self.cabeca

        while atual:
            proximo_elemento = atual.proximo
            atual.proximo = anterior
            anterior = atual
            atual = proximo_elemento

        self.cabeca = anterior

    def push(self, data):
        novo_nodo = Node(data)
        novo_nodo.proximo = self.cabeca
        self.cabeca = novo_nodo

    def print_list(self):
        atual = self.cabeca
        l1 = []
        while atual:
            l1.append(atual.data)
            atual = atual.proximo
        return l1

def sort_reverse(lista):
	"""
	Implementação do Sort Reverse
	:param lista: qualquer ordem
	:return: lista reversa
	"""
	'''n = len(lista)-1
	m = len(lista)
	for i in range(m):
		aux = lista[n]
		lista[n] = i
		lista[i] = aux
		n -= 1'''
	cabeca = ListaEncadeada()
	for i in lista:
		cabeca.push(i)
	cabeca.reverter()
	lista = cabeca.print_list()
	return lista

def main():
	# Definição de argumentos
	parser = argparse.ArgumentParser(description='Naive TPS')
	help_msg = "arquivo de saída.  Padrão:{}".format(DEFAULT_OUT)
	parser.add_argument("--out", "-o", help=help_msg, default=DEFAULT_OUT, type=str)

	help_msg = "semente aleatória. Padrão:{}".format(DEFAULT_SEED)
	parser.add_argument("--seed", "-s", help=help_msg, default=DEFAULT_SEED, type=int)

	help_msg = "n máximo.          Padrão:{}".format(DEFAULT_N_STOP)
	parser.add_argument("--nstop", "-n", help=help_msg, default=DEFAULT_N_STOP, type=int)

	help_msg = "n mínimo.          Padrão:{}".format(DEFAULT_N_START)
	parser.add_argument("--nstart", "-a", help=help_msg, default=DEFAULT_N_START, type=int)

	help_msg = "n passo.           Padrão:{}".format(DEFAULT_N_STEP)
	parser.add_argument("--nstep", "-e", help=help_msg, default=DEFAULT_N_STEP, type=int)

	help_msg = "tentativas.        Padrão:{}".format(DEFAULT_N_STEP)
	parser.add_argument("--trials", "-t", help=help_msg, default=DEFAULT_TRIALS, type=int)

	# Lê argumentos from da linha de comando
	args = parser.parse_args()

	trials = args.trials
	f = open(args.out, "w")
	f.write("#Reverse sort\n")
	f.write("#n time_s_avg time_s_std (for {} trials)\n".format(trials))
	m = 100
	np.random.seed(args.seed)
	for n in range(args.nstart, args.nstop+1, args.nstep): #range(1, 100):
		resultados = [0 for i in range(trials)]
		tempos = [0 for i in range(trials)]
		for trial in range(trials):
			print("\n-------")
			print("n: {} trial: {}".format(n, trial+1))
			entrada = np.random.randint(0, n, n)
			print("Entrada: {}".format(entrada))
			tempo_inicio = timeit.default_timer()
			resultados[trial] = sort_reverse(entrada)
			tempo_fim = timeit.default_timer()
			tempos[trial] = tempo_fim - tempo_inicio
			print("Saída: {}".format(resultados[trial]))
			print('Tempo: {} s'.format(tempos[trial]))
			print("")

		tempos_avg = np.average(tempos)  # calcula média
		tempos_std = np.std(a=tempos, ddof=False)  # ddof=calcula desvio padrao de uma amostra?

		f.write("{} {} {}\n".format(n, tempos_avg, tempos_std))
	f.close()

if __name__ == '__main__':
	sys.exit(main())


import numpy as np

def random_split_data(data, label, proportion):
    """
    Split two numpy arrays into two parts of `proportion` and `1 - proportion`

    Args:
        - data: numpy array, to be split along the first axis
        - proportion: a float less than 1
    """
    assert data.shape[0] == label.shape[0]
    size = data.shape[0]
    s = np.random.permutation(size)
    split_idx = int(proportion * size)
    return data[s[:split_idx]], label[s[:split_idx]], data[s[split_idx:]], label[s[split_idx:]]


def ex_compiler_pass():

    # magictoken.ex_compiler_pass.begin
    from numba import njit
    from numba import ir
    from numba.compiler import CompilerBase, DefaultPassBuilder
    from numba.compiler_machinery import FunctionPass, register_pass
    from numba.untyped_passes import IRProcessing
    from numbers import Number

    # Register this pass with the compiler framework, declare that it will not
    # mutate the control flow graph and that it is not an analysis_only pass (it
    # potentially mutates the IR).
    @register_pass(mutates_CFG=False, analysis_only=False)
    class ConstsAddOne(FunctionPass):
        _name = "consts_add_one" # the common name for the pass

        def __init__(self):
            FunctionPass.__init__(self)

        # implement method to do the work, "state" is the internal compiler
        # state from the CompilerBase instance.
        def run_pass(self, state):
            func_ir = state.func_ir # get the FunctionIR object
            mutated = False # used to record whether this pass mutates the IR
            # walk the blocks
            for blk in func_ir.blocks.values():
                # find the assignment nodes in the block and walk them
                for assgn in blk.find_insts(ir.Assign):
                    # if an assignment value is a ir.Consts
                    if isinstance(assgn.value, ir.Const):
                        const_val = assgn.value
                        # if the value of the ir.Const is a Number
                        if isinstance(const_val.value, Number):
                            # then add one!
                            const_val.value += 1
                            mutated |= True
            return mutated # return True if the IR was mutated, False if not.
    # magictoken.ex_compiler_pass.end

    # magictoken.ex_compiler_defn.begin
    class MyCompiler(CompilerBase): # custom compiler extends from CompilerBase

        def define_pipelines(self):
            # define a new set of pipelines (just one in this case) and for ease
            # base it on an existing pipeline from the DefaultPassBuilder,
            # namely the "nopython" pipeline
            pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
            # Add the new pass to run after IRProcessing
            pm.add_pass_after(ConstsAddOne, IRProcessing)
            # finalize
            pm.finalize()
            # return as an iterable, any number of pipelines may be defined!
            return [pm]
    # magictoken.ex_compiler_defn.end

    # magictoken.ex_compiler_call.begin
    @njit(pipeline_class=MyCompiler) # JIT compile using the custom compiler
    def foo(x):
        a = 10
        b = 20.2
        c = x + a + b
        return c

    print(foo(100)) # 100 + 10 + 20.2 (+ 1 + 1), extra + 1 + 1 from the rewrite!
    # magictoken.ex_compiler_call.end

    # magictoken.ex_compiler_timings.begin
    compile_result = foo.overloads[foo.signatures[0]]
    nopython_times = compile_result.metadata['pipeline_times']['nopython']
    for k in nopython_times.keys():
        if ConstsAddOne._name in k:
            print(nopython_times[k])
    # magictoken.ex_compiler_timings.end

    assert foo(100) == 132.2

ex_compiler_pass()


# ---
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# jupyter:
#   jupytext:
#     formats: ipynb,md:myst,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.10.0
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# # Autodidax: JAX core from scratch
#
# Ever want to learn how JAX works, but the implementation seemed too
# impenetrable? Well, you're in luck! By reading this tutorial, you'll learn
# every big idea in JAX's core system. You'll even get clued into our weird
# jargon!

# ## Part 1: Transformations as interpreters: standard evaluation, `jvp`, and `vmap`
#
# We want to transform functions that look like this:
#
# ```python
# def f(x):
#   y = sin(x) * 2
#   z = - y + x
#   return z
# ```
#
# Think of functions like `sin` and the arithmetic operations underlying the
# infix operators (`mul`, `add`, and `neg`) as primitive operations, meaning
# atomic units of processing rather than compositions.
#
# "Transform" means "interpret differently." Instead of standard interpretation
# where we apply primitive functions to numerical inputs to produce numerical
# outputs, we want to override primitive application and let different values
# flow through our program. For example, we might want to replace the
# application of every primitive with an application of [its JVP
# rule](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html),
# and let primal-tangent pairs flow through our program. Moreover, we want to
# apply a composition of multiple transformations, leading to stacks of
# interpreters.

# ### JAX core machinery
#
# We can implement stacks of interpreters and even have them all discharge on
# the fly as we execute the Python function to be transformed. To start, let's
# define these primitives so that we can intercept their application:

# +
from typing import NamedTuple

class Primitive(NamedTuple):
  name: str

add_p = Primitive('add')
mul_p = Primitive('mul')
neg_p = Primitive("neg")
sin_p = Primitive("sin")
cos_p = Primitive("cos")
reduce_sum_p = Primitive("reduce_sum")
greater_p = Primitive("greater")

def add(x, y): return bind(add_p, x, y)
def mul(x, y): return bind(mul_p, x, y)
def neg(x): return bind(neg_p, x)
def sin(x): return bind(sin_p, x)
def cos(x): return bind(cos_p, x)
def reduce_sum(x, axis=None): return bind(reduce_sum_p, x, axis=axis)
def greater(x, y): return bind(greater_p, x, y)

# -

# We'll set up array data types and infix operator methods in a moment.
#
# A `Primitive` is just an object with a name, to which we attach our
# interpretation rules (one for each transformation). The `bind` function is our
# interception point: it'll figure out which transformation rule to apply, based
# on how the arguments are boxed in tracers and what interpreters are active.
#
# The functions that user code calls, like `add` and `sin`, are just wrappers
# around calls to `bind`. These wrappers let us control how arguments are passed
# to `bind`, and in particular we follow a handy internal convention: when we
# call `bind`, we pass values representing array data as positional arguments,
# and we pass metadata like the `axis` argument to `sum_p` via keyword. This
# calling convention simplifies some core logic (since e.g. instances of the
# `Tracer` class to be defined below can only occurr in positional arguments to
# `bind`). The wrappers can also provide docstrings!
#
# We represent active interpreters as a stack. The stack is just a simple
# `list`, and each element is a container with an integer level (corresponding
# to the element's height in the stack), an interpreter type (which we'll call a
# `trace_type`), and an optional field for any global data the interpreter
# needs. We call each element a `MainTrace`, though maybe "Interpreter" would be
# more descriptive.

# +
from contextlib import contextmanager
from typing import Type, List, Optional, Any

class MainTrace(NamedTuple):
  level: int
  trace_type: Type['Trace']
  global_data: Optional[Any]

trace_stack: List[MainTrace] = []

@contextmanager
def new_main(trace_type: Type['Trace'], global_data=None):
  level = len(trace_stack)
  main = MainTrace(level, trace_type, global_data)
  trace_stack.append(main)

  try:
    yield main
  finally:
    trace_stack.pop()

# -

# When we're about to apply a transformed function, we'll push another
# interpreter onto the stack using `new_main`. Then, as we apply primitives in
# the function, we can think of the `bind` first being interpreted by the trace
# at the top of the stack (i.e. with the highest level). If that first
# interpreter itself binds other primitives in its interpretation rule for the
# primitive, like how the JVP rule of `sin_p` might bind `cos_p` and `mul_p`,
# then those `bind` calls will be handled by the interpreter at the next level
# down.
#
# What goes at the bottom of the interpreter stack? At the bottom, we know all
# the transformation interpreters are finished, and we just want to do standard
# evaluation. So at the bottom we'll put an evaluation interpreter.
#
# Let's sketch out the interface for interpreters, which is based on the `Trace`
# and `Tracer` base classes. A `Tracer` represents a boxed-up value, perhaps
# carrying some extra context data used by the interpreter. A `Trace` handles
# boxing up vales into `Tracers` and also handles primitive application.

class Trace:
  main: MainTrace

  def __init__(self, main: MainTrace) -> None:
    self.main = main

  def pure(self, val): assert False  # must override
  def lift(self, val): assert False  # must override

  def process_primitive(self, primitive, tracers, params):
    assert False  # must override

# The first two methods are about boxing up values in `Tracer`s, which are the
# objects that flow through the Python programs we transform. The last method is
# the callback we'll use to interpret primitive application.
#
# The `Trace` itself doesn't contain any data, other than a reference to its
# corresponding `MainTrace` instance. In fact, multiple instances of a `Trace`
# might be created and discarded during an application of a transformation,
# whereas only a single `MainTrace` instance is created per application of a
# transformation.
#
# As for `Tracer`s themselves, each one carries an abstract value (and forwards
# infix operators to it), and the rest is up to the transformation. (The
# relationship between `Tracer`s and `AbstractValue`s is that there's one
# `Tracer` per transformation, and at least one `AbstractValue` per base type,
# like arrays.)

# +
import numpy as np
from typing import Tuple

class Tracer:
  _trace: Trace

  __array_priority__ = 1000

  @property
  def aval(self):
    assert False  # must override

  def full_lower(self):
    return self  # default implementation

  def __neg__(self): return self.aval._neg(self)
  def __add__(self, other): return self.aval._add(self, other)
  def __radd__(self, other): return self.aval._radd(self, other)
  def __mul__(self, other): return self.aval._mul(self, other)
  def __rmul__(self, other): return self.aval._rmul(self, other)
  def __gt__(self, other): return self.aval._gt(self, other)
  def __bool__(self): return self.aval._bool(self)
  def __nonzero__(self): return self.aval._nonzero(self)

  def __getattr__(self, name):
    try:
      return getattr(self.aval, name)
    except AttributeError:
      raise AttributeError(f"{self.__class__.__name__} has no attribute {name}")

class ShapedArray:
  array_abstraction_level = 1
  shape: Tuple[int]
  dtype: np.dtype

  def __init__(self, shape, dtype):
    self.shape = shape
    self.dtype = dtype

  @property
  def ndim(self):
    return len(self.shape)

  _neg = staticmethod(neg)
  _add = staticmethod(add)
  _radd = staticmethod(add)
  _mul = staticmethod(mul)
  _rmul = staticmethod(mul)
  _gt = staticmethod(greater)

  @staticmethod
  def _bool(tracer):
    raise Exception("ShapedArray can't be unambiguously converted to bool")

  @staticmethod
  def _nonzero(tracer):
    raise Exception("ShapedArray can't be unambiguously converted to bool")

  def str_short(self):
    return f'{self.dtype.name}[{",".join(str(d) for d in self.shape)}]'

class ConcreteArray(ShapedArray):
  array_abstraction_level = 2
  val: np.ndarray

  def __init__(self, val):
    self.val = val
    self.shape = val.shape
    self.dtype = val.dtype

  @staticmethod
  def _bool(tracer):
    return bool(tracer.aval.val)

  @staticmethod
  def _nonzero(tracer):
    return bool(tracer.aval.val)

def get_aval(x):
  if isinstance(x, Tracer):
    return x.aval
  else:
    return ConcreteArray(np.asarray(x))

# -

# Notice that we actually have two `AbstractValue`s for arrays, representing
# different levels of abstraction. A `ShapedArray` represents the set of all
# possible arrays with a given shape and dtype. A `ConcreteArray` represents a
# singleton set consisting of a single array value.
#
# Now that we've set up the trace stack, the Trace/Tracer API for interpreters,
# and abstract values, we can come back to implement `bind`:

def bind(prim, *args, **params):
  top_trace = find_top_trace(args)
  tracers = [full_raise(top_trace, arg) for arg in args]
  out = top_trace.process_primitive(prim, tracers, params)
  return full_lower(out)

# The main action is that we call `find_top_trace` to figure out which
# interpreter should handle this primitive application as a function of the
# arguments and the active traces on the trace stack. We then call that top
# trace's `process_primitive` so that the trace can apply its interpretation
# rule. The calls to `full_raise` just ensure that the inputs are boxed in the
# top trace's `Tracer` instances, and the call to `full_lower` is an optional
# optimization so that we unbox values out of `Tracer`s as much as possible.

# +
from operator import attrgetter

def find_top_trace(xs) -> Trace:
  top_main = max((x._trace.main for x in xs if isinstance(x, Tracer)),
                 default=trace_stack[0], key=attrgetter('level'))
  return top_main.trace_type(top_main)

# -

# In words, `find_top_trace` returns the highest-level interpreter associated
# with the `Tracer`s on its inputs, and otherwise returns the interpreter at the
# bottom of the stack (which is always an evaluation trace, at least for now).
# This corresponds to JAX transformations mostly working by data dependence
# _except_ for the special bottom-of-the-stack interpreter, which interprets
# everything.

# +
def full_lower(val):
  if isinstance(val, Tracer):
    return val.full_lower()
  else:
    return val

def full_raise(trace, val) -> Tracer:
  if not isinstance(val, Tracer):
    return trace.pure(val)
  level = trace.main.level
  if val._trace.main is trace.main:
    return val
  elif val._trace.main.level < level:
    return trace.lift(val)
  elif val._trace.main.level > level:
    raise Exception(f"Can't lift level {val._trace.main.level} to {level}.")
  else:  # val._trace.level == level
    raise Exception(f"Different traces at same level: {val._trace}, {trace}.")

# -

# The logic in `full_raise` serves to box values into `Tracer`s for a particular
# `Trace`, calling different methods on the `Trace` based on context:
# `Trace.pure` is called on non-`Tracer` constants, and `Trace.lift` is called
# for values that are already `Tracer`s from a lower-level interpreter. These
# two methods could share the same implementation, but by distinguishing them in
# the core logic we can provide more information to the `Trace` subclass.
#
# That's it for the JAX core! Now we can start adding interpreters.

# ### Evaluation interpreter
#
# We'll start with the simplest interpreter: the evaluation interpreter that
# will sit at the bottom of the interpreter stack.

# +
class EvalTrace(Trace):
  pure = lift = lambda self, x: x  # no boxing in Tracers needed

  def process_primitive(self, primitive, tracers, params):
    return impl_rules[primitive](*tracers, **params)

trace_stack.append(MainTrace(0, EvalTrace, None))  # special bottom of the stack

impl_rules = {}
impl_rules[add_p] = np.add
impl_rules[mul_p] = np.multiply
impl_rules[neg_p] = np.negative
impl_rules[sin_p] = np.sin
impl_rules[cos_p] = np.cos
impl_rules[reduce_sum_p] = np.sum
impl_rules[greater_p] = np.greater

# -

# With this interpreter, we can evaluate user functions:

# +
def f(x):
  y = sin(x) * 2
  z = - y + x
  return z

print(f(3.0))

# -

# Woo! Like going around in a big circle. But the point of this indirection is
# that now we can add some real transformations.

# ### Forward-mode autodiff with `jvp`
#
# First, a couple of helper functions:

# +
def zeros_like(val):
  return np.zeros_like(val)

def unzip2(pairs):
  lst1, lst2 = [], []
  for x1, x2 in pairs:
    lst1.append(x1)
    lst2.append(x2)
  return lst1, lst2

# -

# The `Tracer` for forward-mode autodiff carries a primal-tangent pair. The
# `Trace` applies JVP rules.

# +
class JVPTracer(Tracer):
  def __init__(self, trace, primal, tangent):
    self._trace = trace
    self.primal = primal
    self.tangent = tangent

  @property
  def aval(self):
    return get_aval(self.primal)

class JVPTrace(Trace):
  pure = lift = lambda self, val: JVPTracer(self, val, zeros_like(val))

  def process_primitive(self, primitive, tracers, params):
    primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)
    jvp_rule = jvp_rules[primitive]
    primal_out, tangent_out = jvp_rule(primals_in, tangents_in, **params)
    return JVPTracer(self, primal_out, tangent_out)

jvp_rules = {}

# -

# Notice both `lift` and `sublift` package a value into a `JVPTracer` with the
# minimal amount of context, which is a zero tangent value.

# Let's add some JVP rules for primitives:

# +
def add_jvp(primals, tangents):
  (x, y), (x_dot, y_dot) = primals, tangents
  return x + y, x_dot + y_dot
jvp_rules[add_p] = add_jvp

def mul_jvp(primals, tangents):
  (x, y), (x_dot, y_dot) = primals, tangents
  return x * y, x_dot * y + x * y_dot
jvp_rules[mul_p] = mul_jvp

def sin_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return sin(x), cos(x) * x_dot
jvp_rules[sin_p] = sin_jvp

def cos_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return cos(x), -sin(x) * x_dot
jvp_rules[cos_p] = cos_jvp

def neg_jvp(primals, tangents):
  (x,), (x_dot,) = primals, tangents
  return neg(x), neg(x_dot)
jvp_rules[neg_p] = neg_jvp

def reduce_sum_jvp(primals, tangents, *, axis):
  (x,), (x_dot,) = primals, tangents
  return reduce_sum(x, axis), reduce_sum(x_dot, axis)
jvp_rules[reduce_sum_p] = reduce_sum_jvp

def greater_jvp(primals, tangents):
  (x, y), _ = primals, tangents
  out_primal = greater(x, y)
  return out_primal, zeros_like(out_primal)
jvp_rules[greater_p] = greater_jvp

# -

# Finally, we add a transformation API to kick off the trace:

def jvp(f, primals, tangents):
  with new_main(JVPTrace) as main:
    trace = JVPTrace(main)
    tracers_in = [JVPTracer(trace, x, t) for x, t in zip(primals, tangents)]
    out = f(*tracers_in)
    tracer_out = full_raise(trace, out)
    primal_out, tangent_out = tracer_out.primal, tracer_out.tangent
  return primal_out, tangent_out

# And with that, we can differentiate!

x = 3.0
y, sin_deriv_at_3 = jvp(sin, (x,), (1.0,))
print(sin_deriv_at_3)
print(cos(3.0))

# +
def f(x):
  y = sin(x) * 2
  z = - y + x
  return z

x, xdot = 3., 1.
y, ydot = jvp(f, (x,), (xdot,))
print(y)
print(ydot)

# +
def deriv(f):
  return lambda x: jvp(f, (x,), (1.,))[1]

print(deriv(sin)(3.))
print(deriv(deriv(sin))(3.))
print(deriv(deriv(deriv(sin)))(3.))
print(deriv(deriv(deriv(deriv(sin))))(3.))

# +
def f(x):
  if x > 0.:  # Python control flow
    return 2. * x
  else:
    return x

print(deriv(f)(3.))
print(deriv(f)(-3.))

# -

# ### Vectorized batching with `vmap`
#
# First, a couple helper functions, one for producing mapped abstract values
# from unmapped ones (by removing an axis), and one for moving batch dimensions
# around:

# +
def mapped_aval(batch_dim, aval):
  shape = list(aval.shape)
  del shape[batch_dim]
  return ShapedArray(tuple(shape), aval.dtype)

def move_batch_axis(axis_size, src, dst, x):
  if src is not_mapped:
    target_shape = list(np.shape(x))
    target_shape.insert(dst, axis_size)
    return np.broadcast_to(np.expand_dims(x, dst), target_shape)
  else:
    return np.moveaxis(x, src, dst)

# -

# The `Tracer` for vectorized batching carries a batched value and an optional
# integer indicating which axis (if any) is the batch axis.

# +
from typing import Union

class NotMapped: pass
not_mapped = NotMapped()

class BatchTracer(Tracer):
  def __init__(self, trace, val, batch_dim: Union[NotMapped, int]):
    self._trace = trace
    self.val = val
    self.batch_dim = batch_dim

  @property
  def aval(self):
    if self.batch_dim is not_mapped:
      return get_aval(self.val)
    else:
      return mapped_aval(self.batch_dim, get_aval(self.val))

  def full_lower(self):
    if self.batch_dim is not_mapped:
      return full_lower(self.val)
    else:
      return self

class BatchTrace(Trace):
  pure = lift = lambda self, val: BatchTracer(self, val, not_mapped)

  def process_primitive(self, primitive, tracers, params):
    vals_in, bdims_in = unzip2((t.val, t.batch_dim) for t in tracers)
    vmap_rule = vmap_rules[primitive]
    val_out, bdim_out = vmap_rule(self.axis_size, vals_in, bdims_in, **params)
    return BatchTracer(self, val_out, bdim_out)

  @property
  def axis_size(self):
    return self.main.global_data

vmap_rules = {}
# -

# Here we've implemented the optional `Tracer.full_lower` method, which lets us
# peel off a batching tracer if it's not needed because it doesn't represent a
# batched value.
#
# For `BatchTrace`, analogous to `JVPTrace`, the methods `pure` and `lift` just
# box a value in a `BatchTracer` with the minimal amount of context, which in
# this case is a `batch_dim` taking the sentinel value `not_mapped`. Notice we
# use the `MainTrace`'s interpreter-global data field to store the batch axis
# size.
#
# Next we can define batching interpreter rules for each primitive:

# +
from functools import partial

def broadcasting_binop_batching_rule(op, axis_size, vals_in, dims_in):
  (x, y), (x_bdim, y_bdim) = vals_in, dims_in
  if x_bdim != y_bdim:
    y = move_batch_axis(axis_size, y_bdim, x_bdim, y)
  return op(x, y), x_bdim
vmap_rules[add_p] = partial(broadcasting_binop_batching_rule, add)
vmap_rules[mul_p] = partial(broadcasting_binop_batching_rule, mul)

def vectorized_unop_batching_rule(op, axis_size, vals_in, dims_in):
  (x,), (x_bdim,) = vals_in, dims_in
  return op(x), x_bdim
vmap_rules[sin_p] = partial(vectorized_unop_batching_rule, sin)
vmap_rules[cos_p] = partial(vectorized_unop_batching_rule, cos)
vmap_rules[neg_p] = partial(vectorized_unop_batching_rule, neg)

def reduce_sum_batching_rule(axis_size, vals_in, dims_in, *, axis):
  (x,), (x_bdim,) = vals_in, dims_in
  new_axis = axis + (x_bdim <= axis)
  out_bdim = x_bdim - (new_axis < x_bdim)
  return reduce_sum(x, new_axis), out_bdim
vmap_rules[reduce_sum_p] = reduce_sum_batching_rule

# -

# Finally, we add a transformation API to kick off the trace:

def vmap(f, in_axes, out_axis):
  def batched_f(*args):
    axis_size, = {x.shape[ax] for x, ax in zip(args, in_axes)
                  if ax is not None}
    with new_main(BatchTrace, axis_size) as main:
      trace = BatchTrace(main)
      tracers_in = [BatchTracer(trace, x, ax) if ax is not None else x
                    for x, ax in zip(args, in_axes)]
      out = f(*tracers_in)
      tracer_out = full_raise(trace, out)
      val_out, batch_dim_out = tracer_out.val, tracer_out.batch_dim
    return move_batch_axis(axis_size, batch_dim_out, out_axis, val_out)
  return batched_f

# +
def add_one_to_a_scalar(scalar):
  assert np.ndim(scalar) == 0
  return 1 + scalar

vector_in = np.arange(3.)
vector_out = vmap(add_one_to_a_scalar, (0,), 0)(vector_in)

print(vector_in)
print(vector_out)

# +
def jacfwd(f, x):
  pushfwd = lambda v: jvp(f, (x,), (v,))[1]
  vecs_in = np.eye(np.size(x)).reshape(np.shape(x) * 2)
  return vmap(pushfwd, (0,), 0)(vecs_in)

def f(x):
  return sin(x)

jacfwd(f, np.arange(3.))
# -

# That's it for `jvp` and `vmap`! Before moving on, let's highlight a few
# simplifications in what we've seen so far compared to the full JAX
# implementation:
# 1. **Fewer, simpler primitives.** More primitives means more interpretation
# rules, and for more complex primitives (like for convolution or advanced
# indexing) each rule is harder to write. But the overarching design is no
# different.
# 1. **Transformations expect arrays in, single array out.**
# 2. **No symbolic zeros in autodiff.**
# 3. **No special call primitives yet.** The core machinery needs to be
#     generalized to handle the most flexible kind of higher-order primitive,
#     used by `jax.custom_jvp` and `jax.custom_vjp`.

# ## Part 2: Jaxprs, for `jit` and `vjp`
#
# The next transformations are the horizon are `jit` for just-in-time
# compilation and `vjp` for reverse-mode autodiff.  (`grad` is just a small
# wrapper around `vjp`.) For `jvp` and `vmap` we only needed each `Tracer` to
# carry a little    bit of extra context, but for both `jit` and `vjp` we need
# much richer context: we need to represent _programs_. That is, we need jaxprs!
#
# Jaxprs are JAX's internal intermediate representation of programs. Jaxprs are
# an explicitly typed, functional, first-order language. We need a program
# representation for `jit` because the purpose of `jit` is to stage computation
# out of Python. For any computation we want to stage out, we need to be able to
# represent it as data, and build it up as we trace a Python function.
# Similarly, `vjp` needs a way to represent the computation for the backward
# pass of reverse-mode autodiff. We use the same jaxpr program representation
# for both needs.
#
# (Building a program representation is the most
# [free](https://en.wikipedia.org/wiki/Free_object) kind of
# trace- transformation, and so except for issues around handling native Python
# control flow, any transformation could be implemented by first tracing to a
# jaxpr and then interpreting the jaxpr.)
#
# The jaxpr term syntax is roughly:
#
# ```
# jaxpr ::=
#   { lambda <binder> , ... .
#     let <eqn>
#         ...
#     in <atom> }
#
# binder ::= <var>:<array_type>
# var ::= a | b | c | ...
# atom ::= <var> | <literal>
# literal ::= <int32> | <float32>
#
# eqn ::= <binder> = <primitive> [ <params> ] <atom> , ...
# ```
#
# The syntax of types is:
#
# ```
# jaxpr_type ::= [<array_type>, ...] -> [<array_type>, ...]
# array_type ::= <dtype>[<shape>]
# dtype ::= f32 | f64 | i32 | i64
# shape ::= <int> , ...
# ```
#
# How do we represent these as Python data structures? We reuse ShapedArrays to
# represent types, and we can represent the term syntax with a few Python
# structs:

# +
from typing import Dict, Set

class Var:
  aval: ShapedArray
  def __init__(self, aval): self.aval = aval

class Lit:
  val: Any
  aval: ShapedArray

  def __init__(self, val):
    self.val = val
    self.aval = raise_to_shaped(get_aval(self.val))

Atom = Union[Var, Lit]

class JaxprEqn(NamedTuple):
  primitive: Primitive
  inputs: List[Atom]
  params: Dict[str, Any]
  out_binder: Var

class Jaxpr(NamedTuple):
  in_binders: List[Var]
  eqns: List[JaxprEqn]
  out: Atom

def raise_to_shaped(aval):
  return ShapedArray(aval.shape, aval.dtype)

# +
class JaxprType:
  in_types: List[ShapedArray]
  out_type: ShapedArray

  def __init__(self, in_types, out_type):
    self.in_types = in_types
    self.out_type = out_type

  def __repr__(self):
    in_types = ', '.join(aval.str_short() for aval in self.in_types)
    out_type = self.out_type.str_short()
    return f'({in_types}) -> {out_type}'

def typecheck_jaxpr(jaxpr: Jaxpr) -> JaxprType:
  env: Set[Var] = set()

  for v in jaxpr.in_binders:
    env.add(v)

  for eqn in jaxpr.eqns:
    in_types = [typecheck_atom(env, x) for x in eqn.inputs]
    out_type = abstract_eval_rules[eqn.primitive](*in_types, **eqn.params)
    if not types_equal(out_type, eqn.out_binder.aval): raise TypeError
    env.add(eqn.out_binder)

  out_type = typecheck_atom(env, jaxpr.out)
  return JaxprType([v.aval for v in jaxpr.in_binders], out_type)

def typecheck_atom(env: Set[Var], x: Atom) -> ShapedArray:
  if isinstance(x, Var):
    if x not in env: raise TypeError("unbound variable")
    return x.aval
  elif isinstance(x, Lit):
    return raise_to_shaped(get_aval(x.val))
  else:
    assert False

def types_equal(a: ShapedArray, b: ShapedArray) -> bool:
  return a.shape == b.shape and a.dtype == b.dtype

# -

# Now that we have jaxprs as a data structure, we need ways to produce these
# from tracing Python code. In general there are two variants of how we trace to
# a jaxpr; `jit` uses one and `vjp` uses the other. We'll start with the one
# used   by `jit`, which is also used by control flow primitives like
# `lax.cond`, `lax.while_loop`, and `lax.scan`.

# +
# NB: the analogous class in JAX is called 'DynamicJaxprTracer'
class JaxprTracer(Tracer):
  __slots__ = ['aval']
  aval: ShapedArray

  def __init__(self, trace, aval):
    self._trace = trace
    self.aval = aval

# NB: the analogous class in JAX is called 'DynamicJaxprTrace'
class JaxprTrace(Trace):
  def new_arg(self, aval: ShapedArray) -> JaxprTracer:
    aval = raise_to_shaped(aval)
    tracer = JaxprTracer(self, aval)
    self.builder.tracer_to_var[id(tracer)] = Var(aval)
    return tracer

  def get_or_make_const_tracer(self, val: Any) -> JaxprTracer:
    tracer = self.builder.const_tracers.get(id(val))
    if tracer is None:
      tracer = JaxprTracer(self, raise_to_shaped(get_aval(val)))
      self.builder.add_const(tracer, val)
    return tracer
  pure = lift = get_or_make_const_tracer

  def process_primitive(self, primitive, tracers, params):
    avals_in = [t.aval for t in tracers]
    aval_out = abstract_eval_rules[primitive](*avals_in, **params)
    out_tracer = JaxprTracer(self, aval_out)
    inputs = [self.builder.getvar(t) for t in tracers]
    outvar = self.builder.add_var(out_tracer)
    self.builder.add_eqn(JaxprEqn(primitive, inputs, params, outvar))
    return out_tracer

  @property
  def builder(self):
    return self.main.global_data

# NB: in JAX, instead of a dict we attach impl rules to the Primitive instance
abstract_eval_rules = {}

# -

# Notice that we keep as interpreter-global data a builder object, which keeps
# track of variables, constants, and eqns  as we build up the jaxpr.

class JaxprBuilder:
  eqns: List[JaxprEqn]
  tracer_to_var: Dict[int, Var]
  const_tracers: Dict[int, JaxprTracer]
  constvals: Dict[Var, Any]

  def __init__(self):
    self.eqns = []
    self.tracer_to_var = {}
    self.const_tracers = {}
    self.constvals = {}

  def add_eqn(self, eqn: JaxprEqn) -> None:
    self.eqns.append(eqn)

  def add_var(self, tracer: JaxprTracer) -> Var:
    var = self.tracer_to_var.get(id(tracer))
    assert var is None
    var = self.tracer_to_var[id(tracer)] = Var(tracer.aval)
    return var

  def getvar(self, tracer: JaxprTracer) -> Var:
    var = self.tracer_to_var.get(id(tracer))
    assert var is not None
    return var

  def add_const(self, tracer: JaxprTracer, val: Any) -> Var:
    var = self.add_var(tracer)
    self.const_tracers[id(val)] = tracer
    self.constvals[var] = val
    return var

  def build(self, in_tracers: List[JaxprTracer], out_tracer: JaxprTracer
            ) -> Tuple[Jaxpr, List[Any]]:
    constvars, constvals = unzip2(self.constvals.items())
    t2v = lambda t: self.tracer_to_var[id(t)]
    in_binders = constvars + [t2v(t) for t in in_tracers]
    jaxpr = Jaxpr(in_binders, self.eqns, t2v(out_tracer))
    typecheck_jaxpr(jaxpr)
    return jaxpr, constvals

# The rules we need for `JaxprTrace.process_primitive` are essentially typing
# rules for primitive applications: given   the primitive, its parameters, and
# types for the inputs, the rule must produce a type for the output, which is
# then   packaged with the output `JaxprTracer`. We can use abstract evaluation
# rules for this same purpose, even though they  can be more general (since
# abstract evaluation rules need to work on ConcreteArray inputs as well). We'll
# reuse these abstract evaluation rules for the other jaxpr-producing trace
# machinery, where the potential extra generality is useful.

def broadcast_shapes(*shapes):
  assert len(shapes) > 1
  for sizes in zip(*shapes):
    sizes = [d for d in sizes if d != 1]
    if sizes[:-1] != sizes[1:]:
      raise Exception
  return tuple(next((d for d in sizes if d != 1), 1) for sizes in zip(*shapes))

# +
def broadcasting_binop_abstract_eval_rule(*avals_in):
  out_dtype = np.result_type(*map(np.result_type, avals_in))
  out_shape = broadcast_shapes(*map(np.shape, avals_in))
  return ShapedArray(out_shape, out_dtype)

abstract_eval_rules[add_p] = broadcasting_binop_abstract_eval_rule
abstract_eval_rules[mul_p] = broadcasting_binop_abstract_eval_rule

def vectorized_unop_abstract_eval_rule(aval_in):
  return ShapedArray(np.shape(aval_in), np.result_type(aval_in))

abstract_eval_rules[sin_p] = vectorized_unop_abstract_eval_rule
abstract_eval_rules[cos_p] = vectorized_unop_abstract_eval_rule
abstract_eval_rules[neg_p] = vectorized_unop_abstract_eval_rule

def reduce_sum_abstract_eval_rule(aval_in, *, axis):
  new_shape = [d for i, d in enumerate(aval_in.shape) if i != axis]
  return ShapedArray(tuple(new_shape), aval_in.dtype)
abstract_eval_rules[reduce_sum_p] = reduce_sum_abstract_eval_rule

# -

# To check our implementation, we can add a `make_jaxpr` transformation and
# first pretty-printer:

def make_jaxpr(f, avals_in):
  builder = JaxprBuilder()
  with new_main(JaxprTrace, builder) as main:
    trace = JaxprTrace(main)
    tracers_in = [trace.new_arg(aval) for aval in avals_in]
    out = f(*tracers_in)
    tracer_out = full_raise(trace, out)
    return builder.build(tracers_in, tracer_out)

# +
from collections import defaultdict
import itertools as it
import string

class PPrint:
  lines: List[Tuple[int, str]]

  def __init__(self, lines):
    self.lines = lines

  def indent(self, indent: int) -> 'PPrint':
    return PPrint([(indent + orig_indent, s) for orig_indent, s in self.lines])

  def __add__(self, rhs: 'PPrint') -> 'PPrint':
    return PPrint(self.lines + rhs.lines)

  def __rshift__(self, rhs: 'PPrint') -> 'PPrint':
    if not rhs.lines: return self
    if not self.lines: return rhs
    indent, s = self.lines[-1]
    indented_block = rhs.indent(indent + len(s))
    common_line = s + ' ' * rhs.lines[0][0] + rhs.lines[0][1]
    return PPrint(self.lines[:-1]
                  + [(indent, common_line)]
                  + indented_block.lines[1:])

  def __str__(self) -> str:
    return '\n'.join(' ' * indent + s for indent, s in self.lines)

def pp(s: Any) -> PPrint:
  return PPrint([(0, line) for line in str(s).splitlines()])

def vcat(ps: List[PPrint]) -> PPrint:
  return sum(ps, pp(''))

def pp_jaxpr(jaxpr: Jaxpr):
  namegen = (''.join(s) for r in it.count(1)
             for s in it.permutations(string.ascii_lowercase, r))
  names = defaultdict(lambda: next(namegen))
  in_binders = ', '.join(var_str(names, x) for x in jaxpr.in_binders)
  eqns = vcat([pp_eqn(names, e) for e in jaxpr.eqns])
  out = names[jaxpr.out] if isinstance(jaxpr.out, Var) else str(jaxpr.out.val)
  return (pp(f'{{ lambda {in_binders} .') +
          ((pp('let ') >> eqns) + pp(f'in {out} }}')).indent(2))

def var_str(names: Dict[Var, str], v: Var) -> str:
  return f'{names[v]}:{v.aval.str_short()}'

def pp_eqn(names: Dict[Var, str], eqn: JaxprEqn) -> PPrint:
  lhs = pp(var_str(names, eqn.out_binder))
  rhs = (pp(eqn.primitive.name) >> pp_params(eqn.params) >>
         pp(' '.join(names[x] if isinstance(x, Var) else str(x.val)
                     for x in eqn.inputs)))
  return lhs >> pp(' = ') >> rhs

def pp_params(params: Dict[str, Any]) -> PPrint:
  items = sorted(params.items())
  if items:
    return pp(' [ ') >> vcat([pp(f'{k}={v}') for k, v in items]) >> pp(' ] ')
  else:
    return pp(' ')

# -

jaxpr, consts = make_jaxpr(lambda x: 2. * x, [raise_to_shaped(get_aval(3.))])
print(pp_jaxpr(jaxpr))
print(typecheck_jaxpr(jaxpr))


"""
.. _tut-fnirs-processing-bad:

Importance of Controls and Parameter Selection
==============================================

.. warning:: DO NOT USE THIS SCRIPT TO PROCESS YOUR DATA

.. sidebar:: Relevant literature

   Huppert TJ. Commentary on the statistical properties of noise and its
   implication on general linear models in functional near-infrared
   spectroscopy. Neurophotonics. 2016;3(1)

This tutorial demonstrates how **NOT** to process your
functional near-infrared spectroscopy (fNIRS)
data.
Here I demonstrate that with a simple process pipelines even random noise
can look like a canonical haemodynamic response.

This tutorial was written to motivate further interest in to what affect
signal processing has on our data. And to provide a concrete example
about why you should understand your analysis and not vary parameters to meet
preconceived ideas of what your data should look like.
I recommend specifying analysis parameters in advance along with your experimental
protocol. Hopefully this tutorial also highlights the importance of
including a control condition in your experimental design.

The
:ref:`MNE fNIRS waveform tutorial <mne:tut-fnirs-processing>`
is used as a basis for this analysis, so most of the document looks similar.
However, at the midpoint we replace the real data with noise and demonstrate
that without careful attention to the analysis parameter it would still
appear as if a fNIRS response is observed.

.. contents:: Page contents
   :local:
   :depth: 2

"""
# sphinx_gallery_thumbnail_number = 7

# Authors: Robert Luke <mail@robertluke.net>
#
# License: BSD (3-clause)

import os
import numpy as np

import mne
import mne_nirs
np.random.seed(1)

fnirs_data_folder = mne.datasets.fnirs_motor.data_path()
fnirs_cw_amplitude_dir = os.path.join(fnirs_data_folder, 'Participant-1')
raw_intensity = mne.io.read_raw_nirx(fnirs_cw_amplitude_dir, verbose=True)
raw_intensity.load_data()

# %%
# Selecting channels appropriate for detecting neural responses
# -------------------------------------------------------------
#
# First we remove channels that are too close together (short channels) to
# detect a neural response (less than 1 cm distance between optodes).
# These short channels can be seen in the figure above.
# To achieve this we pick all the channels that are not considered to be short.

picks = mne.pick_types(raw_intensity.info, meg=False, fnirs=True)
dists = mne.preprocessing.nirs.source_detector_distances(
    raw_intensity.info, picks=picks)
raw_intensity.pick(picks[dists > 0.01])
raw_intensity.plot(n_channels=len(raw_intensity.ch_names),
                   duration=500, show_scrollbars=False)

# %%
# Converting from raw intensity to optical density
# ------------------------------------------------
#
# The raw intensity values are then converted to optical density.

raw_od = mne.preprocessing.nirs.optical_density(raw_intensity)
raw_od.plot(n_channels=len(raw_od.ch_names),
            duration=500, show_scrollbars=False)

# %%
# Converting from optical density to haemoglobin
# ----------------------------------------------
#
# Next we convert the optical density data to haemoglobin concentration using
# the modified Beer-Lambert law.

raw_haemo = mne.preprocessing.nirs.beer_lambert_law(raw_od)
raw_haemo.plot(n_channels=len(raw_haemo.ch_names),
               duration=500, show_scrollbars=False)

# %%
# !!!!!!Replace real data with white noise!!!!!
# ----------------------------------------------
#
# Here we replace the signals with white noise.
# We make the HbR white noise 3 time smaller than HbO as is commonly observed.

raw_haemo._data = np.random.randn(40, 23239) / 1.0e6 * 1
raw_haemo._data[::2, :]= np.random.randn(20, 23239) / 1.0e6 * 3

# %%
# Removing heart rate from signal (BAD DONT COPY)
# -----------------------------------------------
#
# This analysis is BAD and is provided just as a demonstration.
# Do not do this!!
#

fig = raw_haemo.plot_psd(average=True)
fig.suptitle('Before filtering', weight='bold', size='x-large')
fig.subplots_adjust(top=0.88)
raw_haemo = raw_haemo.filter(0.05, 0.1, h_trans_bandwidth=0.2,
                             l_trans_bandwidth=0.02)
fig = raw_haemo.plot_psd(average=True)
fig.suptitle('After filtering', weight='bold', size='x-large')
fig.subplots_adjust(top=0.88)

# %%
# Apply Cui negative correlation method
# -------------------------------------
#
# Here we apply the Cui signal enhancement technique.

raw_haemo = mne_nirs.signal_enhancement.enhance_negative_correlation(raw_haemo)

# %%
# Extract epochs
# --------------
#
# Now that the signal has been converted to relative haemoglobin concentration,
# and the unwanted heart rate component has been removed, we can extract epochs
# related to each of the experimental conditions.
#
# First we extract the events of interest and visualise them to ensure they are
# correct.

events, _ = mne.events_from_annotations(raw_haemo, event_id={'1.0': 1,
                                                             '2.0': 2,
                                                             '3.0': 3})
event_dict = {'Control': 1, 'Tapping/Left': 2, 'Tapping/Right': 3}
fig = mne.viz.plot_events(events, event_id=event_dict,
                          sfreq=raw_haemo.info['sfreq'])
fig.subplots_adjust(right=0.7)  # make room for the legend

# %%
# Next we define the range of our epochs, the rejection criteria,
# baseline correction, and extract the epochs. We visualise the log of which
# epochs were dropped.

reject_criteria = dict(hbo=10e-6)
tmin, tmax = -1.5, 10.5

epochs = mne.Epochs(raw_haemo, events, event_id=event_dict,
                    tmin=tmin, tmax=tmax,
                    reject=reject_criteria, reject_by_annotation=True,
                    proj=True, baseline=(None, 0), preload=True,
                    detrend=1, verbose=True)

# %%
# Plot standard fNIRS response image
# ----------------------------------
#
# Next we generate the most common visualisation of fNIRS data: plotting
# both the HbO and HbR on the same figure to illustrate the relation between
# the two signals.

evoked_dict = {'Tapping/HbO': epochs['Tapping'].average(picks='hbo'),
               'Tapping/HbR': epochs['Tapping'].average(picks='hbr')}

# Rename channels until the encoding of frequency in ch_name is fixed
for condition in evoked_dict:
    evoked_dict[condition].rename_channels(lambda x: x[:-4])

color_dict = dict(HbO='#AA3377', HbR='b')

mne.viz.plot_compare_evokeds(evoked_dict, combine="mean", ci=0.95,
                             colors=color_dict)

# %%
# Summary
# --------------
#
# Above is a waveform that has many of the characteristics that are expected
# from a haeomodynamic response.
# We see a baseline that is around zero, the oxyhaemoglobin is positive
# and the deoxyhaemoglobin is negative,
# the peak occurs around 5 seconds then returns to
# baseline.
# However, it was generated from random noise!
#
# .. sidebar:: PLEASE READ THIS PAPER!!
#
#    Huppert TJ. Commentary on the statistical properties of noise and its
#    implication on general linear models in functional near-infrared
#    spectroscopy. Neurophotonics. 2016;3(1)
#
# As this is an introduction tutorial I do not go in to details about the
# signal process theory behind why the settings above are wrong.
# But if you are interested you should start with a DSP textbook
# on signals and systems.
#
# There are many signal processing parameter choices that interact to produce
# your end result. It is important to have a control condition to ensure
# that you aren't forcing your data to have these fake responses.
# Its also important to understand the effect that each parameter choice
# has on your output. Using a receiver operating characteristic is a good
# approach to understand your analysis.
# This also highlights some of the concerns with the time domain analysis
# approach, and the benefits of the GLM style of analysis
# :ref:`(see MNE-NIRS fNIRS GLM tutorial) <tut-fnirs-hrf>` (and Huppert 2016).
#
# .. warning:: DO NOT USE THIS SCRIPT TO PROCESS YOUR DATA


import unittest
import pandas as pd
import numpy as np
import numpy.testing as np_test
import pandas.util.testing as pd_test
from src.preprocess import times
from datetime import datetime

class UtilTest(unittest.TestCase):

    # round hours to a given interval
    def test_round_hour(self):
        dt = pd.to_datetime('2017-01-01 23:30:00', utc=True)
        # round to 3 hours
        rounded = times.round_hour(dt, 3)
        self.assertEqual('21:00:00', rounded.strftime('%H:%M:%S'))
        # round to 6 hours
        rounded = times.round_hour(dt, 6)
        self.assertEqual('18:00:00', rounded.strftime('%H:%M:%S'))
        # round to 12 hours
        rounded = times.round_hour(dt, 12)
        self.assertEqual('12:00:00', rounded.strftime('%H:%M:%S'))
        # round to 24 hours
        rounded = times.round_hour(dt, 24)
        self.assertEqual('00:00:00', rounded.strftime('%H:%M:%S'))

    # test time sampling for different time modes
    @staticmethod
    def test_aggregate_time():
        # sample all hours
        df = pd.DataFrame(data={
            'time': ['2017-01-01 18:20:00', '2017-01-01 18:40:00', '2017-02-07 23:00:01'],
            'value': [1, 2, 3]
        })
        sample = times.group(ts=df, mode='h')
        expected = pd.DataFrame(data={'time': ['2017-01-01 18:00:00', '2017-02-07 23:00:00'],
                                      'value': [1.5, 3]})
        pd_test.assert_frame_equal(sample, expected)
        # sample all 3 hours
        # 23:00 is expected to become 21:00
        sample = times.group(ts=df, mode='3h')
        expected = pd.DataFrame(data={'time': ['2017-01-01 18:00:00', '2017-02-07 21:00:00'],
                                      'value': [1.5, 3]})
        pd_test.assert_frame_equal(sample, expected)
        # sample a specific month = 2
        df = pd.DataFrame(data={
            'time': ['2016-02-01', '2017-01-01', '2017-02-07', '2017-02-07'],
            'value': [2, 1, 2, 3]})
        sample = times.group(ts=df, mode='m', value=2)
        expected = pd.DataFrame(data={'time': ['2016-02', '2017-02'], 'value': [2, 2.5]})
        pd_test.assert_frame_equal(sample, expected)
        # sample a specific day of week, 2017-01-01 is sunday (0)
        df = pd.DataFrame(data={
            'time': ['2017-01-01 12:00', '2017-01-01 14:00', '2017-01-06 10:00'], 'value': [1, 2, 3]})
        sample = times.group(ts=df, mode='dw', value=0)
        expected = pd.DataFrame(data={'time': ['2017-00-0'], 'value': [1.5]})
        pd_test.assert_frame_equal(sample, expected)

    @staticmethod
    def test_filter_by_time():
        df = pd.DataFrame(data={'time': ['2019-01-01', '2020-01-01']})
        df['time'] = pd.to_datetime(df['time'])
        selected = times.select(df, 'time', from_time='2019-01-01', to_time='2019-01-30')
        expected = pd.DataFrame(data={'time': [pd.datetime(2019, 1, 1)]})
        pd_test.assert_frame_equal(selected, expected)
        excluded = times.exclude(df, 'time', from_time='2019-01-01', to_time='2019-01-30')
        expected = pd.DataFrame(data={'time': [pd.datetime(2020, 1, 1)]})
        pd_test.assert_frame_equal(excluded, expected)

    def test_group_at(self):
        # test grouping values backward or forward from a given value and given hours unit
        yr = '2018-01-01 '
        time = pd.to_datetime([yr + ' 12', yr + ' 15', yr + '16', yr + '17'], utc=True).tolist()
        value = [2, 2, 3, 4]
        # expected groups: 12:00, 15:00
        forward_average = times.group_at(time, value, index=2, direction=1, group_hours=3)
        self.assertEqual(first=3.5, second=forward_average)
        backward_average = times.group_at(time, value, index=2, direction=-1, group_hours=3)
        self.assertEqual(first=2.5, second=backward_average)
        single_average = times.group_at(time, value, index=0, direction=-1, group_hours=3)
        self.assertEqual(first=2, second=single_average)

    @staticmethod
    def test_group_from():
        # test grouping values backward or forward from a given value and given hours unit
        yr = '2018-01-01 '
        time = pd.to_datetime([yr + ' 12', yr + ' 15', yr + '16', yr + '17', yr + '18'], utc=True).tolist()
        value = [2, 2, 3, 4, 5]
        # expected to group (2, 2, 3) into 9:00, 12:00 and 15:00
        # where 9:00 is expected to be the repetition of 12:00
        grouped_back = times.group_from(time, value, index=2, step=-3, group_hours=3)
        np_test.assert_array_equal(grouped_back, [2, 2, 2.5])
        # expected to group (3, 4) forward into 15:00, and (5) into 18:00
        grouped_forward = times.group_from(time, value, index=2, step=2, group_hours=3)
        np_test.assert_array_equal(grouped_forward, [3.5, 5])
        # redo the test including the group boundaries
        # group boundary of 16:00 is 17:00 for backward grouping
        boundary_grouped_back = times.group_from(time, value, index=2, step=-3, group_hours=3,
                                                 whole_group=True)
        np_test.assert_array_equal(boundary_grouped_back, [2, 2, 3])
        # group boundary of 16:00 is [15:00, 17:00) for forward grouping
        boundary_grouped_back = times.group_from(time, value, index=2, step=3, group_hours=3,
                                                 whole_group=True)
        np_test.assert_array_equal(boundary_grouped_back, [3, 5, 5])

    @staticmethod
    def test_running_average():
        # test averaging values backward or forward from a given value and given hours unit
        yr = '2018-01-01 '
        time = pd.to_datetime([yr + ' 12', yr + ' 13', yr + ' 15', yr + '16', yr + '17'], utc=True).tolist()
        value = [2, np.nan, 3, 4, 5]
        # expected to forward-average values into 12:00 and 15:00 (unit: 3 hours)
        forward_averaged = times.running_average(time=time, value=value, group_hours=3)
        np_test.assert_array_equal(x=[2, 2, 3, 3.5, 4], y=forward_averaged)
        # expected to backward-average
        backward_averaged = times.running_average(time=time, value=value, group_hours=3, direction=-1)
        np_test.assert_array_equal(x=[2, 0, 4, 4.5, 5], y=backward_averaged)
        # put whole group average for each member
        whole_averaged = times.running_average(time=time, value=value, group_hours=3, whole_group=True)
        np_test.assert_array_equal(x=[2, 2, 4, 4, 4], y=whole_averaged)

    def test_group_average(self):
        yr = '2018-01-01 '
        time = pd.to_datetime([yr + ' 13', yr + '14', yr + ' 15', yr + '16', yr + '17', yr + '18'], utc=True).tolist()
        value = [2, 3, 4, np.nan, 5, 6]
        expected_group_time = pd.to_datetime([yr + ' 12', yr + '15', yr + ' 18'], utc=True).tolist()
        group_time, group_average, lookup, count = times.group_average(time=time, value=value, group_hours=3)
        # time groups: 12:00, 15:00, and 18:00
        np_test.assert_array_equal(x=expected_group_time, y=group_time)
        np_test.assert_array_equal(x=[2.5, 4.5, 6], y=group_average)
        np_test.assert_array_equal(x=[2, 2, 1], y=count)
        self.assertEqual(0, lookup[expected_group_time[0]])
        self.assertEqual(1, lookup[expected_group_time[1]])
        self.assertEqual(2, lookup[expected_group_time[2]])

    @staticmethod
    def test_split():
        yr = '2018-01-01 '
        time = pd.to_datetime([yr + ' 13', yr + '14', yr + ' 15', yr + '16', yr + '18'], utc=True).tolist()
        value = [2, 3, 4, 5, 6]
        # time groups: 12:00, 15:00, and 18:00
        # averaging from group start until given index
        split = times.split(time=time, value=value, group_hours=3, step=-3, region=(0, -1))
        expected = [[2, 2, 2], [2.5, 2.5, 2.5], [2.5, 2.5, 4], [2.5, 2.5, 4.5], [2.5, 4.5, 6]]
        np_test.assert_array_equal(x=expected, y=split)

        # time groups: 12:00, 15:00, and 18:00
        # averaging from group start until given index
        split = times.split(time=time, value=value, group_hours=3, step=3, region=(0, -1))
        expected = [[2.5, 4.5, 6], [3, 4.5, 6], [4.5, 6, 6], [5, 6, 6], [6, 6, 6]]
        np_test.assert_array_equal(x=expected, y=split)

        # first index skipped
        # total group average is considered for each member regardless of index
        split = times.split(time=time, value=value, group_hours=3, step=3, region=(1, -1), whole_group=True)
        expected = [[2.5, 4.5, 6], [4.5, 6, 6], [4.5, 6, 6], [6, 6, 6]]
        np_test.assert_array_equal(x=expected, y=split)

        split = times.split(time=time, value=value, group_hours=3, step=-3, region=(4, 4), whole_group=True)
        expected = [[2.5, 4.5, 6]]
        np_test.assert_array_equal(x=expected, y=split)

    @staticmethod
    def test_one_hot():
        columns = ['01', '02', '03', '04']
        s = pd.to_datetime(pd.Series(data=[1, 2, 4]), utc=True, format='%H')
        one_hot = times.one_hot(times=s, columns=columns, time_format='%H')
        np_test.assert_array_equal(x=[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]], y=one_hot)

    def test_to_datetime(self):
        d_time = datetime.utcnow()
        d = d_time.date()
        expected = datetime.strptime(d_time.strftime('%y-%m-%d'), '%y-%m-%d')
        self.assertEqual(times.to_datetime(date=d), expected)


"""
Module colormap for creating custom color maps.  For example...
  >>> from pyclaw.plotting import colormaps
  >>> mycmap = colormaps.make_colormap({0:'r', 1.:'b'})  # red to blue
  >>> colormaps.showcolors(mycmap)   # displays resulting colormap

Note that many colormaps are also defined in matplotlib and can be set by
  >>> from matplotlib import cm
  >>> mycmap = cm.get_cmap('Greens')
for example, to get colors ranging from white to green.
See matplotlib._cm for the data defining various maps.
"""

#-------------------------
def make_colormap(colors):
#-------------------------
    """
    Define a new color map based on values specified in the dictionary
    colors, where colors[z] is the color that value z should be mapped to,
    with linear interpolation between the given values of z.

    The z values (dictionary keys) are real numbers and the values
    colors[z] can be either an RGB list, e.g. [1,0,0] for red, or an
    html hex string, e.g. "#ff0000" for red.
    """

    from matplotlib.colors import LinearSegmentedColormap, ColorConverter
    from numpy import sort

    z = sort(colors.keys())
    n = len(z)
    z1 = min(z)
    zn = max(z)
    x0 = (z - z1) / (zn - z1)

    CC = ColorConverter()
    R = []
    G = []
    B = []
    for i in range(n):
        #i'th color at level z[i]:
        Ci = colors[z[i]]
        if type(Ci) == str:
            # a hex string of form '#ff0000' for example (for red)
            RGB = CC.to_rgb(Ci)
        else:
            # assume it's an RGB triple already:
            RGB = Ci
        R.append(RGB[0])
        G.append(RGB[1])
        B.append(RGB[2])

    cmap_dict = {}
    cmap_dict['red'] = [(x0[i],R[i],R[i]) for i in range(len(R))]
    cmap_dict['green'] = [(x0[i],G[i],G[i]) for i in range(len(G))]
    cmap_dict['blue'] = [(x0[i],B[i],B[i]) for i in range(len(B))]
    mymap = LinearSegmentedColormap('mymap',cmap_dict)
    return mymap

def showcolors(cmap):
    from pylab import colorbar, clf, axes, linspace, pcolor, \
         meshgrid, show, axis, title
    #from scitools.easyviz.matplotlib_ import colorbar, clf, axes, linspace,\
                 #pcolor, meshgrid, show, colormap
    clf()
    x = linspace(0,1,21)
    X,Y = meshgrid(x,x)
    pcolor(X,Y,0.5*(X+Y), cmap=cmap, edgecolors='k')
    axis('equal')
    colorbar()
    title('Plot of x+y using colormap')

def schlieren_colormap(color=[0,0,0]):
    """
    For Schlieren plots:
    """
    from numpy import linspace, array
    if color=='k': color = [0,0,0]
    if color=='r': color = [1,0,0]
    if color=='b': color = [0,0,1]
    if color=='g': color = [0,0.5,0]
    color = array([1,1,1]) - array(color)
    s  = linspace(0,1,20)
    colors = {}
    for key in s:
        colors[key] = array([1,1,1]) - key**10 * color
    schlieren_colors = make_colormap(colors)
    return schlieren_colors

# -----------------------------------------------------------------
# Some useful colormaps follow...
# There are also many colormaps in matplotlib.cm

all_white = make_colormap({0.:'w', 1.:'w'})
all_light_red = make_colormap({0.:'#ffdddd', 1.:'#ffdddd'})
all_light_blue = make_colormap({0.:'#ddddff', 1.:'#ddddff'})
all_light_green = make_colormap({0.:'#ddffdd', 1.:'#ddffdd'})
all_light_yellow = make_colormap({0.:'#ffffdd', 1.:'#ffffdd'})

red_white_blue = make_colormap({0.:'r', 0.5:'w', 1.:'b'})
blue_white_red = make_colormap({0.:'b', 0.5:'w', 1.:'r'})
red_yellow_blue = make_colormap({0.:'r', 0.5:'#ffff00', 1.:'b'})
blue_yellow_red = make_colormap({0.:'b', 0.5:'#ffff00', 1.:'r'})
yellow_red_blue = make_colormap({0.:'#ffff00', 0.5:'r', 1.:'b'})
white_red = make_colormap({0.:'w', 1.:'r'})
white_blue = make_colormap({0.:'w', 1.:'b'})

schlieren_grays = schlieren_colormap('k')
schlieren_reds = schlieren_colormap('r')
schlieren_blues = schlieren_colormap('b')
schlieren_greens = schlieren_colormap('g')

#-------------------------------
def make_amrcolors(nlevels=4):
#-------------------------------
    """
    Make lists of colors useful for distinguishing different grids when
    plotting AMR results.

    INPUT::
       nlevels: maximum number of AMR levels expected.
    OUTPUT::
       (linecolors, bgcolors)
       linecolors = list of nlevels colors for grid lines, contour lines
       bgcolors = list of nlevels pale colors for grid background
    """

    # For 4 or less levels:
    linecolors = ['k', 'b', 'r', 'g']
    # Set bgcolors to white, then light shades of blue, red, green:
    bgcolors = ['#ffffff','#ddddff','#ffdddd','#ddffdd']
    # Set bgcolors to light shades of yellow, blue, red, green:
    #bgcolors = ['#ffffdd','#ddddff','#ffdddd','#ddffdd']

    if nlevels > 4:
        linecolors = 4*linecolors  # now has length 16
        bgcolors = 4*bgcolors
    if nlevels <= 16:
        linecolors = linecolors[:nlevels]
        bgcolors = bgcolors[:nlevels]
    else:
        print "*** Warning, suggest nlevels <= 16"

    return (linecolors, bgcolors)