import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.model_selection import cross_val_score, learning_curve, RepeatedKFold
from sklearn.metrics import classification_report, confusion_matrix, auc

def plot_conf_mat(*all_conf_mat, **format_settings):
    """Plots the confusion matrix using Seaborn Heatmap.

    Args:
        all_conf_mat (args as ndarray or pandas dataframe): The confusion matrix generated using sklearn's method.
        format_settings (keywords arguements):
            titles_list (list of str): Titles to be given to each of the plots.
            cmaps_list (list of str): cmap value to decide the color for each of the plots
                                     (refer to Matplotlib for possible values).
            num_of_rows (int): number of rows of the plots.
            num_of_cols (int): number of columns of the plots.

    Returns: None
    """
    num_of_plots = len(all_conf_mat)
    num_of_rows, num_of_cols = get_ideal_plot_dim(num_of_plots)

    all_titles = format_settings.get("titles_list", ["Confusion Matrix"]*num_of_plots)
    all_cmaps = format_settings.get("cmaps_list", ['Blues']*num_of_plots)
    nrows = format_settings.get("nrows", num_of_rows)
    ncols = format_settings.get("ncols", num_of_cols)
    figsize = format_settings.get("figsize", (12, 4))

    fig, axes = plt.subplots(nrows=nrows,
                             ncols=ncols,
                             figsize=figsize)
    if not type(axes) == np.ndarray:
        axes = [axes]

    for ax, conf_mat, title, cmap in zip(axes, all_conf_mat, all_titles, all_cmaps):
        sns.heatmap(conf_mat,
                    annot=True,
                    annot_kws={'size': 12},
                    cbar=False,
                    cmap=cmap,
                    ax=ax)

        ax.set_title(title, fontdict={'size': 16, 'weight': 'bold'}, pad=30)
        ax.set_xlabel('Predicted Label', fontdict={'size': 14, 'weight': 'bold'}, labelpad=20)
        ax.set_ylabel('True Label', fontdict={'size': 14, 'weight': 'bold'}, labelpad=20)

        loc_x = ax.get_xticks()
        loc_y = ax.get_yticks()
        loc_y -= 0.3

        ax.set_xticklabels(['Edible', 'Poisonuous'], fontdict={"fontsize": 14})
        ax.set_yticks(loc_y)
        ax.set_yticklabels(['Edible', 'Poisonuous'], fontdict={"fontsize": 14}, rotation=90)

    plt.tight_layout(rect=[0, 0.03, 1.1, 0.97])

def get_ideal_plot_dim(total_plots):
    """Calculated the best number of rows and columns to create plots clearly.

    Args:
        total_plots (int): Total number of plots to be created.

    Returns:
        nrows (int): Number of rows to be plot.
        ncols (int): Number of columns to be plot.
    """
    nrows = 1
    ncols = 3
    if total_plots <= 3:
        ncols = total_plots
        return nrows, ncols
    else:
        nrows = round(total_plots / 3) + 2
        return nrows, ncols

def metric_evaluation(model, X, y, cv: int = 5, print_results: bool = True):
    """
    The function evalutes the model on five metrics using cross validation scores for each and plots them.
    The six metrics are:
    1. Accuracy
    2. Weighted Precision
    3. Weighted Recall
    4. Weighted f1 score
    6. Area under the ROC curve

    Notes:
        The higher the value of the metric, the better.

    Args:
        model (object): An instantiated object of a sklearn classifier.
        X (np.ndarray or pd.DataFrame): Contains the features from the dataset.
        y (np.ndarray or pd.DataFrame): Contains the target or response varible from the dataset.
        cv (int): number of folds in the cross validation score.
        print_results (bool): prints the performance dataframe if True. Default True.

    Returns:
        performance_df (pd.DataFrame): A dictionary containing the four performance metrics.
    """

    # Setting up random seed to ensure all models are evaluated on same data splits
    np.random.seed(100)

    # Creating a list of metrics
    metrics_list = ['accuracy', 'precision_weighted', 'recall_weighted', 'f1_weighted', 'roc_auc']
    performance_dict = {}

    for metric in metrics_list:
        metric_score = cross_val_score(model, X, y, cv=cv, scoring=metric)
        performance_dict.update({metric: metric_score})

    performance_df = pd.DataFrame(performance_dict).round(3).T
    performance_df.columns = [f"Fold {n}" for n in range(1, len(metric_score)+1)]

    # Instantiating figure
    fig, ax = plt.subplots(figsize=(10,7))

    # plotting data
    performance_df.T.plot.bar(ax=ax)

    # Formatting title and axes
    ax.set_title("Cross Validation scores of the Model",
                fontdict={"fontsize":20, "fontweight":'bold'},
                pad=30)

    ax.set_xlabel("Cross Validation Fold number", fontsize=14, fontweight='bold', labelpad=20)
    ax.set_ylabel("Cross Validated scores of the Metric",
                fontsize=14, fontweight='bold', labelpad=20)

    plt.xticks(fontsize=12, rotation=0)
    plt.yticks(fontsize=12)

    # Formatting legend
    leg = ax.legend(fontsize=12, loc=(1.02, 0.81), frameon=True)
    leg.get_frame().set_color("#F2F2F2")
    leg.get_frame().set_edgecolor("#000000")
    leg.set_title("Estimator", prop={"size": 14, "weight": 'bold'})

    # Displaying Results
    if print_results:
        print("######### Averaged Cross Validation Scores ##########\n"
              "Accuracy score:            {accuracy: 0.2%}\n"
              "Weighted Precision score:  {precision_weighted: 0.2%}\n"
              "Weighted Recall score:     {recall_weighted: 0.2%}\n"
              "Weighted F1 score:         {f1_weighted: 0.2%}\n"
              "ROC Area Under Curve:      {roc_auc: .3f}".format(**dict(performance_df.T.mean())))

    # Print statement to know completion, incase of function being called multiple times in a cell.
    print("Model evaluation complete.")

    return performance_df

def plot_learning_curves(model, X, y):
    """
    The function plots the learning curve for the input model.

    Args:
        model (object): An instantiated object of a sklearn classifier.
        X (np.ndarray or pd.DataFrame): Contains the features from the dataset.
        y (np.ndarray or pd.DataFrame): Contains the target or response varible from the dataset.
    """
    # Setting up random seed to ensure all models are evaluated on same data splits
    np.random.seed(100)

    train_sizes, train_scores, test_scores = learning_curve(estimator=model,
                                                            X=X,
                                                            y=y,
                                                            train_sizes= np.linspace(0.1, 1.0, 10),
                                                            cv=10,
                                                            scoring='recall_weighted',random_state=100)
    train_mean = np.mean(train_scores, axis=1)
    train_std = np.std(train_scores, axis=1)
    test_mean = np.mean(test_scores, axis=1)
    test_std = np.std(test_scores, axis=1)

    plt.plot(train_sizes, train_mean,color='blue', marker='o',
             markersize=5, label='training recall')
    plt.fill_between(train_sizes, train_mean + train_std, train_mean - train_std,
                     alpha=0.15, color='blue')

    plt.plot(train_sizes, test_mean, color='green', linestyle='--', marker='s', markersize=5,
             label='validation recall')
    plt.fill_between(train_sizes, test_mean + test_std, test_mean - test_std,
                     alpha=0.15, color='green')
    plt.grid(True)
    plt.xlabel('Number of training samples')
    plt.ylabel('Recall')
    plt.legend(loc='best')
    plt.show()

def plot_box_plot(model, X, y, cv: int = 5, model_name: str = "Model", scoring: str = "recall_weighted",
                  models: list = None):
    """The function plots a box plot for the scoring parameter defined.

    Args:
        model (object): An instantiated object of a sklearn classifier.
        X (np.ndarray or pd.DataFrame): Contains the features from the dataset.
        y (np.ndarray or pd.DataFrame): Contains the target or response varible from the dataset.
        model_name (str, optional): Name of the model used. If None, uses `Decision Tree`. Defaults to None.
        scoring (str, optional): The scoring parameter to be used. Defaults to "recall_weighted".
        models (list, optional): A list of `model`s. Defaults to [].
    """
    # Setting up random seed to ensure all models are evaluated on same data splits
    np.random.seed(100)

    if models is None:
        models = []
        models.append((model_name, model))

    results =[]
    names=[]
    scoring ='recall_weighted'
    metric = scoring.replace('_', ' ').capitalize()
    print(f'Model Evaluation - {metric}')
    for name, model in models:
        # rkf = RepeatedKFold(n_splits=10, n_repeats=5, random_state=100)
        cv_results = cross_val_score(model, X, y, cv=cv, scoring=scoring)
        results.append(cv_results)
        names.append(name)
        print('{} {:.2f} +/- {:.2f}'.format(name,cv_results.mean(),cv_results.std()))
    print('\n')

    fig = plt.figure(figsize=(5,5))
    fig.suptitle('Boxplot View')
    ax = fig.add_subplot(111)
    sns.boxplot(data=results)
    ax.set_xticklabels(names)
    plt.ylabel(metric)
    plt.xlabel('Model')
    plt.show()
    pass

def full_model_evaluation(model, X, y, cv: int = 5, model_name: str = None):
    """
    The function does the following 4 things:
    1. Plot the learning curve.
    2. Plot a boxplot for weighted recall.
    3. Plot the cross validation scores of 5 metrics.
    """
    # Setting up random seed to ensure all models are evaluated on same data splits
    np.random.seed(100)

    if model_name is None:
        model_name = "Decision Tree"

    # Plot learning curve
    print(f'{model_name} Learning Curve')
    plot_learning_curves(model, X, y)

    # Model Evaluation - Boxplot
    plot_box_plot(model, X, y, cv=cv, model_name=model_name)

    # Evaluate the performance
    metric_evaluation(model, X, y, cv=cv)


import numpy as np
import matplotlib.pyplot as plt

def get_fig(fitness_array, fig_name, show = False):
    plt.figure()
    plt.title('Fitness value in x_th interation')
    plt.xlabel('x_th interation')
    plt.ylabel('fitness value')

    plt.plot(fitness_array)
    # plt.plot(x, head_fitness, color = 'blue')

    # plt.legend(['F_fitness', 'head_slap_fitness'])
    if show:
        plt.show()
    plt.savefig(fig_name)
    plt.close()


# The Python Programming Language: Functions

x = 1
y = 2
x + y

x

`add_numbers` is a function that takes two numbers and adds them together.

def add_numbers(x, y):
    return x + y

add_numbers(1, 2)

`add_numbers` updated to take an optional 3rd parameter. Using `print` allows printing of multiple expressions within a single cell.

def add_numbers(x,y,z=None):
    if (z==None):
        return x+y
    else:
        return x+y+z

print(add_numbers(1, 2))
print(add_numbers(1, 2, 3))

`add_numbers` updated to take an optional flag parameter.

def add_numbers(x, y, z=None, flag=False):
    if (flag):
        print('Flag is true!')
    if (z==None):
        return x + y
    else:
        return x + y + z

print(add_numbers(1, 2, flag=True))

Assign function `add_numbers` to variable `a`.

def add_numbers(x,y):
    return x+y

a = add_numbers
a(1,2)

# The Python Programming Language: Types and Sequences

Use `type` to return the object's type.

type('This is a string')

type(None)

type(1)

type(1.0)

type(add_numbers)

Tuples are an immutable data structure (cannot be altered).

x = (1, 'a', 2, 'b')
type(x)

Lists are a mutable data structure.

x = [1, 'a', 2, 'b']
type(x)

Use `append` to append an object to a list.

x.append(3.3)
print(x)

This is an example of how to loop through each item in the list.

for item in x:
    print(item)

Or using the indexing operator:

i=0
while( i != len(x) ):
    print(x[i])
    i = i + 1

Use `+` to concatenate lists.

[1,2] + [3,4]

Use `*` to repeat lists.

[1]*3

Use the `in` operator to check if something is inside a list.

1 in [1, 2, 3]

Now let's look at strings. Use bracket notation to slice a string.

x = 'This is a string'
print(x[0]) #first character
print(x[0:1]) #first character, but we have explicitly set the end character
print(x[0:2]) #first two characters

This will return the last element of the string.

x[-1]

This will return the slice starting from the 4th element from the end and stopping before the 2nd element from the end.

x[-4:-2]

This is a slice from the beginning of the string and stopping before the 3rd element.

x[:3]

And this is a slice starting from the 3rd element of the string and going all the way to the end.

x[3:]

firstname = 'Christopher'
lastname = 'Brooks'

print(firstname + ' ' + lastname)
print(firstname*3)
print('Chris' in firstname)

`split` returns a list of all the words in a string, or a list split on a specific character.

firstname = 'Christopher Arthur Hansen Brooks'.split(' ')[0] # [0] selects the first element of the list
lastname = 'Christopher Arthur Hansen Brooks'.split(' ')[-1] # [-1] selects the last element of the list
print(firstname)
print(lastname)

Make sure you convert objects to strings before concatenating.

'Chris' + 2

'Chris' + str(2)

Dictionaries associate keys with values.

x = {'Christopher Brooks': 'brooksch@umich.edu', 'Bill Gates': 'billg@microsoft.com'}
x['Christopher Brooks'] # Retrieve a value by using the indexing operator

x['Kevyn Collins-Thompson'] = None
x['Kevyn Collins-Thompson']

Iterate over all of the keys:

for name in x:
    print(x[name])

Iterate over all of the values:

for email in x.values():
    print(email)

Iterate over all of the items in the list:

for name, email in x.items():
    print(name)
    print(email)

You can unpack a sequence into different variables:

x = ('Christopher', 'Brooks', 'brooksch@umich.edu')
fname, lname, email = x

fname

lname

Make sure the number of values you are unpacking matches the number of variables being assigned.

x = ('Christopher', 'Brooks', 'brooksch@umich.edu', 'Ann Arbor')
fname, lname, email = x

# The Python Programming Language: More on Strings

print('Chris' + 2)

print('Chris' + str(2))

Python has a built in method for convenient string formatting.

sales_record = {
'price': 3.24,
'num_items': 4,
'person': 'Chris'}

sales_statement = '{} bought {} item(s) at a price of {} each for a total of {}'

print(sales_statement.format(sales_record['person'],
                             sales_record['num_items'],
                             sales_record['price'],
                             sales_record['num_items']*sales_record['price']))

# Reading and Writing CSV files

Let's import our datafile mpg.csv, which contains fuel economy data for 234 cars.

* mpg : miles per gallon
* class : car classification
* cty : city mpg
* cyl : # of cylinders
* displ : engine displacement in liters
* drv : f = front-wheel drive, r = rear wheel drive, 4 = 4wd
* fl : fuel (e = ethanol E85, d = diesel, r = regular, p = premium, c = CNG)
* hwy : highway mpg
* manufacturer : automobile manufacturer
* model : model of car
* trans : type of transmission
* year : model year

import csv

%precision 2

with open('mpg.csv') as csvfile:
    mpg = list(csv.DictReader(csvfile))

mpg[:3] # The first three dictionaries in our list.

`csv.Dictreader` has read in each row of our csv file as a dictionary. `len` shows that our list is comprised of 234 dictionaries.

len(mpg)

`keys` gives us the column names of our csv.

mpg[0].keys()

This is how to find the average cty fuel economy across all cars. All values in the dictionaries are strings, so we need to convert to float.

sum(float(d['cty']) for d in mpg) / len(mpg)

Similarly this is how to find the average hwy fuel economy across all cars.

sum(float(d['hwy']) for d in mpg) / len(mpg)

Use `set` to return the unique values for the number of cylinders the cars in our dataset have.

cylinders = set(d['cyl'] for d in mpg)
cylinders

Here's a more complex example where we are grouping the cars by number of cylinder, and finding the average cty mpg for each group.

CtyMpgByCyl = []

for c in cylinders: # iterate over all the cylinder levels
    summpg = 0
    cyltypecount = 0
    for d in mpg: # iterate over all dictionaries
        if d['cyl'] == c: # if the cylinder level type matches,
            summpg += float(d['cty']) # add the cty mpg
            cyltypecount += 1 # increment the count
    CtyMpgByCyl.append((c, summpg / cyltypecount)) # append the tuple ('cylinder', 'avg mpg')

CtyMpgByCyl.sort(key=lambda x: x[0])
CtyMpgByCyl

Use `set` to return the unique values for the class types in our dataset.

vehicleclass = set(d['class'] for d in mpg) # what are the class types
vehicleclass

And here's an example of how to find the average hwy mpg for each class of vehicle in our dataset.

HwyMpgByClass = []

for t in vehicleclass: # iterate over all the vehicle classes
    summpg = 0
    vclasscount = 0
    for d in mpg: # iterate over all dictionaries
        if d['class'] == t: # if the cylinder amount type matches,
            summpg += float(d['hwy']) # add the hwy mpg
            vclasscount += 1 # increment the count
    HwyMpgByClass.append((t, summpg / vclasscount)) # append the tuple ('class', 'avg mpg')

HwyMpgByClass.sort(key=lambda x: x[1])
HwyMpgByClass

# The Python Programming Language: Dates and Times

import datetime as dt
import time as tm

`time` returns the current time in seconds since the Epoch. (January 1st, 1970)

tm.time()

Convert the timestamp to datetime.

dtnow = dt.datetime.fromtimestamp(tm.time())
dtnow

Handy datetime attributes:

dtnow.year, dtnow.month, dtnow.day, dtnow.hour, dtnow.minute, dtnow.second # get year, month, day, etc.from a datetime

`timedelta` is a duration expressing the difference between two dates.

delta = dt.timedelta(days = 100) # create a timedelta of 100 days
delta

`date.today` returns the current local date.

today = dt.date.today()

today - delta # the date 100 days ago

today > today-delta # compare dates

# The Python Programming Language: Objects and map()

An example of a class in python:

class Person:
    department = 'School of Information' #a class variable

    def set_name(self, new_name): #a method
        self.name = new_name
    def set_location(self, new_location):
        self.location = new_location

person = Person()
person.set_name('Christopher Brooks')
person.set_location('Ann Arbor, MI, USA')
print('{} live in {} and works in the department {}'.format(person.name, person.location, person.department))

Here's an example of mapping the `min` function between two lists.

store1 = [10.00, 11.00, 12.34, 2.34]
store2 = [9.00, 11.10, 12.34, 2.01]
cheapest = map(min, store1, store2)
cheapest

Now let's iterate through the map object to see the values.

for item in cheapest:
    print(item)

# The Python Programming Language: Lambda and List Comprehensions

Here's an example of lambda that takes in three parameters and adds the first two.

my_function = lambda a, b, c : a + b

my_function(1, 2, 3)

Let's iterate from 0 to 999 and return the even numbers.

my_list = []
for number in range(0, 1000):
    if number % 2 == 0:
        my_list.append(number)
my_list

Now the same thing but with list comprehension.

my_list = [number for number in range(0,1000) if number % 2 == 0]
my_list

# The Python Programming Language: Numerical Python (NumPy)

import numpy as np

## Creating Arrays

Create a list and convert it to a numpy array

mylist = [1, 2, 3]
x = np.array(mylist)
x

Or just pass in a list directly

y = np.array([4, 5, 6])
y

Pass in a list of lists to create a multidimensional array.

m = np.array([[7, 8, 9], [10, 11, 12]])
m

Use the shape method to find the dimensions of the array. (rows, columns)

m.shape

`arange` returns evenly spaced values within a given interval.

n = np.arange(0, 30, 2) # start at 0 count up by 2, stop before 30
n

`reshape` returns an array with the same data with a new shape.

n = n.reshape(3, 5) # reshape array to be 3x5
n

`linspace` returns evenly spaced numbers over a specified interval.

o = np.linspace(0, 4, 9) # return 9 evenly spaced values from 0 to 4
o

`resize` changes the shape and size of array in-place.

o.resize(3, 3)
o

`ones` returns a new array of given shape and type, filled with ones.

np.ones((3, 2))

`zeros` returns a new array of given shape and type, filled with zeros.

np.zeros((2, 3))

`eye` returns a 2-D array with ones on the diagonal and zeros elsewhere.

np.eye(3)

`diag` extracts a diagonal or constructs a diagonal array.

np.diag(y)

Create an array using repeating list (or see `np.tile`)

np.array([1, 2, 3] * 3)

Repeat elements of an array using `repeat`.

np.repeat([1, 2, 3], 3)

#### Combining Arrays

p = np.ones([2, 3], int)
p

Use `vstack` to stack arrays in sequence vertically (row wise).

np.vstack([p, 2*p])

Use `hstack` to stack arrays in sequence horizontally (column wise).

np.hstack([p, 2*p])

## Operations

Use `+`, `-`, `*`, `/` and `**` to perform element wise addition, subtraction, multiplication, division and power.

print(x + y) # elementwise addition     [1 2 3] + [4 5 6] = [5  7  9]
print(x - y) # elementwise subtraction  [1 2 3] - [4 5 6] = [-3 -3 -3]

print(x * y) # elementwise multiplication  [1 2 3] * [4 5 6] = [4  10  18]
print(x / y) # elementwise divison         [1 2 3] / [4 5 6] = [0.25  0.4  0.5]

print(x**2) # elementwise power  [1 2 3] ^2 =  [1 4 9]

**Dot Product:**

$ \begin{bmatrix}x_1 \ x_2 \ x_3\end{bmatrix}
\cdot
\begin{bmatrix}y_1 \\ y_2 \\ y_3\end{bmatrix}
= x_1 y_1 + x_2 y_2 + x_3 y_3$

x.dot(y) # dot product  1*4 + 2*5 + 3*6

z = np.array([y, y**2])
print(len(z)) # number of rows of array

Let's look at transposing arrays. Transposing permutes the dimensions of the array.

z = np.array([y, y**2])
z

The shape of array `z` is `(2,3)` before transposing.

z.shape

Use `.T` to get the transpose.

z.T

The number of rows has swapped with the number of columns.

z.T.shape

Use `.dtype` to see the data type of the elements in the array.

z.dtype

Use `.astype` to cast to a specific type.

z = z.astype('f')
z.dtype

## Math Functions

Numpy has many built in math functions that can be performed on arrays.

a = np.array([-4, -2, 1, 3, 5])

a.sum()

a.max()

a.min()

a.mean()

a.std()

`argmax` and `argmin` return the index of the maximum and minimum values in the array.

a.argmax()

a.argmin()

## Indexing / Slicing

s = np.arange(13)**2
s

Use bracket notation to get the value at a specific index. Remember that indexing starts at 0.

s[0], s[4], s[-1]

Use `:` to indicate a range. `array[start:stop]`

Leaving `start` or `stop` empty will default to the beginning/end of the array.

s[1:5]

Use negatives to count from the back.

s[-4:]

A second `:` can be used to indicate step-size. `array[start:stop:stepsize]`

Here we are starting 5th element from the end, and counting backwards by 2 until the beginning of the array is reached.

s[-5::-2]

Let's look at a multidimensional array.

r = np.arange(36)
r.resize((6, 6))
r

Use bracket notation to slice: `array[row, column]`

r[2, 2]

And use : to select a range of rows or columns

r[3, 3:6]

Here we are selecting all the rows up to (and not including) row 2, and all the columns up to (and not including) the last column.

r[:2, :-1]

This is a slice of the last row, and only every other element.

r[-1, ::2]

We can also perform conditional indexing. Here we are selecting values from the array that are greater than 30. (Also see `np.where`)

r[r > 30]

Here we are assigning all values in the array that are greater than 30 to the value of 30.

r[r > 30] = 30
r

## Copying Data

Be careful with copying and modifying arrays in NumPy!

`r2` is a slice of `r`

r2 = r[:3,:3]
r2

Set this slice's values to zero ([:] selects the entire array)

r2[:] = 0
r2

`r` has also been changed!

r

To avoid this, use `r.copy` to create a copy that will not affect the original array

r_copy = r.copy()
r_copy

Now when r_copy is modified, r will not be changed.

r_copy[:] = 10
print(r_copy, '\n')
print(r)

### Iterating Over Arrays

Let's create a new 4 by 3 array of random numbers 0-9.

test = np.random.randint(0, 10, (4,3))
test

Iterate by row:

for row in test:
    print(row)

Iterate by index:

for i in range(len(test)):
    print(test[i])

Iterate by row and index:

for i, row in enumerate(test):
    print('row', i, 'is', row)

Use `zip` to iterate over multiple iterables.

test2 = test**2
test2

for i, j in zip(test, test2):
    print(i,'+',j,'=',i+j)

book/_build/jupyter_execute/pandas/Week 1-Introduction to Data Science[Coursera].py


# S.D. Peckham
# June 10, 2009
#-------------------

import time
import numpy

def loop_test(n=1000000):

    #---------------------------------
    # For loop using Python's range
    #---------------------------------
    start  = time.time()
    my_sum = 0

    for k in range(n):
        my_sum += 1

    run_time = (time.time() - start)
    print 'run time with range =', run_time

    #---------------------------------
    # For loop using Python's xrange
    #---------------------------------
    start  = time.time()
    my_sum = 0

    for k in xrange(n):
        my_sum += 1

    run_time = (time.time() - start)
    print 'run time with xrange =', run_time

    #--------------------------------
    # For loop using NumPy's arange
    #---------------------------------
    start  = time.time()
    my_sum = 0

    for k in numpy.arange(n):
        my_sum += 1

    run_time = (time.time() - start)
    print 'run time with arange =', run_time


"""
=============================================================
Customizing Figure Layouts Using GridSpec and Other Functions
=============================================================

How to create grid-shaped combinations of axes.

    :func:`~matplotlib.pyplot.subplots`
        Perhaps the primary function used to create figures and axes.
        It's also similar to :func:`.matplotlib.pyplot.subplot`,
        but creates and places all axes on the figure at once.  See also
        `matplotlib.Figure.subplots`.

    :class:`~matplotlib.gridspec.GridSpec`
        Specifies the geometry of the grid that a subplot will be
        placed. The number of rows and number of columns of the grid
        need to be set. Optionally, the subplot layout parameters
        (e.g., left, right, etc.) can be tuned.

    :class:`~matplotlib.gridspec.SubplotSpec`
        Specifies the location of the subplot in the given *GridSpec*.

    :func:`~matplotlib.pyplot.subplot2grid`
        A helper function that is similar to
        :func:`~matplotlib.pyplot.subplot`,
        but uses 0-based indexing and let subplot to occupy multiple cells.
        This function is not covered in this tutorial.

"""

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

############################################################################
# Basic Quickstart Guide
# ======================
#
# These first two examples show how to create a basic 2-by-2 grid using
# both :func:`~matplotlib.pyplot.subplots` and :mod:`~matplotlib.gridspec`.
#
# Using :func:`~matplotlib.pyplot.subplots` is quite simple.
# It returns a :class:`~matplotlib.figure.Figure` instance and an array of
# :class:`~matplotlib.axes.Axes` objects.

fig1, f1_axes = plt.subplots(ncols=2, nrows=2, constrained_layout=True)

############################################################################
# For a simple use case such as this, :mod:`~matplotlib.gridspec` is
# perhaps overly verbose.
# You have to create the figure and :class:`~matplotlib.gridspec.GridSpec`
# instance separately, then pass elements of gridspec instance to the
# :func:`~matplotlib.figure.Figure.add_subplot` method to create the axes
# objects.
# The elements of the gridspec are accessed in generally the same manner as
# numpy arrays.

fig2 = plt.figure(constrained_layout=True)
spec2 = gridspec.GridSpec(ncols=2, nrows=2, figure=fig2)
f2_ax1 = fig2.add_subplot(spec2[0, 0])
f2_ax2 = fig2.add_subplot(spec2[0, 1])
f2_ax3 = fig2.add_subplot(spec2[1, 0])
f2_ax4 = fig2.add_subplot(spec2[1, 1])

#############################################################################
# The power of gridspec comes in being able to create subplots that span
# rows and columns.  Note the
# `Numpy slice <https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html>`_
# syntax for selecing the part of the gridspec each subplot will occupy.
#
# Note that we have also used the convenience method `.Figure.add_gridspec`
# instead of `.gridspec.GridSpec`, potentially saving the user an import,
# and keeping the namespace cleaner.

fig3 = plt.figure(constrained_layout=True)
gs = fig3.add_gridspec(3, 3)
f3_ax1 = fig3.add_subplot(gs[0, :])
f3_ax1.set_title('gs[0, :]')
f3_ax2 = fig3.add_subplot(gs[1, :-1])
f3_ax2.set_title('gs[1, :-1]')
f3_ax3 = fig3.add_subplot(gs[1:, -1])
f3_ax3.set_title('gs[1:, -1]')
f3_ax4 = fig3.add_subplot(gs[-1, 0])
f3_ax4.set_title('gs[-1, 0]')
f3_ax5 = fig3.add_subplot(gs[-1, -2])
f3_ax5.set_title('gs[-1, -2]')

#############################################################################
# :mod:`~matplotlib.gridspec` is also indispensable for creating subplots
# of different widths via a couple of methods.
#
# The method shown here is similar to the one above and initializes a
# uniform grid specification,
# and then uses numpy indexing and slices to allocate multiple
# "cells" for a given subplot.

fig4 = plt.figure(constrained_layout=True)
spec4 = fig4.add_gridspec(ncols=2, nrows=2)
anno_opts = dict(xy=(0.5, 0.5), xycoords='axes fraction',
                 va='center', ha='center')

f4_ax1 = fig4.add_subplot(spec4[0, 0])
f4_ax1.annotate('GridSpec[0, 0]', **anno_opts)
fig4.add_subplot(spec4[0, 1]).annotate('GridSpec[0, 1:]', **anno_opts)
fig4.add_subplot(spec4[1, 0]).annotate('GridSpec[1:, 0]', **anno_opts)
fig4.add_subplot(spec4[1, 1]).annotate('GridSpec[1:, 1:]', **anno_opts)

############################################################################
# Another option is to use the ``width_ratios`` and ``height_ratios``
# parameters. These keyword arguments are lists of numbers.
# Note that absolute values are meaningless, only their relative ratios
# matter. That means that ``width_ratios=[2, 4, 8]`` is equivalent to
# ``width_ratios=[1, 2, 4]`` within equally wide figures.
# For the sake of demonstration, we'll blindly create the axes within
# ``for`` loops since we won't need them later.

fig5 = plt.figure(constrained_layout=True)
widths = [2, 3, 1.5]
heights = [1, 3, 2]
spec5 = fig5.add_gridspec(ncols=3, nrows=3, width_ratios=widths,
                          height_ratios=heights)
for row in range(3):
    for col in range(3):
        ax = fig5.add_subplot(spec5[row, col])
        label = 'Width: {}\nHeight: {}'.format(widths[col], heights[row])
        ax.annotate(label, (0.1, 0.5), xycoords='axes fraction', va='center')

############################################################################
# Learning to use ``width_ratios`` and ``height_ratios`` is particularly
# useful since the top-level function :func:`~matplotlib.pyplot.subplots`
# accepts them within the ``gridspec_kw`` parameter.
# For that matter, any parameter accepted by
# :class:`~matplotlib.gridspec.GridSpec` can be passed to
# :func:`~matplotlib.pyplot.subplots` via the ``gridspec_kw`` parameter.
# This example recreates the previous figure without directly using a
# gridspec instance.

gs_kw = dict(width_ratios=widths, height_ratios=heights)
fig6, f6_axes = plt.subplots(ncols=3, nrows=3, constrained_layout=True,
        gridspec_kw=gs_kw)
for r, row in enumerate(f6_axes):
    for c, ax in enumerate(row):
        label = 'Width: {}\nHeight: {}'.format(widths[c], heights[r])
        ax.annotate(label, (0.1, 0.5), xycoords='axes fraction', va='center')

############################################################################
# The ``subplots`` and ``gridspec`` methods can be combined since it is
# sometimes more convenient to make most of the subplots using ``subplots``
# and then remove some and combine them.  Here we create a layout with
# the bottom two axes in the last column combined.

fig7, f7_axs = plt.subplots(ncols=3, nrows=3)
gs = f7_axs[1, 2].get_gridspec()
# remove the underlying axes
for ax in f7_axs[1:, -1]:
    ax.remove()
axbig = fig7.add_subplot(gs[1:, -1])
axbig.annotate('Big Axes \nGridSpec[1:, -1]', (0.1, 0.5),
               xycoords='axes fraction', va='center')

fig7.tight_layout()

###############################################################################
# Fine Adjustments to a Gridspec Layout
# =====================================
#
# When a GridSpec is explicitly used, you can adjust the layout
# parameters of subplots that are created from the GridSpec.  Note this
# option is not compatible with ``constrained_layout`` or
# `.Figure.tight_layout` which both adjust subplot sizes to fill the
# figure.

fig8 = plt.figure(constrained_layout=False)
gs1 = fig8.add_gridspec(nrows=3, ncols=3, left=0.05, right=0.48, wspace=0.05)
f8_ax1 = fig8.add_subplot(gs1[:-1, :])
f8_ax2 = fig8.add_subplot(gs1[-1, :-1])
f8_ax3 = fig8.add_subplot(gs1[-1, -1])

###############################################################################
# This is similar to :func:`~matplotlib.pyplot.subplots_adjust`, but it only
# affects the subplots that are created from the given GridSpec.
#
# For example, compare the left and right sides of this figure:

fig9 = plt.figure(constrained_layout=False)
gs1 = fig9.add_gridspec(nrows=3, ncols=3, left=0.05, right=0.48,
                        wspace=0.05)
f9_ax1 = fig9.add_subplot(gs1[:-1, :])
f9_ax2 = fig9.add_subplot(gs1[-1, :-1])
f9_ax3 = fig9.add_subplot(gs1[-1, -1])

gs2 = fig9.add_gridspec(nrows=3, ncols=3, left=0.55, right=0.98,
                        hspace=0.05)
f9_ax4 = fig9.add_subplot(gs2[:, :-1])
f9_ax5 = fig9.add_subplot(gs2[:-1, -1])
f9_ax6 = fig9.add_subplot(gs2[-1, -1])

###############################################################################
# GridSpec using SubplotSpec
# ==========================
#
# You can create GridSpec from the :class:`~matplotlib.gridspec.SubplotSpec`,
# in which case its layout parameters are set to that of the location of
# the given SubplotSpec.
#
# Note this is also available from the more verbose
# `.gridspec.GridSpecFromSubplotSpec`.

fig10 = plt.figure(constrained_layout=True)
gs0 = fig10.add_gridspec(1, 2)

gs00 = gs0[0].subgridspec(2, 3)
gs01 = gs0[1].subgridspec(3, 2)

for a in range(2):
    for b in range(3):
        fig10.add_subplot(gs00[a, b])
        fig10.add_subplot(gs01[b, a])

###############################################################################
# A Complex Nested GridSpec using SubplotSpec
# ===========================================
#
# Here's a more sophisticated example of nested GridSpec where we put
# a box around each cell of the outer 4x4 grid, by hiding appropriate
# spines in each of the inner 3x3 grids.

import numpy as np
from itertools import product

def squiggle_xy(a, b, c, d, i=np.arange(0.0, 2*np.pi, 0.05)):
    return np.sin(i*a)*np.cos(i*b), np.sin(i*c)*np.cos(i*d)

fig11 = plt.figure(figsize=(8, 8), constrained_layout=False)

# gridspec inside gridspec
outer_grid = fig11.add_gridspec(4, 4, wspace=0.0, hspace=0.0)

for i in range(16):
    inner_grid = outer_grid[i].subgridspec(3, 3, wspace=0.0, hspace=0.0)
    a, b = int(i/4)+1, i % 4+1
    for j, (c, d) in enumerate(product(range(1, 4), repeat=2)):
        ax = plt.Subplot(fig11, inner_grid[j])
        ax.plot(*squiggle_xy(a, b, c, d))
        ax.set_xticks([])
        ax.set_yticks([])
        fig11.add_subplot(ax)

all_axes = fig11.get_axes()

# show only the outside spines
for ax in all_axes:
    for sp in ax.spines.values():
        sp.set_visible(False)
    if ax.is_first_row():
        ax.spines['top'].set_visible(True)
    if ax.is_last_row():
        ax.spines['bottom'].set_visible(True)
    if ax.is_first_col():
        ax.spines['left'].set_visible(True)
    if ax.is_last_col():
        ax.spines['right'].set_visible(True)

plt.show()

#############################################################################
#
# ------------
#
# References
# """"""""""
#
# The usage of the following functions and methods is shown in this example:

matplotlib.pyplot.subplots
matplotlib.figure.Figure.add_gridspec
matplotlib.figure.Figure.add_subplot
matplotlib.gridspec.GridSpec
matplotlib.gridspec.SubplotSpec.subgridspec
matplotlib.gridspec.GridSpecFromSubplotSpec


# ---
# jupyter:
#   jupytext:
#     formats: py:percent,ipynb
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # 1 Initial Approach with Dynamic Programming
# This notebook contains some intial trials of hands-on, basic reinforcement learning algorithms to familiarize with the topic and eventually arrive at the bare frame of the environment for ramp-up scenarios. The set-up follows the discourse in [Artificial Intelligence: Reinforcement Learning in Python](https://www.udemy.com/share/1013kmBEQbdF1XTXQ=/) by Lazy Programmer Inc roughly, altough the environment conforms with OpenAI's toolkit for RL, [Gym](https://gym.openai.com/).

# %%
import gym
import numpy as np
import pandas as pd
from gym import spaces
from IPython.core.display import HTML, display

# %% [markdown]
# ## Review helper functions

# %% [markdown]
# - With `display_side_by_side()`, tables are displayed side by side to ease interpretations and save space as the table index is similar.
# - Future value at each possible combination of state and episode is generated and displayed with `pivot_v()`
# - The policy at each possible combination of state and episode is generated and displayed with `pivot_p()`

# %%
def display_side_by_side(dfs: list, captions: list):
    output = ""
    combined = dict(zip(captions, dfs))
    for caption, df in combined.items():
        output += (
            df.style.set_table_attributes("style='display:inline'")
            .set_caption(caption)
            ._repr_html_()
        )
        output += "\xa0\xa0\xa0"
    display(HTML(output))

def pivot_v(V):
    v_table = pd.Series(V).reset_index()
    v_table.columns = ["Episode", "Action", "Value"]
    v_table.sort_values(by=["Episode", "Action"], inplace=True)
    v_table_p = v_table.pivot(index="Episode", columns="Action")
    return v_table, v_table_p

def pivot_p(policy):
    p_table = pd.Series(policy).reset_index()
    p_table.columns = ["Episode", "State", "Policy_S"]
    p_table.sort_values(by=["Episode"], inplace=True)
    p_table = p_table.pivot(index="Episode", columns="State")
    p_table = p_table.assign(demand=demand)
    return p_table

# %% [markdown]
# ## Definitions of constants
#
# All constants needed for the environment include
# - the **action list**, a discrete set of actions to move between states
# - the **legal changes**, a constraining factor which actions may follow upon which state
# - **demand**, for now just a list which is the trajectory that needs to be followed during an episode, only the choice of actions is open to the agent

# %%
action_list = [
    "OPERATE",  # 0
    "PREPARE",  # 1
    "PARK",  # 2
    "STORE",  # 3
]

legal_changes = {
    0: {0, 1, 2},  # from OPERATE to OPERATE, PREPARE, or PARK
    1: {1, 0, 2},  # from PREPARE to PREPARE, OPERATE, or PARK
    2: {2, 1, 3},  # from PARK to PARK, PREPARE, or STORE
    3: {3, 2},  # from STORE to STORE or PARK
}
demand = [
    0,
    0,
    0,
    0,
    0,
    0,
    0,  # week 1
    0,
    0,
    0,
    2,
    0,
    3,
    1,  # week 2
    0,
    0,
    60,
    25,
    32,
    87,
    56,  # week 3
    92,
    83,
    29,
    40,
    86,
    70,
    45,  # week 4
]

# %% [markdown]
# ## Environment Definition
# The environment is in accordance with Gym. It's a sequence of X days, on each day a single action can be performed. That's the entire logic encoded in the environment!

# %%
class RampupEnv(gym.Env):
    def __init__(self, horizon=10, verbose=False):
        super(RampupEnv, self).__init__()

        self.fleet_size = 1
        self.horizon = horizon

        self.state_time = 0  # time component of state
        self.state_status = np.zeros(self.horizon)  # status component of state

        # action_space: The Space object corresponding to valid actions
        # observation_space: The Space object corresponding to valid observations
        self.action_space = spaces.Discrete(len(action_list))
        self.observation_space = spaces.Tuple(
            (spaces.Discrete(self.fleet_size), spaces.Discrete(self.horizon))
        )

        self.total_reward = 0
        self.verbose = verbose

    def translate_action(self, action):
        if action == 0:
            action_description = "OPERATE"
            if demand[self.state_time] > 80:
                reward = 15000
            else:
                reward = -3000
        elif action == 1:
            action_description = "PREPARE"
            reward = -2000
        elif action == 2:
            action_description = "PARK"
            reward = -1000
        elif action == 3:
            action_description = "STORE"
            reward = -500

        if self.verbose:
            print(f"Action {action} leads to state {action_description}")

        return action_description, reward

    def step(self, action):
        # Increment time component of state
        self.state_time += 1

        action_description, reward = self.translate_action(action)

        # Adjust the status component of state
        self.state_status[self.state_time] = action

        obs = self.state_status[0 : self.state_time]
        done = bool(self.state_time == self.horizon - 1)
        info = {}

        if self.verbose:
            print(f"Action taken: {action_description}")

        self.total_reward += reward

        return obs, reward, done, info

    def undo_step(self, o, r, d, i):
        self.state_time -= 1
        self.total_reward -= r

        return o, r, d, i

    def render(self):
        if self.verbose:
            print(f"Reward so far: {self.total_reward}")

    def reset(self):
        # Initialize the agent in the first field
        self.state_time = 0
        self.total_reward = 0

    def set_state(self, state):
        state_time, state_status = state
        self.state_time = state_time
        self.state_status[self.state_time] = state_status

# %% [markdown]
# ## Random Walk Episode
# The length of the episode is determined by `horizon`, actions are taken at random. We observe, that games mostly lead to a negative total reward, thus it takes a policy better than a random policy to perform well.

# %%
horizon = 7 * 4
env = RampupEnv(horizon=horizon)

for step in range(horizon):
    action = env.action_space.sample()
    obs, reward, done, info = env.step(action)
    env.render()
    if done:
        print(
            f"Random walk episode is complete.\nTrajectory:\n{obs}\nTotal reward: {env.total_reward}"
        )
        break

# %% [markdown]
# ## Value Iteration
# All possible state transitions (transitions between possible combinations of state time and state actions) are iteratively evaluated for their value, until the improvement in value converges to a defined threshold. It takes 23 iterations for the sample environment with the provided data. The resulting table shows the value for each state-action.

# %%
h = 7 * 4  # time horizon
env = RampupEnv(horizon=h, verbose=False)
all_states = set(
    [(s, a) for s in range(len(env.state_status)) for a in range(env.action_space.n)]
)

SMALL_ENOUGH = 1e-3  # threshold for convergence
GAMMA = 1.0  # discount factor

V = {}
for state in all_states:
    V[state] = 0

biggest_change = 0
iters = 0

while True:
    iters += 1
    biggest_change = 0
    for state in all_states:
        old_v = V[state]
        env.set_state(state)
        state_time, state_status = state
        if env.state_time < env.horizon - 1:
            new_v = 0  # answer is accumulated
            p_a = 1.0 / len(legal_changes[state_status])  # equal probability
            for a in legal_changes[state_status]:
                env.set_state(state)
                o, r, d, i = env.step(a)
                new_v += p_a * (r + GAMMA * V[(env.state_time, a)])
                o, r, d, i = env.undo_step(o, r, d, i)
            V[state] = new_v
            biggest_change = max(biggest_change, np.abs(old_v - V[state]))

        if env.verbose:
            print(f"State: {state}, Old V: {old_v}")

    if biggest_change < SMALL_ENOUGH:
        print(f"Done in {iters} iterations")
        v_table, v_table_p = pivot_v(V)
        display(v_table_p)
        break

# %% [markdown]
# ## Policy Iteration
# Value iteration stops as the improvement of value converges, policy iteration then adds a policy improvement step. In the case of the example, there is no improvement visible, but we do get the policy for the values from value iterations inside of the attached policy improvement step.

# %%
h = 7 * 4  # time horizon
env = RampupEnv(horizon=h, verbose=False)
all_states = set(
    [(s, a) for s in range(len(env.state_status)) for a in range(env.action_space.n)]
)

SMALL_ENOUGH = 1e-3
GAMMA = 1

policy = {}
for state in all_states:
    policy[state] = env.action_space.sample()

V = {}
for state in all_states:
    V[state] = 0

iters1, iters2 = 0, 0
rep1, rep2 = "", ""

while True:
    iters1 = 0
    # Policy Evaluation Step
    while True:
        iters1 += 1
        biggest_change = 0
        for state in all_states:
            old_v = V[state]
            env.set_state(state)
            state_time, state_status = state
            if env.state_time < env.horizon - 1:
                new_v = 0  # answer is accumulated
                p_a = 1.0 / len(legal_changes[state_status])  # equal probability
                for a in legal_changes[state_status]:
                    env.set_state(state)
                    o, r, d, i = env.step(a)
                    new_v += p_a * (r + GAMMA * V[(env.state_time, a)])
                    o, r, d, i = env.undo_step(o, r, d, i)
                V[state] = new_v
                biggest_change = max(biggest_change, np.abs(old_v - V[state]))

            if env.verbose:
                print(f"State: {state}, Old V: {old_v}")

        if biggest_change < SMALL_ENOUGH:
            rep1 += f"{str(iters1)}, "
            iters2 += 1
            break

    # Policy Improvement Step
    is_policy_converged = True
    for state in all_states:
        for state in policy:
            old_a = policy[state]
            new_a = None
            best_value = float("-inf")
            env.set_state(state)
            state_time, state_status = state
            if env.state_time < env.horizon - 1:
                # loop through all possible actions to find the best current action
                for a in legal_changes[state_status]:
                    env.set_state(state)
                    o, r, d, i = env.step(a)
                    v = r + GAMMA * V[(env.state_time, a)]
                    # print(v)
                    if v > best_value:
                        best_value = v
                        new_a = a
                policy[state] = new_a
                if new_a != old_a:
                    is_policy_converged = False
    if is_policy_converged:
        print(
            f"Policy evaluation and improvement steps were entered {iters2} times. \
Policy evaluation iterated {rep1[:-2]} times."
        )
        v_table, v_table_p = pivot_v(V)
        p_table = pivot_p(policy)
        display_side_by_side([v_table_p, p_table], ["Value Table", "Policy Table"])
        break

# %%


#
# author: Jungtaek Kim (jtkim@postech.ac.kr)
# last updated: August 17, 2021
#
"""test_trees_trees_common"""

import typing
import pytest
import numpy as np

from bayeso import constants
from bayeso.trees import trees_common as package_target

TEST_EPSILON = 1e-7

def test_get_inputs_from_leaf_typing():
    annos = package_target.get_inputs_from_leaf.__annotations__

    assert annos['leaf'] == list
    assert annos['return'] == np.ndarray

def test_get_inputs_from_leaf():
    leaf = [
        (np.array([1.0, 2.0, 3.0]), np.array([2.0])),
        (np.array([2.0, 1.0, 3.0]), np.array([1.0])),
        (np.array([3.0, 0.0, 3.0]), np.array([0.5])),
        (np.array([4.0, -1.0, 3.0]), np.array([1.5])),
    ]

    with pytest.raises(AssertionError) as error:
        package_target.get_inputs_from_leaf(123)
    with pytest.raises(AssertionError) as error:
        package_target.get_inputs_from_leaf('abc')

    inputs = package_target.get_inputs_from_leaf(leaf)
    assert len(inputs.shape) == 2
    assert inputs.shape[0] == len(leaf)
    assert inputs.shape[1] == leaf[0][0].shape[0]

def test_get_outputs_from_leaf_typing():
    annos = package_target.get_outputs_from_leaf.__annotations__

    assert annos['leaf'] == list
    assert annos['return'] == np.ndarray

def test_get_outputs_from_leaf():
    leaf = [
        (np.array([1.0, 2.0, 3.0]), np.array([2.0])),
        (np.array([2.0, 1.0, 3.0]), np.array([1.0])),
        (np.array([3.0, 0.0, 3.0]), np.array([0.5])),
        (np.array([4.0, -1.0, 3.0]), np.array([1.5])),
    ]

    with pytest.raises(AssertionError) as error:
        package_target.get_outputs_from_leaf(123)
    with pytest.raises(AssertionError) as error:
        package_target.get_outputs_from_leaf('abc')

    outputs = package_target.get_outputs_from_leaf(leaf)
    assert len(outputs.shape) == 2
    assert outputs.shape[0] == len(leaf)
    assert outputs.shape[1] == leaf[0][1].shape[0]

def test__mse_typing():
    annos = package_target._mse.__annotations__

    assert annos['Y'] == np.ndarray
    assert annos['return'] == float

def test__mse():
    Y = np.array([
        [1.0],
        [2.0],
        [6.0],
    ])

    with pytest.raises(AssertionError) as error:
        package_target._mse(123)
    with pytest.raises(AssertionError) as error:
        package_target._mse('abc')

    output = package_target._mse(np.zeros((0, 1)))
    assert output == 1e8

    output = package_target._mse(np.array([]))
    assert output == 1e8

    output = package_target._mse(Y)
    assert output == 14.0 / 3

def test_mse_typing():
    annos = package_target.mse.__annotations__

    assert annos['left_right'] == tuple
    assert annos['return'] == float

def test_mse():
    left = [
        (np.array([1.0, 2.0, 3.0]), np.array([0.5])),
        (np.array([2.0, 2.0, 1.0]), np.array([0.1])),
        (np.array([3.0, 0.0, 1.0]), np.array([0.2])),
        (np.array([4.0, 0.0, 3.0]), np.array([0.3])),
    ]
    right = [
        (np.array([10.0, 20.0, 30.0]), np.array([0.6])),
        (np.array([20.0, 20.0, 10.0]), np.array([0.9])),
        (np.array([30.0, 0.0, 10.0]), np.array([0.9])),
        (np.array([40.0, 0.0, 30.0]), np.array([0.8])),
    ]

    with pytest.raises(AssertionError) as error:
        package_target.mse(123)
    with pytest.raises(AssertionError) as error:
        package_target.mse('abc')
    with pytest.raises(AssertionError) as error:
        package_target.mse(np.zeros((4, 1)))

    output = package_target.mse(([], []))
    assert output == 2e8

    output = package_target.mse((left, []))
    assert output == 0.021875 + 1e8

    output = package_target.mse(([], right))
    assert output == 1e8 + 0.015

    output = package_target.mse((left, right))
    assert np.abs(output - (0.021875 + 0.015)) < TEST_EPSILON

def test_subsample_typing():
    annos = package_target.subsample.__annotations__

    assert annos['X'] == np.ndarray
    assert annos['Y'] == np.ndarray
    assert annos['ratio_sampling'] == float
    assert annos['replace_samples'] == bool
    assert annos['return'] == constants.TYPING_TUPLE_TWO_ARRAYS

def test_subsample():
    np.random.seed(42)

    X = np.reshape(np.arange(0, 40), (10, 4))
    Y = np.random.randn(10, 1)
    ratio_sampling = 0.5
    replace_samples = False

    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, Y, ratio_sampling, 123)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, Y, ratio_sampling, 'abc')
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, Y, 1, replace_samples)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, Y, 'abc', replace_samples)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, 'abc', ratio_sampling, replace_samples)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, 123, ratio_sampling, replace_samples)
    with pytest.raises(AssertionError) as error:
        package_target.subsample('abc', Y, ratio_sampling, replace_samples)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(123, Y, ratio_sampling, replace_samples)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, Y, 4.0, False)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, Y, -0.5, False)
    with pytest.raises(AssertionError) as error:
        package_target.subsample(X, Y, -0.5, True)

    X_, Y_ = package_target.subsample(X, Y, ratio_sampling, replace_samples)
    print(X_)
    print(Y_)

    X_truth = np.array([
        [24, 25, 26, 27],
        [8, 9, 10, 11],
        [32, 33, 34, 35],
        [28, 29, 30, 31],
        [36, 37, 38, 39],
    ])
    Y_truth = np.array([
        [1.57921282],
        [0.64768854],
        [-0.46947439],
        [0.76743473],
        [0.54256004],
    ])

    assert np.all(np.abs(X_truth - X_) < TEST_EPSILON)
    assert np.all(np.abs(Y_truth - Y_) < TEST_EPSILON)
    assert X_.shape[0] == Y_.shape[0] == X_truth.shape[0] == Y_truth.shape[0]
    assert X_.shape[0] == Y_.shape[0] == int(ratio_sampling * X.shape[0])

    X_, Y_ = package_target.subsample(X, Y, 1.2, True)
    print(X_)
    print(Y_)

    X_truth = np.array([
        [36, 37, 38, 39],
        [8, 9, 10, 11],
        [24, 25, 26, 27],
        [12, 13, 14, 15],
        [32, 33, 34, 35],
        [8, 9, 10, 11],
        [16, 17, 18, 19],
        [8, 9, 10, 11],
        [24, 25, 26, 27],
        [16, 17, 18, 19],
        [32, 33, 34, 35],
        [24, 25, 26, 27],
    ])
    Y_truth = np.array([
        [0.54256004],
        [0.64768854],
        [1.57921282],
        [1.52302986],
        [-0.46947439],
        [0.64768854],
        [-0.23415337],
        [0.64768854],
        [1.57921282],
        [-0.23415337],
        [-0.46947439],
        [1.57921282],
    ])

    assert np.all(np.abs(X_truth - X_) < TEST_EPSILON)
    assert np.all(np.abs(Y_truth - Y_) < TEST_EPSILON)
    assert X_.shape[0] == Y_.shape[0] == X_truth.shape[0] == Y_truth.shape[0]
    assert X_.shape[0] == Y_.shape[0] == int(1.2 * X.shape[0])

def test__split_left_right_typing():
    annos = package_target._split_left_right.__annotations__

    assert annos['X'] == np.ndarray
    assert annos['Y'] == np.ndarray
    assert annos['dim_to_split'] == int
    assert annos['val_to_split'] == float
    assert annos['return'] == tuple

def test__split_left_right():
    np.random.seed(42)

    X = np.reshape(np.arange(0, 40), (10, 4))
    Y = np.random.randn(10, 1)
    dim_to_split = 1
    val_to_split = 14.0

    with pytest.raises(AssertionError) as error:
        package_target._split_left_right(X, Y, dim_to_split, 'abc')
    with pytest.raises(AssertionError) as error:
        package_target._split_left_right(X, Y, dim_to_split, 123)
    with pytest.raises(AssertionError) as error:
        package_target._split_left_right(X, Y, 0.5, val_to_split)
    with pytest.raises(AssertionError) as error:
        package_target._split_left_right(X, Y, 'abc', val_to_split)
    with pytest.raises(AssertionError) as error:
        package_target._split_left_right(X, 'abc', dim_to_split, val_to_split)
    with pytest.raises(AssertionError) as error:
        package_target._split_left_right('abc', Y, dim_to_split, val_to_split)

    left, right = package_target._split_left_right(X, Y, dim_to_split, val_to_split)
    print(left)
    print(right)

    left_ = [
        (np.array([0, 1, 2, 3]), np.array([0.49671415])),
        (np.array([4, 5, 6, 7]), np.array([-0.1382643])),
        (np.array([8, 9, 10, 11]), np.array([0.64768854])),
        (np.array([12, 13, 14, 15]), np.array([1.52302986]))
    ]
    right_ = [
        (np.array([16, 17, 18, 19]), np.array([-0.23415337])),
        (np.array([20, 21, 22, 23]), np.array([-0.23413696])),
        (np.array([24, 25, 26, 27]), np.array([1.57921282])),
        (np.array([28, 29, 30, 31]), np.array([0.76743473])),
        (np.array([32, 33, 34, 35]), np.array([-0.46947439])),
        (np.array([36, 37, 38, 39]), np.array([0.54256004]))
    ]

    assert len(left_) == 4
    assert len(right_) == 6
    assert (len(left_) + len(right_)) == X.shape[0]

def test__split_typing():
    annos = package_target._split.__annotations__

    assert annos['X'] == np.ndarray
    assert annos['Y'] == np.ndarray
    assert annos['num_features'] == int
    assert annos['split_random_location'] == bool
    assert annos['return'] == dict

def test__split():
    np.random.seed(42)

    X = np.reshape(np.arange(0, 40), (10, 4))
    Y = np.random.randn(10, 1)
    num_features = 2
    split_random_location = False

    with pytest.raises(AssertionError) as error:
        package_target._split(X, Y, num_features, 'abc')
    with pytest.raises(AssertionError) as error:
        package_target._split(X, Y, 'abc', split_random_location)
    with pytest.raises(AssertionError) as error:
        package_target._split(X, Y, 2.0, split_random_location)
    with pytest.raises(AssertionError) as error:
        package_target._split(X, 'abc', num_features, split_random_location)
    with pytest.raises(AssertionError) as error:
        package_target._split('abc', Y, num_features, split_random_location)

    dict_split = package_target._split(X, Y, num_features, split_random_location)
    print(dict_split)
    print(dict_split['index'])
    print(dict_split['value'])
    print(dict_split['left_right'])

    assert isinstance(dict_split, dict)
    assert dict_split['index'] == 1
    assert dict_split['value'] == 35.0

    assert np.all(dict_split['left_right'][0][0][0] == np.array([0, 1, 2, 3]))
    assert np.abs(dict_split['left_right'][0][0][1] - np.array([0.49671415])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][1][0] == np.array([4, 5, 6, 7]))
    assert np.abs(dict_split['left_right'][0][1][1] - np.array([-0.1382643])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][2][0] == np.array([8, 9, 10, 11]))
    assert np.abs(dict_split['left_right'][0][2][1] - np.array([0.64768854])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][3][0] == np.array([12, 13, 14, 15]))
    assert np.abs(dict_split['left_right'][0][3][1] - np.array([1.52302986])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][4][0] == np.array([16, 17, 18, 19]))
    assert np.abs(dict_split['left_right'][0][4][1] - np.array([-0.23415337])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][5][0] == np.array([20, 21, 22, 23]))
    assert np.abs(dict_split['left_right'][0][5][1] - np.array([-0.23413696])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][6][0] == np.array([24, 25, 26, 27]))
    assert np.abs(dict_split['left_right'][0][6][1] - np.array([1.57921282])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][7][0] == np.array([28, 29, 30, 31]))
    assert np.abs(dict_split['left_right'][0][7][1] - np.array([0.76743473])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][0][8][0] == np.array([32, 33, 34, 35]))
    assert np.abs(dict_split['left_right'][0][8][1] - np.array([-0.46947439])) < TEST_EPSILON

    assert np.all(dict_split['left_right'][1][0][0] == np.array([36, 37, 38, 39]))
    assert np.abs(dict_split['left_right'][1][0][1] - np.array([0.54256004])) < TEST_EPSILON

    dict_split = package_target._split(X, Y, num_features, True)
    print(dict_split)
    print(dict_split['index'])
    print(dict_split['value'])
    print(dict_split['left_right'])

    assert isinstance(dict_split, dict)
    assert dict_split['index'] == 0
    assert dict_split['value'] == 0.2543869879098266

    X = np.ones(X.shape)

    dict_split = package_target._split(X, Y, num_features, True)
    print(dict_split)
    print(dict_split['index'])
    print(dict_split['value'])
    print(dict_split['left_right'])

    assert isinstance(dict_split, dict)
    assert dict_split['index'] == 1
    assert dict_split['value'] == 1.0

def test_split_typing():
    annos = package_target.split.__annotations__

    assert annos['node'] == dict
    assert annos['depth_max'] == int
    assert annos['size_min_leaf'] == int
    assert annos['num_features'] == int
    assert annos['split_random_location'] == bool
    assert annos['cur_depth'] == int
    assert annos['return'] == constants.TYPE_NONE

def test_split():
    np.random.seed(42)

    X = np.reshape(np.arange(0, 40), (10, 4))
    Y = np.random.randn(10, 1)
    depth_max = 4
    size_min_leaf = 2
    num_features = 2
    split_random_location = False

    node = package_target._split(X, Y, num_features, split_random_location)

    with pytest.raises(AssertionError) as error:
        package_target.split(node, depth_max, size_min_leaf, num_features, split_random_location, 'abc')
    with pytest.raises(AssertionError) as error:
        package_target.split(node, depth_max, size_min_leaf, num_features, split_random_location, 1.0)
    with pytest.raises(AssertionError) as error:
        package_target.split(node, depth_max, size_min_leaf, num_features, 'abc', 1)
    with pytest.raises(AssertionError) as error:
        package_target.split(node, depth_max, size_min_leaf, 'abc', split_random_location, 1)
    with pytest.raises(AssertionError) as error:
        package_target.split(node, depth_max, 'abc', num_features, split_random_location, 1)
    with pytest.raises(AssertionError) as error:
        package_target.split(node, 'abc', size_min_leaf, num_features, split_random_location, 1)
    with pytest.raises(AssertionError) as error:
        package_target.split(X, depth_max, size_min_leaf, num_features, split_random_location, 1)
    with pytest.raises(AssertionError) as error:
        package_target.split('abc', depth_max, size_min_leaf, num_features, split_random_location, 1)

    package_target.split(node, depth_max, size_min_leaf, num_features, split_random_location, 1)
    assert isinstance(node, dict)

def test__predict_by_tree_typing():
    annos = package_target._predict_by_tree.__annotations__

    assert annos['bx'] == np.ndarray
    assert annos['tree'] == dict
    assert annos['return'] == constants.TYPING_TUPLE_TWO_FLOATS

def test__predict_by_tree():
    np.random.seed(42)

    X = np.reshape(np.arange(0, 40), (10, 4))
    Y = np.random.randn(10, 1)
    depth_max = 4
    size_min_leaf = 2
    num_features = 2
    split_random_location = True

    node = package_target._split(X, Y, num_features, split_random_location)
    package_target.split(node, depth_max, size_min_leaf, num_features, split_random_location, 1)

    with pytest.raises(AssertionError) as error:
        package_target._predict_by_tree(np.array([4.0, 2.0, 3.0, 1.0]), 'abc')
    with pytest.raises(AssertionError) as error:
        package_target._predict_by_tree(X, node)
    with pytest.raises(AssertionError) as error:
        package_target._predict_by_tree('abc', node)

    mean, std = package_target._predict_by_tree(np.array([4.0, 2.0, 3.0, 1.0]), node)
    print(mean)
    print(std)

    assert mean == 0.179224925920024
    assert std == 0.31748922709120864

def test__predict_by_trees_typing():
    annos = package_target._predict_by_trees.__annotations__

    assert annos['bx'] == np.ndarray
    assert annos['list_trees'] == list
    assert annos['return'] == constants.TYPING_TUPLE_TWO_FLOATS

def test__predict_by_trees():
    np.random.seed(42)

    X = np.reshape(np.arange(0, 40), (10, 4))
    Y = np.random.randn(10, 1)
    depth_max = 4
    size_min_leaf = 2
    num_features = 2
    split_random_location = True

    node_1 = package_target._split(X, Y, num_features, split_random_location)
    package_target.split(node_1, depth_max, size_min_leaf, num_features, split_random_location, 1)

    node_2 = package_target._split(X, Y, num_features, split_random_location)
    package_target.split(node_2, depth_max, size_min_leaf, num_features, split_random_location, 1)

    node_3 = package_target._split(X, Y, num_features, split_random_location)
    package_target.split(node_3, depth_max, size_min_leaf, num_features, split_random_location, 1)

    list_trees = [node_1, node_2, node_3]

    with pytest.raises(AssertionError) as error:
        package_target._predict_by_trees(np.array([4.0, 2.0, 3.0, 1.0]), node_1)
    with pytest.raises(AssertionError) as error:
        package_target._predict_by_trees(np.array([4.0, 2.0, 3.0, 1.0]), 'abc')
    with pytest.raises(AssertionError) as error:
        package_target._predict_by_trees(X, list_trees)
    with pytest.raises(AssertionError) as error:
        package_target._predict_by_trees('abc', list_trees)

    mean, std = package_target._predict_by_trees(np.array([4.0, 2.0, 3.0, 1.0]), list_trees)
    print(mean)
    print(std)

    assert mean == 0.12544669602080652
    assert std == 0.3333040901154691

def test_predict_by_trees_typing():
    annos = package_target.predict_by_trees.__annotations__

    assert annos['X'] == np.ndarray
    assert annos['list_trees'] == list
    assert annos['return'] == constants.TYPING_TUPLE_TWO_ARRAYS

def test_predict_by_trees():
    np.random.seed(42)

    X = np.reshape(np.arange(0, 40), (10, 4))
    Y = np.random.randn(10, 1)
    depth_max = 4
    size_min_leaf = 2
    num_features = 2
    split_random_location = True

    node_1 = package_target._split(X, Y, num_features, split_random_location)
    package_target.split(node_1, depth_max, size_min_leaf, num_features, split_random_location, 1)

    node_2 = package_target._split(X, Y, num_features, split_random_location)
    package_target.split(node_2, depth_max, size_min_leaf, num_features, split_random_location, 1)

    node_3 = package_target._split(X, Y, num_features, split_random_location)
    package_target.split(node_3, depth_max, size_min_leaf, num_features, split_random_location, 1)

    list_trees = [node_1, node_2, node_3]

    with pytest.raises(AssertionError) as error:
        package_target.predict_by_trees(np.array([4.0, 2.0, 3.0, 1.0]), node_1)
    with pytest.raises(AssertionError) as error:
        package_target.predict_by_trees(np.array([4.0, 2.0, 3.0, 1.0]), 'abc')
    with pytest.raises(AssertionError) as error:
        package_target.predict_by_trees(np.array([4.0, 2.0, 3.0, 1.0]), list_trees)
    with pytest.raises(AssertionError) as error:
        package_target.predict_by_trees('abc', list_trees)

    means, stds = package_target.predict_by_trees(X, list_trees)
    print(means)
    print(stds)

    means_truth = np.array([
        [0.33710618],
        [0.1254467],
        [0.68947573],
        [0.9866563],
        [0.16257799],
        [0.16258346],
        [0.85148454],
        [0.55084716],
        [0.30721653],
        [0.30721653],
    ])

    stds_truth = np.array([
        [0.29842457],
        [0.33330409],
        [0.44398864],
        [0.72536523],
        [0.74232577],
        [0.74232284],
        [0.83388663],
        [0.5615399],
        [0.64331582],
        [0.64331582],
    ])

    assert isinstance(means, np.ndarray)
    assert isinstance(stds, np.ndarray)
    assert len(means.shape) == 2
    assert len(stds.shape) == 2
    assert means.shape[0] == stds.shape[0] == X.shape[0]
    assert means.shape[1] == stds.shape[1] == 1

    assert np.all(np.abs(means - means_truth) < TEST_EPSILON)
    assert np.all(np.abs(stds - stds_truth) < TEST_EPSILON)

    X = np.random.randn(1000, 4)

    means, stds = package_target.predict_by_trees(X, list_trees)

    assert isinstance(means, np.ndarray)
    assert isinstance(stds, np.ndarray)
    assert len(means.shape) == 2
    assert len(stds.shape) == 2
    assert means.shape[0] == stds.shape[0] == X.shape[0]
    assert means.shape[1] == stds.shape[1] == 1

def test_compute_sigma_typing():
    annos = package_target.compute_sigma.__annotations__

    assert annos['preds_mu_leaf'] == np.ndarray
    assert annos['preds_sigma_leaf'] == np.ndarray
    assert annos['min_sigma'] == float
    assert annos['return'] == np.ndarray

def test_compute_sigma():
    means_leaf = np.array([
        1.0,
        2.0,
        3.0,
        9.0,
        8.0,
        4.0,
        5.0,
        6.0,
        7.0,
        10.0,
    ])

    stds_leaf = np.array([
        -1.0,
        0.0,
        1.0,
        2.0,
        1.0,
        1.0,
        4.0,
        3.0,
        4.0,
        -2.0,
    ])
    min_sigma = 0.0

    with pytest.raises(AssertionError) as error:
        package_target.compute_sigma(means_leaf, stds_leaf, min_sigma='abc')
    with pytest.raises(AssertionError) as error:
        package_target.compute_sigma(means_leaf, stds_leaf, min_sigma=4)
    with pytest.raises(AssertionError) as error:
        package_target.compute_sigma(means_leaf, 'abc', min_sigma=min_sigma)
    with pytest.raises(AssertionError) as error:
        package_target.compute_sigma(means_leaf, np.array([[1.0], [2.0], [1.0]]), min_sigma=min_sigma)
    with pytest.raises(AssertionError) as error:
        package_target.compute_sigma('abc', stds_leaf, min_sigma=min_sigma)
    with pytest.raises(AssertionError) as error:
        package_target.compute_sigma(np.array([[1.0], [2.0], [1.0]]), stds_leaf, min_sigma=min_sigma)

    sigma = package_target.compute_sigma(means_leaf, stds_leaf, min_sigma=min_sigma)
    print(sigma)

    sigma_truth = np.mean(means_leaf**2 + np.maximum(stds_leaf, min_sigma)**2)
    print(sigma_truth)
    sigma_truth -= np.mean(means_leaf)**2
    print(sigma_truth)
    sigma_truth = np.sqrt(sigma_truth)
    print(sigma_truth)

    assert sigma == sigma_truth == 3.612478373637688