# coding: utf-8

# ## Rover Project Test Notebook
# This notebook contains the functions from the lesson and provides the scaffolding you need to test out your mapping methods.  The steps you need to complete in this notebook for the project are the following:
#
# * First just run each of the cells in the notebook, examine the code and the results of each.
# * Run the simulator in "Training Mode" and record some data. Note: the simulator may crash if you try to record a large (longer than a few minutes) dataset, but you don't need a ton of data, just some example images to work with.
# * Change the data directory path (2 cells below) to be the directory where you saved data
# * Test out the functions provided on your data
# * Write new functions (or modify existing ones) to report and map out detections of obstacles and rock samples (yellow rocks)
# * Populate the `process_image()` function with the appropriate steps/functions to go from a raw image to a worldmap.
# * Run the cell that calls `process_image()` using `moviepy` functions to create video output
# * Once you have mapping working, move on to modifying `perception.py` and `decision.py` to allow your rover to navigate and map in autonomous mode!
#
# **Note: If, at any point, you encounter frozen display windows or other confounding issues, you can always start again with a clean slate by going to the "Kernel" menu above and selecting "Restart & Clear Output".**
#
# **Run the next cell to get code highlighting in the markdown cells.**

# In[2]:

get_ipython().run_cell_magic('HTML', '', '<style> code {background-color : orange !important;} </style>')

# In[1]:

get_ipython().run_line_magic('matplotlib', 'inline')
#%matplotlib qt # Choose %matplotlib qt to plot to an interactive window (note it may show up behind your browser)
# Make some of the relevant imports
import cv2 # OpenCV for perspective transform
import numpy as np
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import scipy.misc # For saving images as needed
import glob  # For reading in a list of images from a folder
import imageio
imageio.plugins.ffmpeg.download()

# ## Quick Look at the Data
# There's some example data provided in the `test_dataset` folder.  This basic dataset is enough to get you up and running but if you want to hone your methods more carefully you should record some data of your own to sample various scenarios in the simulator.
#
# Next, read in and display a random image from the `test_dataset` folder

# In[3]:

path = '../test_dataset/IMG/*'
img_list = glob.glob(path)
# Grab a random image and display it
idx = np.random.randint(0, len(img_list)-1)
image = mpimg.imread(img_list[idx])
plt.imshow(image)

# ## Calibration Data
# Read in and display example grid and rock sample calibration images.  You'll use the grid for perspective transform and the rock image for creating a new color selection that identifies these samples of interest.

# In[4]:

# In the simulator you can toggle on a grid on the ground for calibration
# You can also toggle on the rock samples with the 0 (zero) key.
# Here's an example of the grid and one of the rocks
example_grid = '../calibration_images/example_grid1.jpg'
example_rock = '../calibration_images/example_rock1.jpg'
grid_img = mpimg.imread(example_grid)
rock_img = mpimg.imread(example_rock)

fig = plt.figure(figsize=(12,3))
plt.subplot(121)
plt.imshow(grid_img)
plt.subplot(122)
plt.imshow(rock_img)

# ## Perspective Transform
#
# Define the perspective transform function from the lesson and test it on an image.

# In[6]:

# Define a function to perform a perspective transform
# I've used the example grid image above to choose source points for the
# grid cell in front of the rover (each grid cell is 1 square meter in the sim)
# Define a function to perform a perspective transform
def perspect_transform(img, src, dst):

    M = cv2.getPerspectiveTransform(src, dst)
    warped = cv2.warpPerspective(img, M, (img.shape[1], img.shape[0]))# keep same size as input image

    return warped

# Define calibration box in source (actual) and destination (desired) coordinates
# These source and destination points are defined to warp the image
# to a grid where each 10x10 pixel square represents 1 square meter
# The destination box will be 2*dst_size on each side
dst_size = 5
# Set a bottom offset to account for the fact that the bottom of the image
# is not the position of the rover but a bit in front of it
# this is just a rough guess, feel free to change it!
bottom_offset = 6
source = np.float32([[14, 140], [301 ,140],[200, 96], [118, 96]])
destination = np.float32([[image.shape[1]/2 - dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  [image.shape[1]/2 - dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  ])
warped = perspect_transform(grid_img, source, destination)
plt.imshow(warped)
#scipy.misc.imsave('../output/warped_example.jpg', warped)

# ## Color Thresholding
# Define the color thresholding function from the lesson and apply it to the warped image
#
# **TODO:** Ultimately, you want your map to not just include navigable terrain but also obstacles and the positions of the rock samples you're searching for.  Modify this function or write a new function that returns the pixel locations of obstacles (areas below the threshold) and rock samples (yellow rocks in calibration images), such that you can map these areas into world coordinates as well.
# **Hints and Suggestion:**
# * For obstacles you can just invert your color selection that you used to detect ground pixels, i.e., if you've decided that everything above the threshold is navigable terrain, then everthing below the threshold must be an obstacle!
#
#
# * For rocks, think about imposing a lower and upper boundary in your color selection to be more specific about choosing colors.  You can investigate the colors of the rocks (the RGB pixel values) in an interactive matplotlib window to get a feel for the appropriate threshold range (keep in mind you may want different ranges for each of R, G and B!).  Feel free to get creative and even bring in functions from other libraries.  Here's an example of [color selection](http://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_imgproc/py_colorspaces/py_colorspaces.html) using OpenCV.
#
# * **Beware However:** if you start manipulating images with OpenCV, keep in mind that it defaults to `BGR` instead of `RGB` color space when reading/writing images, so things can get confusing.

# In[7]:

# Identify pixels above the threshold
# Threshold of RGB > 160 does a nice job of identifying ground pixels only
def color_thresh(img, rgb_thresh=(160, 160, 160)):
    # Create an array of zeros same xy size as img, but single channel
    color_select = np.zeros_like(img[:,:,0])
    # Require that each pixel be above all three threshold values in RGB
    # above_thresh will now contain a boolean array with "True"
    # where threshold was met
    above_thresh = (img[:,:,0] > rgb_thresh[0])                 & (img[:,:,1] > rgb_thresh[1])                 & (img[:,:,2] > rgb_thresh[2])
    # Index the array of zeros with the boolean array and set to 1
    color_select[above_thresh] = 1
    # Return the binary image
    return color_select

fig = plt.figure(figsize=(12,7))
fig.tight_layout()

threshed = color_thresh(warped)
plt.subplot(221); plt.imshow(threshed, cmap='gray')

#scipy.misc.imsave('../output/warped_threshed.jpg', threshed*255)
plt.show()

# In[8]:

# Identify pixels above the range using hsv color space
def color_in_range_by_hsv(bgr_img, hsv_min, hsv_max):
    hsv_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2HSV)  # rgb to hsv
    lower = np.array(hsv_min, np.uint8)   # to numpy array
    upper = np.array(hsv_max, np.uint8)   # to numpy array
    return cv2.inRange(hsv_img, lower, upper)  # filter

def rgb_to_hsv_color(rgb_color):
    hsv_color = cv2.cvtColor(np.uint8([[rgb_color]]), cv2.COLOR_RGB2HSV)[0][0]

    print('origin rgb -> hsv color(sv%):', rgb_color, hsv_color, (hsv_color[1]/255, hsv_color[2]/255))
    return hsv_color

# rock color in hsv
hsv_rock_color = rgb_to_hsv_color((150, 130, 18))

fig = plt.figure(figsize=(21,7))
fig.tight_layout()

# For rocks (yellow)
plt.subplot(231); plt.imshow(rock_img)

# color in range
hsv_rock_min = (60, 120, 110)
hsv_rock_max = (135, 255, 255)
threshed_hsb_rock_img = color_in_range_by_hsv(rock_img, hsv_rock_min, hsv_rock_max)

plt.subplot(232); plt.imshow(threshed_hsb_rock_img, cmap='gray')

# masked
masked_hsb_rock_img = cv2.bitwise_and(rock_img, rock_img, mask=threshed_hsb_rock_img)
plt.subplot(233); plt.imshow(masked_hsb_rock_img)

# perspect
warped_with_rock = perspect_transform(rock_img, source, destination)
plt.subplot(234); plt.imshow(warped_with_rock)

# color in range
threshed_warped_rock_img = color_in_range_by_hsv(warped_with_rock, hsv_rock_min, hsv_rock_max)
plt.subplot(235); plt.imshow(threshed_warped_rock_img, cmap='gray')

# masked
masked_warped_rock_img = cv2.bitwise_and(warped_with_rock, warped_with_rock, mask=threshed_warped_rock_img)
plt.subplot(236); plt.imshow(masked_warped_rock_img)

# In[9]:

## Find rock (岩を探す)
def color_in_range_by_hsv(bgr_img, hsv_min, hsv_max):
    # ref. https://docs.opencv.org/3.1.0/d7/d1b/group__imgproc__misc.html#ga397ae87e1288a81d2363b61574eb8cab
    hsv_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2HSV)  # rgb to hsv

    # ref. https://docs.opencv.org/3.1.0/d2/de8/group__core__array.html#ga48af0ab51e36436c5d04340e036ce981
    lower = np.array(hsv_min, np.uint8)   # to numpy array
    upper = np.array(hsv_max, np.uint8)   # to numpy array
    return cv2.inRange(hsv_img, lower, upper)  # 0 or 255 before filtering

def find_rocks_by_hsv(img, hsv_min = (60, 120, 110), hsv_max = (135, 255, 255)):
    # return 0 or 255 image
    return color_in_range_by_hsv(img, hsv_min, hsv_max)

threshed_rocks_img = find_rocks_by_hsv(rock_img)
not_rocks_img = find_rocks_by_hsv(image)
print('rocks:', threshed_rocks_img.any(), np.count_nonzero(threshed_rocks_img))
print('not rocks: ', not_rocks_img.any(), np.count_nonzero(not_rocks_img))

fig = plt.figure(figsize=(14,7))
plt.subplot(121); plt.imshow(rock_img)
plt.subplot(122); plt.imshow(threshed_rocks_img, cmap='gray')

# In[10]:

# Identify pixels above the range
def color_in_range_by_bgr(bgr_img, bgr_min, bgr_max):
    lower = np.array(bgr_min, np.uint8)   # to numpy array
    upper = np.array(bgr_max, np.uint8)   # to numpy array
    return cv2.inRange(bgr_img, lower, upper)  # filter

rgb_to_hsv_color((0, 0, 0)) # dark obstacle
rgb_to_hsv_color((73, 68, 65)) # light obstacle
rgb_to_hsv_color((114, 105, 99)) # sky
rgb_to_hsv_color((197, 178, 163)) # ground
rgb_to_hsv_color((189, 157, 0)) # rock

# In[11]:

## For obstacles only (mask) 障害のみ
# make fill image
full_img = np.ones_like(rock_img[:,:,0])
# transform image
warped_full_img = perspect_transform(full_img, source, destination)

fig = plt.figure(figsize=(14,7))
fig.tight_layout()

plt.subplot(221); plt.imshow(rock_img)
plt.subplot(222); plt.imshow(warped_with_rock)
plt.subplot(223); plt.imshow(full_img, cmap='gray')
plt.subplot(224); plt.imshow(warped_full_img, cmap='gray')  #

# In[21]:

## For navigable terrain only 道のみ

#hsv_navi_min = (10, 30, 120)
#hsv_navi_max = (255, 230, 255)

fig = plt.figure(figsize=(20,10))
fig.tight_layout()

plt.subplot(221); plt.imshow(rock_img)
plt.subplot(222); plt.imshow(warped_with_rock)

normal_threshed_navi = color_thresh(warped_with_rock)
#_img = color_in_range_by_hsv(warped_with_rock, hsv_navi_min, hsv_navi_max)

plt.subplot(223); plt.imshow(normal_threshed_navi, cmap='gray')
threshed_navi = normal_threshed_navi

# Reduce mask
kernel = np.ones((5,5), np.uint8)
# kernel = np.array([
#     [0,1,1,1,0],
#     [0,1,1,1,0],
#     [0,1,1,1,0],
#     [0,0,0,0,0],
#     [0,0,0,0,0]], np.uint8)
threshed_navi_img = cv2.erode(normal_threshed_navi, kernel)

plt.subplot(224); plt.imshow(threshed_navi_img, cmap='gray')

# ## Coordinate Transformations
# Define the functions used to do coordinate transforms and apply them to an image.

# In[13]:

# Define a function to convert from image coords to rover coords
def rover_coords(binary_img):
    # Identify nonzero pixels
    ypos, xpos = binary_img.nonzero()
    # Calculate pixel positions with reference to the rover position being at the
    # center bottom of the image.
    x_pixel = -(ypos - binary_img.shape[0]).astype(np.float)
    y_pixel = -(xpos - binary_img.shape[1]/2 ).astype(np.float)
    return x_pixel, y_pixel

# Define a function to convert to radial coords in rover space
def to_polar_coords(x_pixel, y_pixel):
    # Convert (x_pixel, y_pixel) to (distance, angle)
    # in polar coordinates in rover space
    # Calculate distance to each pixel
    dist = np.sqrt(x_pixel**2 + y_pixel**2)
    # Calculate angle away from vertical for each pixel
    angles = np.arctan2(y_pixel, x_pixel)
    return dist, angles

# Define a function to map rover space pixels to world space
def rotate_pix(xpix, ypix, yaw):
    # Convert yaw to radians
    yaw_rad = yaw * np.pi / 180
    xpix_rotated = (xpix * np.cos(yaw_rad)) - (ypix * np.sin(yaw_rad))

    ypix_rotated = (xpix * np.sin(yaw_rad)) + (ypix * np.cos(yaw_rad))
    # Return the result
    return xpix_rotated, ypix_rotated

def translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale):
    # Apply a scaling and a translation
    xpix_translated = (xpix_rot / scale) + xpos
    ypix_translated = (ypix_rot / scale) + ypos
    # Return the result
    return xpix_translated, ypix_translated

# Define a function to apply rotation and translation (and clipping)
# Once you define the two functions above this function should work
def pix_to_world(xpix, ypix, xpos, ypos, yaw, world_size, scale):
    # Apply rotation
    xpix_rot, ypix_rot = rotate_pix(xpix, ypix, yaw)
    # Apply translation
    xpix_tran, ypix_tran = translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale)
    # Perform rotation, translation and clipping all at once
    x_pix_world = np.clip(np.int_(xpix_tran), 0, world_size - 1)
    y_pix_world = np.clip(np.int_(ypix_tran), 0, world_size - 1)
    # Return the result
    return x_pix_world, y_pix_world

# Grab another random image
idx = np.random.randint(0, len(img_list)-1)
image = mpimg.imread(img_list[idx])
warped = perspect_transform(image, source, destination)
threshed = color_thresh(warped)

# Calculate pixel values in rover-centric coords and distance/angle to all pixels
xpix, ypix = rover_coords(threshed)
dist, angles = to_polar_coords(xpix, ypix)
mean_dir = np.mean(angles)

# Do some plotting
fig = plt.figure(figsize=(12,9))
plt.subplot(221)
plt.imshow(image)
plt.subplot(222)
plt.imshow(warped)
plt.subplot(223)
plt.imshow(threshed, cmap='gray')
plt.subplot(224)
plt.plot(xpix, ypix, '.')
plt.ylim(-160, 160)
plt.xlim(0, 160)
arrow_length = 100
x_arrow = arrow_length * np.cos(mean_dir)
y_arrow = arrow_length * np.sin(mean_dir)
plt.arrow(0, 0, x_arrow, y_arrow, color='red', zorder=2, head_width=10, width=2)

# ## Read in saved data and ground truth map of the world
# The next cell is all setup to read your saved data into a `pandas` dataframe.  Here you'll also read in a "ground truth" map of the world, where white pixels (pixel value = 1) represent navigable terrain.
#
# After that, we'll define a class to store telemetry data and pathnames to images.  When you instantiate this class (`data = Databucket()`) you'll have a global variable called `data` that you can refer to for telemetry and map data within the `process_image()` function in the following cell.
#

# In[14]:

# Import pandas and read in csv file as a dataframe
import pandas as pd
# Change the path below to your data directory
# If you are in a locale (e.g., Europe) that uses ',' as the decimal separator
# change the '.' to ','
df = pd.read_csv('../test_dataset/robot_log.csv', delimiter=';', decimal='.')
csv_img_list = df["Path"].tolist() # Create list of image pathnames
# Read in ground truth map and create a 3-channel image with it
ground_truth = mpimg.imread('../calibration_images/map_bw.png')
ground_truth_3d = np.dstack((ground_truth*0, ground_truth*255, ground_truth*0)).astype(np.float)

# Creating a class to be the data container
# Will read in saved data from csv file and populate this object
# Worldmap is instantiated as 200 x 200 grids corresponding
# to a 200m x 200m space (same size as the ground truth map: 200 x 200 pixels)
# This encompasses the full range of output position values in x and y from the sim
class Databucket():
    def __init__(self):
        self.images = csv_img_list
        self.xpos = df["X_Position"].values
        self.ypos = df["Y_Position"].values
        self.yaw = df["Yaw"].values
        self.count = 0 # This will be a running index
        self.worldmap = np.zeros((200, 200, 3)).astype(np.float)
        self.ground_truth = ground_truth_3d # Ground truth worldmap

# Instantiate a Databucket().. this will be a global variable/object
# that you can refer to in the process_image() function below
data = Databucket()

# In[15]:

############# Test of make worldmap ################

# a camera image to navigable terrain/obstacles/rock samples
_worldmap = np.copy(data.worldmap)
# todo: update _worldmap
print(_worldmap.shape)

_index = 0
_img = mpimg.imread(data.images[_index])
_scale = 10
_xpos = data.xpos[_index]
_ypos = data.ypos[_index]
_yaw = data.yaw[_index]
print(_img.shape, _xpos, _ypos, _yaw)

# fill for obstacles terrain to
_full_img = np.ones_like(_img[:,:,0])
_warped_full_img = perspect_transform(_full_img, source, destination)
_xpix, _ypix = rover_coords(_warped_full_img)
_x_world, _y_world = pix_to_world(_xpix, _ypix, _xpos, _ypos, _yaw, _worldmap.shape[0], _scale)
_worldmap[_y_world,_x_world,:] = (255,0,0)

# fill for navigable terrain to blue
_warped_img = perspect_transform(_img, source, destination)
_threshed_navi = color_thresh(_warped_img)
_xpix, _ypix = rover_coords(_threshed_navi)
_x_world, _y_world = pix_to_world(_xpix, _ypix, _xpos, _ypos, _yaw, _worldmap.shape[0], _scale)
_worldmap[_y_world,_x_world,:] = (0,0,255)

# fill for rock samples
_threshed_hsb_rock_img = color_in_range_by_hsv(_warped_img, hsv_rock_min, hsv_rock_max)
_xpix, _ypix = rover_coords(_threshed_hsb_rock_img)
_x_world, _y_world = pix_to_world(_xpix, _ypix, _xpos, _ypos, _yaw, _worldmap.shape[0], _scale)
_worldmap[_y_world,_x_world,:] = (255,255,255)

# overlay
_map_add = cv2.addWeighted(_worldmap, 1, data.ground_truth, 0.5, 0)
_flippud_map_add = np.flipud(_map_add)

# show
fig = plt.figure(figsize=(18,5))
print(np.max(_worldmap), np.max(data.ground_truth), np.max(_map_add))
plt.subplot(131); plt.imshow(np.uint8(_worldmap))
plt.subplot(132); plt.imshow(np.uint8(data.ground_truth))
plt.subplot(133); plt.imshow(np.uint8(_map_add))

# ## Write a function to process stored images
#
# Modify the `process_image()` function below by adding in the perception step processes (functions defined above) to perform image analysis and mapping.  The following cell is all set up to use this `process_image()` function in conjunction with the `moviepy` video processing package to create a video from the images you saved taking data in the simulator.
#
# In short, you will be passing individual images into `process_image()` and building up an image called `output_image` that will be stored as one frame of video.  You can make a mosaic of the various steps of your analysis process and add text as you like (example provided below).
#
#
#
# To start with, you can simply run the next three cells to see what happens, but then go ahead and modify them such that the output video demonstrates your mapping process.  Feel free to get creative!

# In[16]:

# Define a function to pass stored images to
# reading rover position and yaw angle from csv file
# This function will be used by moviepy to create an output video
def process_image(img):
    # Example of how to use the Databucket() object defined above
    # to print the current x, y and yaw values
    # print(data.xpos[data.count], data.ypos[data.count], data.yaw[data.count])

    # DONE:
    # 1) Define source and destination points for perspective transform
    dst_size = 5
    bottom_offset = 6
    source = np.float32([[14, 140], [301 ,140],[200, 96], [118, 96]])
    destination = np.float32([[img.shape[1]/2 - dst_size, img.shape[0] - bottom_offset],
                  [img.shape[1]/2 + dst_size, img.shape[0] - bottom_offset],
                  [img.shape[1]/2 + dst_size, img.shape[0] - 2*dst_size - bottom_offset],
                  [img.shape[1]/2 - dst_size, img.shape[0] - 2*dst_size - bottom_offset],
                  ])

    scale = 10
    world_size = data.worldmap.shape[0]
    xpos, ypos, yaw = (data.xpos[data.count], data.ypos[data.count], data.yaw[data.count])

    # 2) Apply perspective transform
    warped_img = perspect_transform(img, source, destination)
    filled_img = np.ones_like(img[:,:,0])
    warped_obst_img = perspect_transform(filled_img, source, destination)

    # 3) Apply color threshold to identify navigable terrain/obstacles/rock samples
    threshed_navi_img = color_thresh(warped_img)
    # reserve mask only(obst_img - img)
    threshed_obst_img =         np.absolute(np.float32(threshed_navi_img) - 1) * warped_obst_img

    # 4) Convert thresholded image pixel values to rover-centric coords
    navi_xpix, navi_ypix = rover_coords(threshed_navi_img)
    obst_xpix, obst_ypix = rover_coords(threshed_obst_img)

    # 5) Convert rover-centric pixel values to world coords
    obst_x_world, obst_y_world = pix_to_world(
        obst_xpix, obst_ypix, xpos, ypos, yaw, world_size, scale)
    navi_x_world, navi_y_world = pix_to_world(
        navi_xpix, navi_ypix, xpos, ypos, yaw, world_size, scale)

    # 6) Update worldmap (to be displayed on right side of screen)
    # Example: data.worldmap[obst_y_world, obst_x_world, 0] += 1
    #          data.worldmap[rock_y_world, rock_x_world, 1] += 1
    #          data.worldmap[navi_y_world, navi_x_world, 2] += 1

    # navigatable terrain as BLUE
    data.worldmap[navi_y_world, navi_x_world, 2] = 255
    # obstacles as RED
    data.worldmap[obst_y_world, obst_x_world, 0] = 255
    # override blue on red
    data.worldmap[data.worldmap[:,:,2] > 0, 0] = 0

    # Fill by white if rocks in image
    rocks_img = find_rocks_by_hsv(warped_img)
    if rocks_img.any():
        rock_xpix, rock_ypix = rover_coords(rocks_img)
        rock_x_world, rock_y_world = pix_to_world(
            rock_xpix, rock_ypix, xpos, ypos, yaw, world_size, scale)
        # rocks as WHITE
        data.worldmap[rock_y_world, rock_x_world, :] = 255

    # 7) Make a mosaic image, below is some example code
    # First create a blank image (can be whatever shape you like)
    output_image = np.zeros((img.shape[0] + data.worldmap.shape[0], img.shape[1]*2, 3))

    # Next you can populate regions of the image with various output
    # Here I'm putting the original image in the upper left hand corner
    # 左上に元画像を貼り付ける
    output_image[0:img.shape[0], 0:img.shape[1]] = img

    # Let's create more images to add to the mosaic, first a warped image
    warped = perspect_transform(img, source, destination)
    # Add the warped image in the upper right hand corner
    # 右上に透視変換した画像を貼り付ける
    output_image[0:img.shape[0], img.shape[1]:] = warped

    # Overlay worldmap with ground truth map
    map_add = cv2.addWeighted(data.worldmap, 1, data.ground_truth, 0.5, 0)
    # Flip map overlay so y-axis points upward and add to output_image
    # 左下に、 実際の地面と 3 つの identify 情報が合成されたマップを貼り付ける
    output_image[img.shape[0]:, 0:data.worldmap.shape[1]] = np.flipud(map_add)

    # World map only
    # 右下に、 ワールドマップ だけのマップを貼り付ける
    output_image[img.shape[0]:,
                 img.shape[1]:img.shape[1] + data.worldmap.shape[1]] \
        = np.flipud(data.worldmap)

    # Then putting some text over the image
    cv2.putText(output_image,"Populate this image with your analyses to make a video!", (20, 20),
                cv2.FONT_HERSHEY_COMPLEX, 0.4, (255, 255, 255), 1)
    if data.count < len(data.images) - 1:
        data.count += 1 # Keep track of the index in the Databucket()

    return output_image

# ## Make a video from processed image data
# Use the [moviepy](https://zulko.github.io/moviepy/) library to process images and create a video.
#

# In[17]:

# Import everything needed to edit/save/watch video clips
from moviepy.editor import VideoFileClip
from moviepy.editor import ImageSequenceClip

# Define pathname to save the output video
output = '../output/test_mapping.mp4'
data = Databucket() # Re-initialize data in case you're running this cell multiple times
clip = ImageSequenceClip(data.images, fps=60) # Note: output video will be sped up because
                                          # recording rate in simulator is fps=25
new_clip = clip.fl_image(process_image) #NOTE: this function expects color images!!
get_ipython().run_line_magic('time', 'new_clip.write_videofile(output, audio=False)')

# ### This next cell should function as an inline video player
# If this fails to render the video, try running the following cell (alternative video rendering method).  You can also simply have a look at the saved mp4 in your `/output` folder

# In[392]:

from IPython.display import HTML
HTML("""
<video width="960" height="540" controls>
  <source src="{0}">
</video>
""".format(output))

# ### Below is an alternative way to create a video in case the above cell did not work.

# In[387]:

import io
import base64
video = io.open(output, 'r+b').read()
encoded_video = base64.b64encode(video)
HTML(data='''<video alt="test" controls>
                <source src="data:video/mp4;base64,{0}" type="video/mp4" />
             </video>'''.format(encoded_video.decode('ascii')))


import numpy as np # the Numpy library
import matplotlib.pyplot as plt #Matplotlib's pyplot

import sys #gives access to a C-like sys library
import os #gives access to operating system

print(sys.argv) #prints any command line arguments, incl program name
print(os.getcwd()) #prints the current working directory


"""
本脚本对cifar10数据进行解析，转换成图片，生成训练图片和测试图片。
"""

import urllib.request
import os
import sys
import tarfile
import glob
import pickle
import numpy as np
import cv2

# 通过这个函数完成对数据集的下载和解压
# tarball_url 表示cifar10数据集的下载链接
# dataset_dir 表示存储的路径

# 执行下面的代码可以完成数据集的下载和解压
# DATA_URL = 'http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
# DATA_DIR = 'data'

# download_and_uncompress_tarball(DATA_URL, DATA_DIR)
def download_and_uncompress_tarball(tarball_url, dataset_dir):
    """Downloads the `tarball_url` and uncompresses it locally.
  Args:
    tarball_url: The URL of a tarball file.
    dataset_dir: The directory where the temporary files are stored.
  """
    # tarball_url='http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
    filename = tarball_url.split('/')[-1]  # 文件名，通过/拆分字符串，取最后一节，也就是cifar-10-python.tar.gz
    # dataset_dir = 'data'
    # os.path.join()路径拼接，/data/cifar-10-python.tar.gz
    filepath = os.path.join(dataset_dir, filename)

    # 定义进度函数，分块下载
    # count 第几个块
    # block_size 每个块的大小
    # total_size 总的大小
    def _progress(count, block_size, total_size):
        # 等价于print() print底层调用的就是sys.stdout.write()
        sys.stdout.write('\r>> Downloading %s %.1f%%' % (
            filename, float(count * block_size) / float(total_size) * 100.0))
        sys.stdout.flush()

    # 如果文件还不存在，才下载，否则不下载
    if not os.path.isfile(filepath):
        # urlretrieve(url, filename=None, reporthook=None, data=None)
        # 参数url：下载链接地址
        # 参数filename：指定了保存本地路径（如果参数未指定，urllib会生成一个临时文件保存数据。）
        # 参数reporthook：是一个回调函数，当连接上服务器、以及相应的数据块传输完毕时会触发该回调，我们可以利用这个回调函数来显示当前的下载进度。
        # 参数data：指post导服务器的数据，该方法返回一个包含两个元素的(filename, headers)元组，filename表示保存到本地的路径，header表示服务器的响应头
        filepath, _ = urllib.request.urlretrieve(tarball_url, filepath, _progress)
        print()

    else:
        print('File already existed!')

    # 获取文件的信息
    statinfo = os.stat(filepath)
    print('Successfully downloaded', filename, statinfo.st_size, 'bytes.')
    # 将gz文件解压到dataset_dir指定的文件夹
    tarfile.open(filepath, 'r:gz').extractall(dataset_dir)

# 预定义10个分类
classification = ['airplane',
                  'automobile',
                  'bird',
                  'cat',
                  'deer',
                  'dog',
                  'frog',
                  'horse',
                  'ship',
                  'truck']

# pickle用于序列化
# 用于python特有的类型和python的数据类型间进行转换
# pickle提供四个功能：dumps,dump,loads,load
# pickle可以存储的数据类型
# - 所有python支持的原生类型：布尔值，整数，浮点数，复数，字符串，字节，None。
# - 由任何原生类型组成的列表，元组，字典和集合。
# - 函数，类，类的实例
def unpickle(file):
    with open(file, 'rb') as fo:  # 打开文件
        dict = pickle.load(fo, encoding='bytes')  # 是以字典的方式序列化数据
    return dict

# DATA_URL = 'http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
# DATA_DIR = 'data'

# download_and_uncompress_tarball(DATA_URL, DATA_DIR)

# 选择cifar10的路径，用相对路径
folders = 'data/cifar-10-batches-py'

# 获取目录下所有匹配到的训练集文件
trfiles = glob.glob(folders + "/data_batch*")

data = []  # 二进制的数据
labels = []  # 标签列表
for file in trfiles:  # 对于每个文件
    dt = unpickle(file)  # 得到反序列化后的数据 {"data": [byte], "labels": [int8]}
    data += list(dt[b"data"])  # 转化为list
    labels += list(dt[b"labels"])

# labels形如[1, 2, 3, 4, 6, ...]
# data 形如：... array([163, 173, 158, ..., 101, 100,  95], dtype=uint8) ...
print(labels)
print(len(data))

# 3通道32*32的图像数据
# [-1, 3, 32, 32]表示将data重新整理成3*32*32的图片，-1表示转化的数量根据实际情况确定
# numpy.reshape(a, newshape)
# a : 数组——需要处理的数据。
# 新的格式——整数或整数数组，如(2,3)表示2行3列。新的形状应该与原来的形状兼容，即行数和列数相乘后等于a中元素的数量。如果是整数，则结果将是长度的一维数组，所以这个整数必须等于a中元素数量。若这里是一个整数数组，那么其中一个数据可以为-1。在这种情况下，这个个值python会自动从根据第二个数值和剩余维度推断出来。
imgs = np.reshape(data, [-1, 3, 32, 32])

# 遍历所有图片
# imgs.shape[0]表示图片的数量
for i in range(imgs.shape[0]):
    im_data = imgs[i, ...]
    im_data = np.transpose(im_data, [1, 2, 0])  # 修改通道顺序，把通道移动到最后[32, 32, 3]
    im_data = cv2.cvtColor(im_data, cv2.COLOR_RGB2BGR)  # 修改图像通道为BGR

    # 把图像数据写成文件: data/image/test/airplane
    f = "{}/{}".format("data/image/train", classification[labels[i]])

    # 确保文件夹存在
    if not os.path.exists(f):
        os.mkdir(f)

    # 文件名形如：data/image/train/airplane/1.jpg
    cv2.imwrite("{}/{}.jpg".format(f, str(i)), im_data)


import numpy as np

from olpy.exceptions import NotFittedError

class LabelEncoder:
    """Encodes an output vector to match the specifications.

    Given that online learning algorithms usually work on output
    vectors with entries (-1, 1), this function performs this action
    for the user.

    Attributes:
        y (:obj:`array` of `ndarray`): the data to be transformed.
        positive_label (:obj:`int`, optional): The number in the output
            field that represents the positive label. The value passed
            should be different than -1. Defaults to 1.
        labels (:obj:`tuple`): represents the labels that are present
            in the dataset. This can be used at prediction time.
    """
    def __init__(self, positive_label=1):
        self.y = None
        self.positive_label = positive_label
        self.labels = None

    def fit(self, y):
        """Fits the output vector y.

        This method parses the parsed value and sets the necessary
        values to transform it later.

        Args:
            y (:obj:`list` or `numpy.ndarray`): the data to be transformed.

        Returns:
            self: the current instance.

        Raises:
            ValueError: if the number of labels is different than 2.
            AssertionError: if the positive label is not found in the
                labels.
        """
        self.y = y
        labels = np.unique(self.y)

        # First check that we have two values
        if len(labels) != 2:
            raise ValueError(
                'Expected two labels. Got {} instead'.format(len(labels))
                )

        # Let's check now that the specified positive label is in the array
        assert self.positive_label in labels,\
             'The positive label ({}) has not been found in the labels'

        self.labels = (
            set(labels).difference(
                {self.positive_label}
                ).pop(), self.positive_label)

        return self

    def transform(self, return_labels=True):
        """Transforms the data.

        Based on the information collected while fitting, this fuction
        returns the transformed labels that can be used directly for
        training.

        Args:
            return_labels (bool, optional): whether the labels should
                be returned or not. Default `True`.

        Returns:
            `numpy.ndarray` if return_labels is True else None

        Raises:
            NotFittedError if the encoder was not already fitted.
        """
        # All is okay. Now we can change
        if self.y is None:
            raise NotFittedError(
                    None,
                    None,
                    'Attempted to transform an unfitted encorder.'
                )
        if return_labels:
            return (
                np.array(
                    [1 if self.y[i] == self.positive_label
                    else -1 for i in range(self.y.shape[0])]),
                self.labels
            )
        else:
            return np.array(
                [1 if self.y[i] == self.positive_label
                else -1 for i in range(self.y.shape[0])])

    def fit_transform(self, y, return_labels=True):
        """Fits and transforms the data.

        Combines the actions of `fit` and `transform` methods.

        Args:
            return_labels (:obj:`bool`, optional): whether the labels should
                be returned or not. Default `True`.
            y (:obj:`array` of `ndarray`): the data to be transformed.

        Returns:
            `numpy.ndarray` if return_labels is True else None

        Raises:
            ValueError: if the number of labels is different than 2.
            AssertionError: if the positive label is not found in the
                    labels.
            NotFittedError if the encoder was not already fitted.

        """
        self.fit(y)
        return self.transform(return_labels=return_labels)


"""
.. _scratch:

Exploratory data analysis with NWB
==================================

This example will focus on the basics of working with an :py:class:`~pynwb.file.NWBFile` to do more than
storing standardized data for use and exchange. For example, you may want to store results from intermediate
analyses or one-off analyses with unknown utility. This functionality is primarily accomplished with linking
and scratch space.

.. note::
    The scratch space is explicitly for non-standardized data that is not intended for reuse
    by others. Standard NWB:N types, and extension if required, should always be used for any data that you
    intend to share. As such, published data should not include scratch data and a user should be able
    to ignore any data stored in scratch to use a file.

"""

####################
# Raw data
# --------
#
# To demonstrate linking and scratch space, lets assume we are starting with some acquired data.
#

# sphinx_gallery_thumbnail_path = 'figures/gallery_thumbnails_scratch.png'
from pynwb import NWBFile, TimeSeries, NWBHDF5IO
from datetime import datetime
from dateutil.tz import tzlocal
import numpy as np

# set up the NWBFile
start_time = datetime(2019, 4, 3, 11, tzinfo=tzlocal())
create_date = datetime(2019, 4, 15, 12, tzinfo=tzlocal())

nwb = NWBFile(session_description='demonstrate NWBFile scratch',  # required
              identifier='NWB456',  # required
              session_start_time=start_time,  # required
              file_create_date=create_date)  # optional

# make some fake data
timestamps = np.linspace(0, 100, 1024)
data = np.sin(0.333 * timestamps) + np.cos(0.1 * timestamps) + np.random.randn(len(timestamps))
test_ts = TimeSeries(name='raw_timeseries', data=data, unit='m', timestamps=timestamps)

# add it to the NWBFile
nwb.add_acquisition(test_ts)

with NWBHDF5IO('raw_data.nwb', mode='w') as io:
    io.write(nwb)

####################
# .. _basic_copying:
#
# Copying an NWB file
# -------------------
#
# To copy a file, we must first read the file.
#
raw_io = NWBHDF5IO('raw_data.nwb', 'r')
nwb_in = raw_io.read()

####################
# And then create a shallow copy the file with the :py:func:`~pynwb.file.NWBFile.copy` method
# of :py:class:`~pynwb.file.NWBFile` .

nwb_proc = nwb_in.copy()

####################
#
# Now that we have a copy, lets process some data, and add the results as a :py:class:`~pynwb.base.ProcessingModule`
# to our copy of the file. [#]_

import scipy.signal as sps

mod = nwb_proc.create_processing_module('filtering_module', "a module to store filtering results")

ts1 = nwb_in.acquisition['raw_timeseries']
filt_data = sps.correlate(ts1.data, np.ones(128), mode='same') / 128
ts2 = TimeSeries(name='filtered_timeseries', data=filt_data, unit='m', timestamps=ts1)

mod.add_container(ts2)

####################
#
# Now write the copy, which contains the processed data. [#]_

with NWBHDF5IO('processed_data.nwb', mode='w', manager=raw_io.manager) as io:
    io.write(nwb_proc)

####################
#
# .. [#]
#    .. note::
#       Notice here that we are reusing the timestamps to the original TimeSeries.
#
# .. [#]
#    .. note::
#       The ``processed_data.nwb`` file (i.e., our copy) stores our processing module and contains external
#       links to all data in our original file, i.e., the data from our raw file is being linked to, not copied.
#       This allows us to isolate our processing data in a separate file while still allowing us to access the
#       raw data from our ``processed_data.nwb`` file, without having to duplicate the data.
#

####################
# .. _basic_scratch:
#
# Adding scratch data
# -------------------
#
# You may end up wanting to store results from some one-off analysis, and writing an extension
# to get your data into an NWBFile is too much over head. This is facilitated by the scratch space
# in NWB:N. [#]_
#
# First, lets read our processed data and then make a copy

proc_io = NWBHDF5IO('processed_data.nwb', 'r')
nwb_proc_in = proc_io.read()

####################
#
# Now make a copy to put our scratch data into [#]_

nwb_scratch = nwb_proc_in.copy()

####################
#
# Now lets do an analysis for which we do not have a specification, but we would like to store
# the results for.

filt_ts = nwb_scratch.modules['filtering_module']['filtered_timeseries']

fft = np.fft.fft(filt_ts.data)

nwb_scratch.add_scratch(fft, name='dft_filtered', description='discrete Fourier transform from filtered data')

####################
#
# Finally, write the results.

with NWBHDF5IO('scratch_analysis.nwb', 'w', manager=proc_io.manager) as io:
    io.write(nwb_scratch)

####################
#
# To get your results back, you can index into :py:attr:`~pynwb.file.NWBFile.scratch` or use
# :py:func:`~pynwb.file.NWBFile.get_scratch`:

scratch_io = NWBHDF5IO('scratch_analysis.nwb', 'r')
nwb_scratch_in = scratch_io.read()

fft_in = nwb_scratch_in.scratch['dft_filtered']

fft_in = nwb_scratch_in.get_scratch('dft_filtered')

####################
#
# .. [#]
#    .. note::
#       This scratch space only exists if you add scratch data.
#
# .. [#]
#    .. note::
#       We recommend writing scratch data into copies of files only. This will make it easier to
#       isolate and discard scratch data and avoids updating files that store precious data.


import itertools
import copy as cp
import numpy as np
import dill

from pySDC.core.Controller import controller
from pySDC.core import Step as stepclass
from pySDC.core.Errors import ControllerError, CommunicationError

class controller_nonMPI(controller):
    """

    PFASST controller, running serialized version of PFASST in blocks (MG-style)

    """

    def __init__(self, num_procs, controller_params, description):
        """
        Initialization routine for PFASST controller

        Args:
           num_procs: number of parallel time steps (still serial, though), can be 1
           controller_params: parameter set for the controller and the steps
           description: all the parameters to set up the rest (levels, problems, transfer, ...)
        """

        if 'predict' in controller_params:
            raise ControllerError('predict flag is ignored, use predict_type instead')

        # call parent's initialization routine
        super(controller_nonMPI, self).__init__(controller_params)

        self.MS = [stepclass.step(description)]

        # try to initialize via dill.copy (much faster for many time-steps)
        try:
            for _ in range(num_procs - 1):
                self.MS.append(dill.copy(self.MS[0]))
        # if this fails (e.g. due to un-picklable data in the steps), initialize seperately
        except dill.PicklingError and TypeError:
            self.logger.warning('Need to initialize steps separately due to pickling error')
            for _ in range(num_procs - 1):
                self.MS.append(stepclass.step(description))

        if self.params.dump_setup:
            self.dump_setup(step=self.MS[0], controller_params=controller_params, description=description)

        if num_procs > 1 and len(self.MS[0].levels) > 1:
            for S in self.MS:
                for L in S.levels:
                    if not L.sweep.coll.right_is_node:
                        raise ControllerError("For PFASST to work, we assume uend^k = u_M^k")

        if all(len(S.levels) == len(self.MS[0].levels) for S in self.MS):
            self.nlevels = len(self.MS[0].levels)
        else:
            raise ControllerError('all steps need to have the same number of levels')

        if self.nlevels == 0:
            raise ControllerError('need at least one level')

        self.nsweeps = []
        for nl in range(self.nlevels):

            if all(S.levels[nl].params.nsweeps == self.MS[0].levels[nl].params.nsweeps for S in self.MS):
                self.nsweeps.append(self.MS[0].levels[nl].params.nsweeps)

        if self.nlevels > 1 and self.nsweeps[-1] > 1:
            raise ControllerError('this controller cannot do multiple sweeps on coarsest level')

        if self.nlevels == 1 and self.params.predict_type is not None:
            self.logger.warning('you have specified a predictor type but only a single level.. '
                                'predictor will be ignored')

    def check_iteration_estimator(self, MS):
        """
        Method to check the iteration estimator

        Args:
            MS (list): list of currently active steps
        """
        diff_new = 0.0
        Kest_loc = 99

        # go through active steps and compute difference, Ltilde, Kest up to this step
        for S in MS:
            L = S.levels[0]

            for m in range(1, L.sweep.coll.num_nodes + 1):
                diff_new = max(diff_new, abs(L.uold[m] - L.u[m]))

            if S.status.iter == 1:
                S.status.diff_old_loc = diff_new
                S.status.diff_first_loc = diff_new
            elif S.status.iter > 1:
                Ltilde_loc = min(diff_new / S.status.diff_old_loc, 0.9)
                S.status.diff_old_loc = diff_new
                alpha = 1 / (1 - Ltilde_loc) * S.status.diff_first_loc
                Kest_loc = np.log(S.params.errtol / alpha) / np.log(Ltilde_loc) * 1.05  # Safety factor!
                self.logger.debug(f'LOCAL: {L.time:8.4f}, {S.status.iter}: {int(np.ceil(Kest_loc))}, '
                                  f'{Ltilde_loc:8.6e}, {Kest_loc:8.6e}, {Ltilde_loc ** S.status.iter * alpha:8.6e}')
                # You should not stop prematurely on earlier steps, since later steps may need more accuracy to reach
                # the tolerance themselves. The final Kest_loc is the one that counts.
                # if np.ceil(Kest_loc) <= S.status.iter:
                #     S.status.force_done = True

        # set global Kest as last local one, force stop if done
        for S in MS:
            if S.status.iter > 1:
                Kest_glob = Kest_loc
                if np.ceil(Kest_glob) <= S.status.iter:
                    S.status.force_done = True

    def run(self, u0, t0, Tend):
        """
        Main driver for running the serial version of SDC, MSSDC, MLSDC and PFASST (virtual parallelism)

        Args:
           u0: initial values
           t0: starting time
           Tend: ending time

        Returns:
            end values on the finest level
            stats object containing statistics for each step, each level and each iteration
        """

        # some initializations and reset of statistics
        uend = None
        num_procs = len(self.MS)
        self.hooks.reset_stats()

        # initial ordering of the steps: 0,1,...,Np-1
        slots = list(range(num_procs))

        # initialize time variables of each step
        time = [t0 + sum(self.MS[j].dt for j in range(p)) for p in slots]

        # determine which steps are still active (time < Tend)
        active = [time[p] < Tend - 10 * np.finfo(float).eps for p in slots]

        if not any(active):
            raise ControllerError('Nothing to do, check t0, dt and Tend.')

        # compress slots according to active steps, i.e. remove all steps which have times above Tend
        active_slots = list(itertools.compress(slots, active))

        # initialize block of steps with u0
        self.restart_block(active_slots, time, u0)

        self.hooks.post_setup(step=None, level_number=None)

        # call pre-run hook
        for S in self.MS:
            self.hooks.pre_run(step=S, level_number=0)

        # main loop: as long as at least one step is still active (time < Tend), do something
        while any(active):

            MS_active = [self.MS[p] for p in active_slots]
            done = False
            while not done:
                done = self.pfasst(MS_active)

            # uend is uend of the last active step in the list
            uend = self.MS[active_slots[-1]].levels[0].uend

            for p in active_slots:
                time[p] += num_procs * self.MS[p].dt

            # determine new set of active steps and compress slots accordingly
            active = [time[p] < Tend - 10 * np.finfo(float).eps for p in slots]
            active_slots = list(itertools.compress(slots, active))

            # restart active steps (reset all values and pass uend to u0)
            self.restart_block(active_slots, time, uend)

        # call post-run hook
        for S in self.MS:
            self.hooks.post_run(step=S, level_number=0)

        return uend, self.hooks.return_stats()

    def restart_block(self, active_slots, time, u0):
        """
        Helper routine to reset/restart block of (active) steps

        Args:
            active_slots: list of active steps
            time: list of new times
            u0: initial value to distribute across the steps

        """

        # loop over active slots (not directly, since we need the previous entry as well)
        for j in range(len(active_slots)):

            # get slot number
            p = active_slots[j]

            # store current slot number for diagnostics
            self.MS[p].status.slot = p
            # store link to previous step
            self.MS[p].prev = self.MS[active_slots[j - 1]]
            # resets step
            self.MS[p].reset_step()
            # determine whether I am the first and/or last in line
            self.MS[p].status.first = active_slots.index(p) == 0
            self.MS[p].status.last = active_slots.index(p) == len(active_slots) - 1
            # initialize step with u0
            self.MS[p].init_step(u0)
            # reset some values
            self.MS[p].status.done = False
            self.MS[p].status.prev_done = False
            self.MS[p].status.iter = 0
            self.MS[p].status.stage = 'SPREAD'
            self.MS[p].status.force_done = False
            self.MS[p].status.time_size = len(active_slots)

            for l in self.MS[p].levels:
                l.tag = None
                l.status.sweep = 1

        for p in active_slots:
            for lvl in self.MS[p].levels:
                lvl.status.time = time[p]

    @staticmethod
    def recv(target, source, tag=None):
        """
        Receive function

        Args:
            target: level which will receive the values
            source: level which initiated the send
            tag: identifier to check if this message is really for me
        """

        if tag is not None and source.tag != tag:
            raise CommunicationError('source and target tag are not the same, got %s and %s' % (source.tag, tag))
        # simply do a deepcopy of the values uend to become the new u0 at the target
        target.u[0] = target.prob.dtype_u(source.uend)
        # re-evaluate f on left interval boundary
        target.f[0] = target.prob.eval_f(target.u[0], target.time)

    @staticmethod
    def send(source, tag):
        """
        Send function

        Args:
            source: level which has the new values
            tag: identifier for this message
        """
        # sending here means computing uend ("one-sided communication")
        source.sweep.compute_end_point()
        source.tag = cp.deepcopy(tag)

    def pfasst(self, local_MS_active):
        """
        Main function including the stages of SDC, MLSDC and PFASST (the "controller")

        For the workflow of this controller, check out one of our PFASST talks or the pySDC paper

        This method changes self.MS directly by accessing active steps through local_MS_active. Nothing is returned.

        Args:
            local_MS_active (list): all active steps
        """

        def spread(local_MS_running):
            """
            Spreading phase

            Args:
                local_MS_running (list): list of currently running steps
            """

            for S in local_MS_running:

                # first stage: spread values
                self.hooks.pre_step(step=S, level_number=0)

                # call predictor from sweeper
                S.levels[0].sweep.predict()

                if self.params.use_iteration_estimator:
                    # store pervious iterate to compute difference later on
                    S.levels[0].uold[:] = S.levels[0].u[:]

                # update stage
                if len(S.levels) > 1:  # MLSDC or PFASST with predict
                    S.status.stage = 'PREDICT'
                else:
                    S.status.stage = 'IT_CHECK'

        def predict(local_MS_running):
            """
            Predictor phase

            Args:
                local_MS_running (list): list of currently running steps
            """

            for S in local_MS_running:
                self.hooks.pre_predict(step=S, level_number=0)

            if self.params.predict_type is None:
                pass

            elif self.params.predict_type == 'fine_only':

                # do a fine sweep only
                for S in local_MS_running:
                    S.levels[0].sweep.update_nodes()

            elif self.params.predict_type == 'libpfasst_style':

                # loop over all steps
                for S in local_MS_running:

                    # restrict to coarsest level
                    for l in range(1, len(S.levels)):
                        S.transfer(source=S.levels[l - 1], target=S.levels[l])

                # run in serial on coarse level
                for S in local_MS_running:

                    self.hooks.pre_comm(step=S, level_number=len(S.levels) - 1)
                    # receive from previous step (if not first)
                    if not S.status.first:
                        self.logger.debug('Process %2i receives from %2i on level %2i with tag %s -- PREDICT' %
                                          (S.status.slot, S.prev.status.slot, len(S.levels) - 1, 0))
                        self.recv(S.levels[-1], S.prev.levels[-1], tag=(len(S.levels), 0, S.prev.status.slot))
                    self.hooks.post_comm(step=S, level_number=len(S.levels) - 1)

                    # do the coarse sweep
                    S.levels[-1].sweep.update_nodes()

                    self.hooks.pre_comm(step=S, level_number=len(S.levels) - 1)
                    # send to succ step
                    if not S.status.last:
                        self.logger.debug('Process %2i provides data on level %2i with tag %s -- PREDICT'
                                          % (S.status.slot, len(S.levels) - 1, 0))
                        self.send(S.levels[-1], tag=(len(S.levels), 0, S.status.slot))
                    self.hooks.post_comm(step=S, level_number=len(S.levels) - 1, add_to_stats=True)

                # go back to fine level, sweeping
                for l in range(self.nlevels - 1, 0, -1):

                    for S in local_MS_running:
                        # prolong values
                        S.transfer(source=S.levels[l], target=S.levels[l - 1])

                        if l - 1 > 0:
                            S.levels[l - 1].sweep.update_nodes()

                # end with a fine sweep
                for S in local_MS_running:
                    S.levels[0].sweep.update_nodes()

            elif self.params.predict_type == 'pfasst_burnin':

                # loop over all steps
                for S in local_MS_running:

                    # restrict to coarsest level
                    for l in range(1, len(S.levels)):
                        S.transfer(source=S.levels[l - 1], target=S.levels[l])

                # loop over all steps
                for q in range(len(local_MS_running)):

                    # loop over last steps: [1,2,3,4], [2,3,4], [3,4], [4]
                    for p in range(q, len(local_MS_running)):
                        S = local_MS_running[p]

                        # do the sweep with new values
                        S.levels[-1].sweep.update_nodes()

                        self.hooks.pre_comm(step=S, level_number=len(S.levels) - 1)
                        # send updated values on coarsest level
                        self.logger.debug('Process %2i provides data on level %2i with tag %s -- PREDICT'
                                          % (S.status.slot, len(S.levels) - 1, 0))
                        self.send(S.levels[-1], tag=(len(S.levels), 0, S.status.slot))
                        self.hooks.post_comm(step=S, level_number=len(S.levels) - 1)

                    # loop over last steps: [2,3,4], [3,4], [4]
                    for p in range(q + 1, len(local_MS_running)):
                        S = local_MS_running[p]
                        # receive values sent during previous sweep
                        self.hooks.pre_comm(step=S, level_number=len(S.levels) - 1)
                        self.logger.debug('Process %2i receives from %2i on level %2i with tag %s -- PREDICT' %
                                          (S.status.slot, S.prev.status.slot, len(S.levels) - 1, 0))
                        self.recv(S.levels[-1], S.prev.levels[-1], tag=(len(S.levels), 0, S.prev.status.slot))
                        self.hooks.post_comm(step=S, level_number=len(S.levels) - 1,
                                             add_to_stats=(p == len(local_MS_running) - 1))

                # loop over all steps
                for S in local_MS_running:

                    # interpolate back to finest level
                    for l in range(len(S.levels) - 1, 0, -1):
                        S.transfer(source=S.levels[l], target=S.levels[l - 1])

                # end this with a fine sweep
                for S in local_MS_running:
                    S.levels[0].sweep.update_nodes()

            elif self.params.predict_type == 'fmg':
                # TODO: implement FMG predictor
                raise NotImplementedError('FMG predictor is not yet implemented')

            else:
                raise ControllerError('Wrong predictor type, got %s' % self.params.predict_type)

            for S in local_MS_running:
                self.hooks.post_predict(step=S, level_number=0)

            for S in local_MS_running:
                # update stage
                S.status.stage = 'IT_CHECK'

        def it_check(local_MS_running):
            """
            Key routine to check for convergence/termination

            Args:
                local_MS_running (list): list of currently running steps
            """

            for S in local_MS_running:

                # send updated values forward
                self.hooks.pre_comm(step=S, level_number=0)
                if not S.status.last:
                    self.logger.debug('Process %2i provides data on level %2i with tag %s'
                                      % (S.status.slot, 0, S.status.iter))
                    self.send(S.levels[0], tag=(0, S.status.iter, S.status.slot))

                # receive values
                if not S.status.prev_done and not S.status.first:
                    self.logger.debug('Process %2i receives from %2i on level %2i with tag %s' %
                                      (S.status.slot, S.prev.status.slot, 0, S.status.iter))
                    self.recv(S.levels[0], S.prev.levels[0], tag=(0, S.status.iter, S.prev.status.slot))
                self.hooks.post_comm(step=S, level_number=0)

                S.levels[0].sweep.compute_residual()

            if self.params.use_iteration_estimator:
                self.check_iteration_estimator(local_MS_running)

            for S in local_MS_running:

                S.status.done = self.check_convergence(S)

                if S.status.iter > 0:
                    self.hooks.post_iteration(step=S, level_number=0)

            for S in local_MS_running:
                if not S.status.first:
                    self.hooks.pre_comm(step=S, level_number=0)
                    S.status.prev_done = S.prev.status.done  # "communicate"
                    self.hooks.post_comm(step=S, level_number=0, add_to_stats=True)
                    S.status.done = S.status.done and S.status.prev_done

                if self.params.all_to_done:
                    self.hooks.pre_comm(step=S, level_number=0)
                    S.status.done = all([T.status.done for T in local_MS_running])
                    self.hooks.post_comm(step=S, level_number=0, add_to_stats=True)

                if not S.status.done:
                    # increment iteration count here (and only here)
                    S.status.iter += 1
                    self.hooks.pre_iteration(step=S, level_number=0)

                    if self.params.use_iteration_estimator:
                        # store pervious iterate to compute difference later on
                        S.levels[0].uold[:] = S.levels[0].u[:]

                    if len(S.levels) > 1:  # MLSDC or PFASST
                        S.status.stage = 'IT_DOWN'
                    else:  # SDC or MSSDC
                        if len(local_MS_running) == 1 or self.params.mssdc_jac:  # SDC or parallel MSSDC (Jacobi-like)
                            S.status.stage = 'IT_FINE'
                        else:
                            S.status.stage = 'IT_COARSE'  # serial MSSDC (Gauss-like)
                else:
                    S.levels[0].sweep.compute_end_point()
                    self.hooks.post_step(step=S, level_number=0)
                    S.status.stage = 'DONE'

        def it_fine(local_MS_running):
            """
            Fine sweeps

            Args:
                local_MS_running (list): list of currently running steps
            """

            for S in local_MS_running:
                S.levels[0].status.sweep = 0

            for k in range(self.nsweeps[0]):

                for S in local_MS_running:
                    S.levels[0].status.sweep += 1

                for S in local_MS_running:
                    # send updated values forward
                    self.hooks.pre_comm(step=S, level_number=0)
                    if not S.status.last:
                        self.logger.debug('Process %2i provides data on level %2i with tag %s'
                                          % (S.status.slot, 0, S.status.iter))
                        self.send(S.levels[0], tag=(0, S.status.iter, S.status.slot))

                    # # receive values
                    if not S.status.prev_done and not S.status.first:
                        self.logger.debug('Process %2i receives from %2i on level %2i with tag %s' %
                                          (S.status.slot, S.prev.status.slot, 0, S.status.iter))
                        self.recv(S.levels[0], S.prev.levels[0], tag=(0, S.status.iter, S.prev.status.slot))
                    self.hooks.post_comm(step=S, level_number=0, add_to_stats=(k == self.nsweeps[0] - 1))

                for S in local_MS_running:
                    # standard sweep workflow: update nodes, compute residual, log progress
                    self.hooks.pre_sweep(step=S, level_number=0)
                    S.levels[0].sweep.update_nodes()
                    S.levels[0].sweep.compute_residual()
                    self.hooks.post_sweep(step=S, level_number=0)

            for S in local_MS_running:
                # update stage
                S.status.stage = 'IT_CHECK'

        def it_down(local_MS_running):
            """
            Go down the hierarchy from finest to coarsest level

            Args:
                local_MS_running (list): list of currently running steps
            """

            for S in local_MS_running:
                S.transfer(source=S.levels[0], target=S.levels[1])

            for l in range(1, self.nlevels - 1):

                # sweep on middle levels (not on finest, not on coarsest, though)

                for _ in range(self.nsweeps[l]):

                    for S in local_MS_running:

                        # send updated values forward
                        self.hooks.pre_comm(step=S, level_number=l)
                        if not S.status.last:
                            self.logger.debug('Process %2i provides data on level %2i with tag %s'
                                              % (S.status.slot, l, S.status.iter))
                            self.send(S.levels[l], tag=(l, S.status.iter, S.status.slot))

                        # # receive values
                        if not S.status.prev_done and not S.status.first:
                            self.logger.debug('Process %2i receives from %2i on level %2i with tag %s' %
                                              (S.status.slot, S.prev.status.slot, l, S.status.iter))
                            self.recv(S.levels[l], S.prev.levels[l], tag=(l, S.status.iter, S.prev.status.slot))
                        self.hooks.post_comm(step=S, level_number=l)

                    for S in local_MS_running:
                        self.hooks.pre_sweep(step=S, level_number=l)
                        S.levels[l].sweep.update_nodes()
                        S.levels[l].sweep.compute_residual()
                        self.hooks.post_sweep(step=S, level_number=l)

                for S in local_MS_running:
                    # transfer further down the hierarchy
                    S.transfer(source=S.levels[l], target=S.levels[l + 1])

            for S in local_MS_running:
                # update stage
                S.status.stage = 'IT_COARSE'

        def it_coarse(local_MS_running):
            """
            Coarse sweep

            Args:
                local_MS_running (list): list of currently running steps
            """

            for S in local_MS_running:

                # receive from previous step (if not first)
                self.hooks.pre_comm(step=S, level_number=len(S.levels) - 1)
                if not S.status.first and not S.status.prev_done:
                    self.logger.debug('Process %2i receives from %2i on level %2i with tag %s' %
                                      (S.status.slot, S.prev.status.slot, len(S.levels) - 1, S.status.iter))
                    self.recv(S.levels[-1], S.prev.levels[-1], tag=(len(S.levels), S.status.iter, S.prev.status.slot))
                self.hooks.post_comm(step=S, level_number=len(S.levels) - 1)

                # do the sweep
                self.hooks.pre_sweep(step=S, level_number=len(S.levels) - 1)
                S.levels[-1].sweep.update_nodes()
                S.levels[-1].sweep.compute_residual()
                self.hooks.post_sweep(step=S, level_number=len(S.levels) - 1)

                # send to succ step
                self.hooks.pre_comm(step=S, level_number=len(S.levels) - 1)
                if not S.status.last:
                    self.logger.debug('Process %2i provides data on level %2i with tag %s'
                                      % (S.status.slot, len(S.levels) - 1, S.status.iter))
                    self.send(S.levels[-1], tag=(len(S.levels), S.status.iter, S.status.slot))
                self.hooks.post_comm(step=S, level_number=len(S.levels) - 1, add_to_stats=True)

                # update stage
                if len(S.levels) > 1:  # MLSDC or PFASST
                    S.status.stage = 'IT_UP'
                else:  # MSSDC
                    S.status.stage = 'IT_CHECK'

        def it_up(local_MS_running):
            """
            Prolong corrections up to finest level (parallel)

            Args:
                local_MS_running (list): list of currently running steps
            """

            for l in range(self.nlevels - 1, 0, -1):

                for S in local_MS_running:
                    # prolong values
                    S.transfer(source=S.levels[l], target=S.levels[l - 1])

                # on middle levels: do communication and sweep as usual
                if l - 1 > 0:

                    for k in range(self.nsweeps[l - 1]):

                        for S in local_MS_running:

                            # send updated values forward
                            self.hooks.pre_comm(step=S, level_number=l - 1)
                            if not S.status.last:
                                self.logger.debug('Process %2i provides data on level %2i with tag %s'
                                                  % (S.status.slot, l - 1, S.status.iter))
                                self.send(S.levels[l - 1], tag=(l - 1, S.status.iter, S.status.slot))

                            # # receive values
                            if not S.status.prev_done and not S.status.first:
                                self.logger.debug('Process %2i receives from %2i on level %2i with tag %s' %
                                                  (S.status.slot, S.prev.status.slot, l - 1, S.status.iter))
                                self.recv(S.levels[l - 1], S.prev.levels[l - 1], tag=(l - 1, S.status.iter,
                                                                                S.prev.status.slot))
                            self.hooks.post_comm(step=S, level_number=l - 1,
                                                 add_to_stats=(k == self.nsweeps[l - 1] - 1))

                        for S in local_MS_running:
                            self.hooks.pre_sweep(step=S, level_number=l - 1)
                            S.levels[l - 1].sweep.update_nodes()
                            S.levels[l - 1].sweep.compute_residual()
                            self.hooks.post_sweep(step=S, level_number=l - 1)

            for S in local_MS_running:
                # update stage
                S.status.stage = 'IT_FINE'

        def default(local_MS_running):
            """
            Default routine to catch wrong status

            Args:
                local_MS_running (list): list of currently running steps
            """
            raise ControllerError('Unknown stage, got %s' % local_MS_running[0].status.stage)  # TODO

        # if all stages are the same (or DONE), continue, otherwise abort
        stages = [S.status.stage for S in local_MS_active if S.status.stage != 'DONE']
        if stages[1:] == stages[:-1]:
            stage = stages[0]
        else:
            raise ControllerError('not all stages are equal')

        self.logger.debug(stage)

        MS_running = [S for S in local_MS_active if S.status.stage != 'DONE']

        switcher = {
            'SPREAD': spread,
            'PREDICT': predict,
            'IT_CHECK': it_check,
            'IT_FINE': it_fine,
            'IT_DOWN': it_down,
            'IT_COARSE': it_coarse,
            'IT_UP': it_up
        }

        switcher.get(stage, default)(MS_running)

        return all([S.status.done for S in local_MS_active])


def monitor(f, input='print', output='print'):
    """
    Returns a wrapped copy of *f* that monitors evaluation by calling
    *input* with every input (*args*, *kwargs*) passed to *f* and
    *output* with every value returned from *f*. The default action
    (specify using the special string value ``'print'``) is to print
    inputs and outputs to stdout, along with the total evaluation
    count::

        >>> from mpmath import *
        >>> mp.dps = 5; mp.pretty = False
        >>> diff(monitor(exp), 1)   # diff will eval f(x-h) and f(x+h)
        in  0 (mpf('0.99999999906867742538452148'),) {}
        out 0 mpf('2.7182818259274480055282064')
        in  1 (mpf('1.0000000009313225746154785'),) {}
        out 1 mpf('2.7182818309906424675501024')
        mpf('2.7182808')

    To disable either the input or the output handler, you may
    pass *None* as argument.

    Custom input and output handlers may be used e.g. to store
    results for later analysis::

        >>> mp.dps = 15
        >>> input = []
        >>> output = []
        >>> findroot(monitor(sin, input.append, output.append), 3.0)
        mpf('3.1415926535897932')
        >>> len(input)  # Count number of evaluations
        9
        >>> print input[3], output[3]
        ((mpf('3.1415076583334066'),), {}) 8.49952562843408e-5
        >>> print input[4], output[4]
        ((mpf('3.1415928201669122'),), {}) -1.66577118985331e-7

    """
    if not input:
        input = lambda v: None
    elif input == 'print':
        incount = [0]
        def input(value):
            args, kwargs = value
            print "in  %s %r %r" % (incount[0], args, kwargs)
            incount[0] += 1
    if not output:
        output = lambda v: None
    elif output == 'print':
        outcount = [0]
        def output(value):
            print "out %s %r" % (outcount[0], value)
            outcount[0] += 1
    def f_monitored(*args, **kwargs):
        input((args, kwargs))
        v = f(*args, **kwargs)
        output(v)
        return v
    return f_monitored

def timing(f, *args, **kwargs):
    """
    Returns time elapsed for evaluating ``f()``. Optionally arguments
    may be passed to time the execution of ``f(*args, **kwargs)``.

    If the first call is very quick, ``f`` is called
    repeatedly and the best time is returned.
    """
    once = kwargs.get('once')
    if 'once' in kwargs:
        del kwargs['once']
    if args or kwargs:
        if len(args) == 1 and not kwargs:
            arg = args[0]
            g = lambda: f(arg)
        else:
            g = lambda: f(*args, **kwargs)
    else:
        g = f
    from timeit import default_timer as clock
    t1=clock(); v=g(); t2=clock(); t=t2-t1
    if t > 0.05 or once:
        return t
    for i in range(3):
        t1=clock();
        # Evaluate multiple times because the timer function
        # has a significant overhead
        g();g();g();g();g();g();g();g();g();g()
        t2=clock()
        t=min(t,(t2-t1)/10)
    return t


import pathlib

import numpy_financial as npf
import pandas as pd

import streamlit as st

def main():
    instructions_path = str(pathlib.Path(__file__).parent.resolve() / "instructions.md")
    readme_text = st.markdown(get_file_content_as_string(instructions_path))

    # Once we have the dependencies, add a selector for the app mode on the sidebar.
    st.sidebar.title("What to do")
    app_mode = st.sidebar.selectbox(
        "Choose a chapter",
        ["Home", "Compounding machine", "Finance starter kit", "Investing"],
    )
    if app_mode == "Home":
        st.sidebar.success("To continue select a chapter from the toolbar.")
    elif app_mode == "Compounding machine":
        readme_text.empty()
        compounding_machine()
    elif app_mode == "Finance starter kit":
        readme_text.empty()
        finance_starter_kit()
    elif app_mode == "Investing":
        readme_text.empty()
        investing()

def compounding_machine():

    st.write('''
    # Compounding: the eighth wonder of the world

    What is compounding? Compounding is exponential increase in value due to earning interest on both the principal and the accumulated interest.

    Let's break it down quickly with an example. Let's say that you lent me $1000 for three years, and I promised to pay you 10% for each year for this favor.

    * The $1000 you loaned me is called the **principal**
    * The 10% I pay you each year is the **interest** rate on the loan

    With this arrangement, there are two ways that I might pay you back: simple interest and compound interest. In simple interest, I would simply pay you $100 (10% of $1000) at the end of each year, and return the $1000 at the end of the third year. Compound interest is more interesting. If we were to use compound interest, I would pay you 10% of the principal in the first year, then 10% of the principal *plus* the first year's interest, then 10% of the principal *plus* the first and second year's interest (and then I would also return the $1000). If you were calculating, that would be $100 in the first year (10% of $1000), $110 in the second year (10% of $1000+$100), and then $121 in the third year (10% of $1000+$100+$110).

    As a result, you would end up with $1300 with simple interest and $1331 with compound interest. It might not seem like a big difference, but if you were to change it to a 20 year loan with the same initial principal of $1000 and 10% interest, you would end up with $3000 with simple interest and $6727.50 with compound interest. As the time horizon gets longer and longer, compounding's exponential growth becomes extremely obvious.

    "Compounding is the eighth wonder of the world." This quote is often apocryphally attributed to Albert Einstein, but the misattribution should not detract from its simplicity and profoundness.

    ''')

    st.write(
    """
    ## Try it out

    Since compounding is an exponential process, it's difficult to get a handle on it unless you try some simulations. I've designed a simple tool here that will let you play with the parameters of a compounding experiment:
    * Starting principal: The amount of money you start with.
    * Years: The number of years you want to run the compounding experiment for.
    * Annual contribution: The amount you want to add to the compounding machine each year.
    * Interest rate: The amount of interest you get paid each year.

    Try extending the number of years we compound for. What happens when we contribute a small amount of money each year? And what happens to the final value when we change the interest rate around?

    """
    )

    with st.form(key="columns_in_form"):
        col1, col2 = st.beta_columns(2)
        with col1:
            principal_input = st.number_input(
                "Starting principal", min_value=0, value=1000, key="principal"
            )
            add_sum_input = st.number_input(
                "Annual contribution", min_value=0, value=0, key="add_sum"
            )
        with col2:
            years_input = st.number_input(
                "Years", min_value=0, max_value=100, value=10, key="years"
            )
            interest_rate_input = st.number_input(
                "Interest rate (%)", min_value=0, value=8, key="interest_rate"
            )
        submitted = st.form_submit_button("Run")

    future_value = npf.fv(
        rate=interest_rate_input / 100, nper=years_input, pmt=-add_sum_input, pv=-principal_input, when="begin"
    )

    st.write("You'll have this left: $", "{:,}".format(round(future_value, 2)))

    @st.cache
    def graph_compounding(years, principal, add_sum, interest_rate):
        # graphing
        year_list = []
        value = []

        for year in range(1, int(years)+1):
            principal = (principal+add_sum)*(1+interest_rate/100)
            year_list.append(year)
            value.append(principal)

        df_compounding = pd.DataFrame(zip(year_list, value), columns=['years', 'value'])
        df_compounding.index = df_compounding['years'].values
        del df_compounding['years']
        return df_compounding.round(2)

    df_compounding = graph_compounding(years_input, principal_input, add_sum_input, interest_rate_input)

    st.write("### Growth over time")

    st.bar_chart(df_compounding)

    st.write(
        '''
        It's not obvious how much compounding affects growth until we have an actual tangible demonstration like this. I understood it from a philosophical level before, but I never truly intuited and viscerally understood the power of compounding until I ran these calculations myself.

        ## Chess and compounding

        You might have heard this story before, but there is a famous legend about the origin of chess. Supposedly, when the inventor of the game showed it to the emperor of India, the emperor was so impressed by the new game that he promised any reward to the inventor.

        The inventor replied, "I only wish for this. Give me one grain of rice for the first square of the chessboard, two grains for the next square, four for the next, eight for the next, and so on for all 64 squares, with each square having double the number of grains as the square before."

        The emperor agreed, amazed that the man had asked for such a small reward — or so he thought. After a week, his treasurer came back and informed him that the reward would add up to an astronomical sum, far greater than all the rice that could conceivably be produced in many centuries!

        While this example of compounding is negative (for the king, at least), the very nature of compounding is neither good nor bad — it is simply a mathematical truth that exists. Compounding can either make or break you, but our job is to make compounding work for you in all facets of life. Harnessing compounding requires you to save or work a little bit harder today to reap a much greater reward tomorrow. Compounding is implementing delayed gratification and requires you to take a long-term view of life.

        **In sum: use your youth and long time horizon — save (but not so much that you can't enjoy the present!) and plant the seeds for your own compounding machine.** We want to look back and thank our previous selves for doing this and setting up such a practice.

        As they say, the best time to plant a tree was twenty years ago; the next best time is today.

        '''
    )

def finance_starter_kit():
    st.write('''
    # Finance starter kit

    Your financial strength is a function of four simple variables:

    1. **Income**: the money that you make
    2. **Expenses**: the money that you spend
    3. **Savings**: the pool of money you have saved
    4. **Investments**: money available for your compounding machine

    However, the first step is to understand exactly what is happening in each of these domains. Figure out how much income and expenses you have each month, and do a tally of the amount of money in your savings and investment accounts. The oft-repeated quote — What is measured is managed — is paramount here.

    Once you figure out what is going in or out of each of these buckets, I'd encourage you to think about how you might slowly begin to improve each facet. Here are some heuristics and good rules of thumb:

    ## Income

    - mostly comes from your job, and also from alternative sources (freelancing)
    - often neglected, but think about ways you can improve your income (negotiation of your salary, etc.)
    - ensure that your monthly net income (income minus expenses) is positive

    ## Expenses

    - figure out what you enjoy, and feel comfortable spending in that category (to within reason)
        - cut your expenses on everything else
    - focus your time on saving expenses on big ticket items (e.g. car payments, mortgages, subscription services, student loans)
    - other expenses: credit card debt, student loan debt, taxes
    - figure out what your average monthly expense is

    ## Savings

    - rule of thumb: your savings account should contain 3-6 months worth of expenses (computed)
        - do not put any money into investing until you have at least 3 months of savings

    ## Investments

    This is one of the most neglected parts of personal finance, yet it's arguably the most important (because this is where we get our compounding machine)!

    Go to the next section to learn more!

    ''')

def investing():

    st.write('''
    # Investing

    **Beep boop bop 🤖 This page is currently under construction.**
    ''')

def get_file_content_as_string(path):
    with open(path, "r") as file:
        data = file.read()
    return data

if __name__ == "__main__":
    main()


"""
Tests that the file header is properly handled or inferred
during parsing for all of the parsers defined in parsers.py
"""

from collections import namedtuple
from io import StringIO

import numpy as np
import pytest

from pandas.errors import ParserError

from pandas import (
    DataFrame,
    Index,
    MultiIndex,
)
import pandas._testing as tm

# TODO(1.4): Change me to xfails at release time
skip_pyarrow = pytest.mark.usefixtures("pyarrow_skip")

@skip_pyarrow
def test_read_with_bad_header(all_parsers):
    parser = all_parsers
    msg = r"but only \d+ lines in file"

    with pytest.raises(ValueError, match=msg):
        s = StringIO(",,")
        parser.read_csv(s, header=[10])

def test_negative_header(all_parsers):
    # see gh-27779
    parser = all_parsers
    data = """1,2,3,4,5
6,7,8,9,10
11,12,13,14,15
"""
    with pytest.raises(
        ValueError,
        match="Passing negative integer to header is invalid. "
        "For no header, use header=None instead",
    ):
        parser.read_csv(StringIO(data), header=-1)

@pytest.mark.parametrize("header", [([-1, 2, 4]), ([-5, 0])])
def test_negative_multi_index_header(all_parsers, header):
    # see gh-27779
    parser = all_parsers
    data = """1,2,3,4,5
        6,7,8,9,10
        11,12,13,14,15
        """
    with pytest.raises(
        ValueError, match="cannot specify multi-index header with negative integers"
    ):
        parser.read_csv(StringIO(data), header=header)

@pytest.mark.parametrize("header", [True, False])
def test_bool_header_arg(all_parsers, header):
    # see gh-6114
    parser = all_parsers
    data = """\
MyColumn
a
b
a
b"""
    msg = "Passing a bool to header is invalid"
    with pytest.raises(TypeError, match=msg):
        parser.read_csv(StringIO(data), header=header)

def test_no_header_prefix(all_parsers):
    parser = all_parsers
    data = """1,2,3,4,5
6,7,8,9,10
11,12,13,14,15
"""
    result = parser.read_csv(StringIO(data), prefix="Field", header=None)
    expected = DataFrame(
        [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]],
        columns=["Field0", "Field1", "Field2", "Field3", "Field4"],
    )
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
def test_header_with_index_col(all_parsers):
    parser = all_parsers
    data = """foo,1,2,3
bar,4,5,6
baz,7,8,9
"""
    names = ["A", "B", "C"]
    result = parser.read_csv(StringIO(data), names=names)

    expected = DataFrame(
        [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
        index=["foo", "bar", "baz"],
        columns=["A", "B", "C"],
    )
    tm.assert_frame_equal(result, expected)

def test_header_not_first_line(all_parsers):
    parser = all_parsers
    data = """got,to,ignore,this,line
got,to,ignore,this,line
index,A,B,C,D
foo,2,3,4,5
bar,7,8,9,10
baz,12,13,14,15
"""
    data2 = """index,A,B,C,D
foo,2,3,4,5
bar,7,8,9,10
baz,12,13,14,15
"""

    result = parser.read_csv(StringIO(data), header=2, index_col=0)
    expected = parser.read_csv(StringIO(data2), header=0, index_col=0)
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
def test_header_multi_index(all_parsers):
    parser = all_parsers
    expected = tm.makeCustomDataframe(5, 3, r_idx_nlevels=2, c_idx_nlevels=4)

    data = """\
C0,,C_l0_g0,C_l0_g1,C_l0_g2

C1,,C_l1_g0,C_l1_g1,C_l1_g2
C2,,C_l2_g0,C_l2_g1,C_l2_g2
C3,,C_l3_g0,C_l3_g1,C_l3_g2
R0,R1,,,
R_l0_g0,R_l1_g0,R0C0,R0C1,R0C2
R_l0_g1,R_l1_g1,R1C0,R1C1,R1C2
R_l0_g2,R_l1_g2,R2C0,R2C1,R2C2
R_l0_g3,R_l1_g3,R3C0,R3C1,R3C2
R_l0_g4,R_l1_g4,R4C0,R4C1,R4C2
"""
    result = parser.read_csv(StringIO(data), header=[0, 1, 2, 3], index_col=[0, 1])
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize(
    "kwargs,msg",
    [
        (
            {"index_col": ["foo", "bar"]},
            (
                "index_col must only contain "
                "row numbers when specifying "
                "a multi-index header"
            ),
        ),
        (
            {"index_col": [0, 1], "names": ["foo", "bar"]},
            ("cannot specify names when specifying a multi-index header"),
        ),
        (
            {"index_col": [0, 1], "usecols": ["foo", "bar"]},
            ("cannot specify usecols when specifying a multi-index header"),
        ),
    ],
)
def test_header_multi_index_invalid(all_parsers, kwargs, msg):
    data = """\
C0,,C_l0_g0,C_l0_g1,C_l0_g2

C1,,C_l1_g0,C_l1_g1,C_l1_g2
C2,,C_l2_g0,C_l2_g1,C_l2_g2
C3,,C_l3_g0,C_l3_g1,C_l3_g2
R0,R1,,,
R_l0_g0,R_l1_g0,R0C0,R0C1,R0C2
R_l0_g1,R_l1_g1,R1C0,R1C1,R1C2
R_l0_g2,R_l1_g2,R2C0,R2C1,R2C2
R_l0_g3,R_l1_g3,R3C0,R3C1,R3C2
R_l0_g4,R_l1_g4,R4C0,R4C1,R4C2
"""
    parser = all_parsers

    with pytest.raises(ValueError, match=msg):
        parser.read_csv(StringIO(data), header=[0, 1, 2, 3], **kwargs)

_TestTuple = namedtuple("_TestTuple", ["first", "second"])

@skip_pyarrow
@pytest.mark.parametrize(
    "kwargs",
    [
        {"header": [0, 1]},
        {
            "skiprows": 3,
            "names": [
                ("a", "q"),
                ("a", "r"),
                ("a", "s"),
                ("b", "t"),
                ("c", "u"),
                ("c", "v"),
            ],
        },
        {
            "skiprows": 3,
            "names": [
                _TestTuple("a", "q"),
                _TestTuple("a", "r"),
                _TestTuple("a", "s"),
                _TestTuple("b", "t"),
                _TestTuple("c", "u"),
                _TestTuple("c", "v"),
            ],
        },
    ],
)
def test_header_multi_index_common_format1(all_parsers, kwargs):
    parser = all_parsers
    expected = DataFrame(
        [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]],
        index=["one", "two"],
        columns=MultiIndex.from_tuples(
            [("a", "q"), ("a", "r"), ("a", "s"), ("b", "t"), ("c", "u"), ("c", "v")]
        ),
    )
    data = """,a,a,a,b,c,c
,q,r,s,t,u,v
,,,,,,
one,1,2,3,4,5,6
two,7,8,9,10,11,12"""

    result = parser.read_csv(StringIO(data), index_col=0, **kwargs)
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
@pytest.mark.parametrize(
    "kwargs",
    [
        {"header": [0, 1]},
        {
            "skiprows": 2,
            "names": [
                ("a", "q"),
                ("a", "r"),
                ("a", "s"),
                ("b", "t"),
                ("c", "u"),
                ("c", "v"),
            ],
        },
        {
            "skiprows": 2,
            "names": [
                _TestTuple("a", "q"),
                _TestTuple("a", "r"),
                _TestTuple("a", "s"),
                _TestTuple("b", "t"),
                _TestTuple("c", "u"),
                _TestTuple("c", "v"),
            ],
        },
    ],
)
def test_header_multi_index_common_format2(all_parsers, kwargs):
    parser = all_parsers
    expected = DataFrame(
        [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]],
        index=["one", "two"],
        columns=MultiIndex.from_tuples(
            [("a", "q"), ("a", "r"), ("a", "s"), ("b", "t"), ("c", "u"), ("c", "v")]
        ),
    )
    data = """,a,a,a,b,c,c
,q,r,s,t,u,v
one,1,2,3,4,5,6
two,7,8,9,10,11,12"""

    result = parser.read_csv(StringIO(data), index_col=0, **kwargs)
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
@pytest.mark.parametrize(
    "kwargs",
    [
        {"header": [0, 1]},
        {
            "skiprows": 2,
            "names": [
                ("a", "q"),
                ("a", "r"),
                ("a", "s"),
                ("b", "t"),
                ("c", "u"),
                ("c", "v"),
            ],
        },
        {
            "skiprows": 2,
            "names": [
                _TestTuple("a", "q"),
                _TestTuple("a", "r"),
                _TestTuple("a", "s"),
                _TestTuple("b", "t"),
                _TestTuple("c", "u"),
                _TestTuple("c", "v"),
            ],
        },
    ],
)
def test_header_multi_index_common_format3(all_parsers, kwargs):
    parser = all_parsers
    expected = DataFrame(
        [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]],
        index=["one", "two"],
        columns=MultiIndex.from_tuples(
            [("a", "q"), ("a", "r"), ("a", "s"), ("b", "t"), ("c", "u"), ("c", "v")]
        ),
    )
    expected = expected.reset_index(drop=True)
    data = """a,a,a,b,c,c
q,r,s,t,u,v
1,2,3,4,5,6
7,8,9,10,11,12"""

    result = parser.read_csv(StringIO(data), index_col=None, **kwargs)
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
def test_header_multi_index_common_format_malformed1(all_parsers):
    parser = all_parsers
    expected = DataFrame(
        np.array([[2, 3, 4, 5, 6], [8, 9, 10, 11, 12]], dtype="int64"),
        index=Index([1, 7]),
        columns=MultiIndex(
            levels=[["a", "b", "c"], ["r", "s", "t", "u", "v"]],
            codes=[[0, 0, 1, 2, 2], [0, 1, 2, 3, 4]],
            names=["a", "q"],
        ),
    )
    data = """a,a,a,b,c,c
q,r,s,t,u,v
1,2,3,4,5,6
7,8,9,10,11,12"""

    result = parser.read_csv(StringIO(data), header=[0, 1], index_col=0)
    tm.assert_frame_equal(expected, result)

@skip_pyarrow
def test_header_multi_index_common_format_malformed2(all_parsers):
    parser = all_parsers
    expected = DataFrame(
        np.array([[2, 3, 4, 5, 6], [8, 9, 10, 11, 12]], dtype="int64"),
        index=Index([1, 7]),
        columns=MultiIndex(
            levels=[["a", "b", "c"], ["r", "s", "t", "u", "v"]],
            codes=[[0, 0, 1, 2, 2], [0, 1, 2, 3, 4]],
            names=[None, "q"],
        ),
    )

    data = """,a,a,b,c,c
q,r,s,t,u,v
1,2,3,4,5,6
7,8,9,10,11,12"""

    result = parser.read_csv(StringIO(data), header=[0, 1], index_col=0)
    tm.assert_frame_equal(expected, result)

@skip_pyarrow
def test_header_multi_index_common_format_malformed3(all_parsers):
    parser = all_parsers
    expected = DataFrame(
        np.array([[3, 4, 5, 6], [9, 10, 11, 12]], dtype="int64"),
        index=MultiIndex(levels=[[1, 7], [2, 8]], codes=[[0, 1], [0, 1]]),
        columns=MultiIndex(
            levels=[["a", "b", "c"], ["s", "t", "u", "v"]],
            codes=[[0, 1, 2, 2], [0, 1, 2, 3]],
            names=[None, "q"],
        ),
    )
    data = """,a,a,b,c,c
q,r,s,t,u,v
1,2,3,4,5,6
7,8,9,10,11,12"""

    result = parser.read_csv(StringIO(data), header=[0, 1], index_col=[0, 1])
    tm.assert_frame_equal(expected, result)

@skip_pyarrow
def test_header_multi_index_blank_line(all_parsers):
    # GH 40442
    parser = all_parsers
    data = [[None, None], [1, 2], [3, 4]]
    columns = MultiIndex.from_tuples([("a", "A"), ("b", "B")])
    expected = DataFrame(data, columns=columns)
    data = "a,b\nA,B\n,\n1,2\n3,4"
    result = parser.read_csv(StringIO(data), header=[0, 1])
    tm.assert_frame_equal(expected, result)

@skip_pyarrow
@pytest.mark.parametrize(
    "data,header", [("1,2,3\n4,5,6", None), ("foo,bar,baz\n1,2,3\n4,5,6", 0)]
)
def test_header_names_backward_compat(all_parsers, data, header):
    # see gh-2539
    parser = all_parsers
    expected = parser.read_csv(StringIO("1,2,3\n4,5,6"), names=["a", "b", "c"])

    result = parser.read_csv(StringIO(data), names=["a", "b", "c"], header=header)
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
@pytest.mark.parametrize("kwargs", [{}, {"index_col": False}])
def test_read_only_header_no_rows(all_parsers, kwargs):
    # See gh-7773
    parser = all_parsers
    expected = DataFrame(columns=["a", "b", "c"])

    result = parser.read_csv(StringIO("a,b,c"), **kwargs)
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize(
    "kwargs,names",
    [
        ({}, [0, 1, 2, 3, 4]),
        ({"prefix": "X"}, ["X0", "X1", "X2", "X3", "X4"]),
        (
            {"names": ["foo", "bar", "baz", "quux", "panda"]},
            ["foo", "bar", "baz", "quux", "panda"],
        ),
    ],
)
def test_no_header(all_parsers, kwargs, names):
    parser = all_parsers
    data = """1,2,3,4,5
6,7,8,9,10
11,12,13,14,15
"""
    expected = DataFrame(
        [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]], columns=names
    )
    result = parser.read_csv(StringIO(data), header=None, **kwargs)
    tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize("header", [["a", "b"], "string_header"])
def test_non_int_header(all_parsers, header):
    # see gh-16338
    msg = "header must be integer or list of integers"
    data = """1,2\n3,4"""
    parser = all_parsers

    with pytest.raises(ValueError, match=msg):
        parser.read_csv(StringIO(data), header=header)

@skip_pyarrow
def test_singleton_header(all_parsers):
    # see gh-7757
    data = """a,b,c\n0,1,2\n1,2,3"""
    parser = all_parsers

    expected = DataFrame({"a": [0, 1], "b": [1, 2], "c": [2, 3]})
    result = parser.read_csv(StringIO(data), header=[0])
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
@pytest.mark.parametrize(
    "data,expected",
    [
        (
            "A,A,A,B\none,one,one,two\n0,40,34,0.1",
            DataFrame(
                [[0, 40, 34, 0.1]],
                columns=MultiIndex.from_tuples(
                    [("A", "one"), ("A", "one.1"), ("A", "one.2"), ("B", "two")]
                ),
            ),
        ),
        (
            "A,A,A,B\none,one,one.1,two\n0,40,34,0.1",
            DataFrame(
                [[0, 40, 34, 0.1]],
                columns=MultiIndex.from_tuples(
                    [("A", "one"), ("A", "one.1"), ("A", "one.1.1"), ("B", "two")]
                ),
            ),
        ),
        (
            "A,A,A,B,B\none,one,one.1,two,two\n0,40,34,0.1,0.1",
            DataFrame(
                [[0, 40, 34, 0.1, 0.1]],
                columns=MultiIndex.from_tuples(
                    [
                        ("A", "one"),
                        ("A", "one.1"),
                        ("A", "one.1.1"),
                        ("B", "two"),
                        ("B", "two.1"),
                    ]
                ),
            ),
        ),
    ],
)
def test_mangles_multi_index(all_parsers, data, expected):
    # see gh-18062
    parser = all_parsers

    result = parser.read_csv(StringIO(data), header=[0, 1])
    tm.assert_frame_equal(result, expected)

@skip_pyarrow
@pytest.mark.parametrize("index_col", [None, [0]])
@pytest.mark.parametrize(
    "columns", [None, (["", "Unnamed"]), (["Unnamed", ""]), (["Unnamed", "NotUnnamed"])]
)
def test_multi_index_unnamed(all_parsers, index_col, columns):
    # see gh-23687
    #
    # When specifying a multi-index header, make sure that
    # we don't error just because one of the rows in our header
    # has ALL column names containing the string "Unnamed". The
    # correct condition to check is whether the row contains
    # ALL columns that did not have names (and instead were given
    # placeholder ones).
    parser = all_parsers
    header = [0, 1]

    if index_col is None:
        data = ",".join(columns or ["", ""]) + "\n0,1\n2,3\n4,5\n"
    else:
        data = ",".join([""] + (columns or ["", ""])) + "\n,0,1\n0,2,3\n1,4,5\n"

    if columns is None:
        msg = (
            r"Passed header=\[0,1\] are too "
            r"many rows for this multi_index of columns"
        )
        with pytest.raises(ParserError, match=msg):
            parser.read_csv(StringIO(data), header=header, index_col=index_col)
    else:
        result = parser.read_csv(StringIO(data), header=header, index_col=index_col)
        exp_columns = []

        for i, col in enumerate(columns):
            if not col:  # Unnamed.
                col = f"Unnamed: {i if index_col is None else i + 1}_level_0"

            exp_columns.append(col)

        columns = MultiIndex.from_tuples(zip(exp_columns, ["0", "1"]))
        expected = DataFrame([[2, 3], [4, 5]], columns=columns)
        tm.assert_frame_equal(result, expected)

@skip_pyarrow
def test_read_csv_multiindex_columns(all_parsers):
    # GH#6051
    parser = all_parsers

    s1 = "Male, Male, Male, Female, Female\nR, R, L, R, R\n.86, .67, .88, .78, .81"
    s2 = (
        "Male, Male, Male, Female, Female\n"
        "R, R, L, R, R\n"
        ".86, .67, .88, .78, .81\n"
        ".86, .67, .88, .78, .82"
    )

    mi = MultiIndex.from_tuples(
        [
            ("Male", "R"),
            (" Male", " R"),
            (" Male", " L"),
            (" Female", " R"),
            (" Female", " R.1"),
        ]
    )
    expected = DataFrame(
        [[0.86, 0.67, 0.88, 0.78, 0.81], [0.86, 0.67, 0.88, 0.78, 0.82]], columns=mi
    )

    df1 = parser.read_csv(StringIO(s1), header=[0, 1])
    tm.assert_frame_equal(df1, expected.iloc[:1])
    df2 = parser.read_csv(StringIO(s2), header=[0, 1])
    tm.assert_frame_equal(df2, expected)

@skip_pyarrow
def test_read_csv_multi_header_length_check(all_parsers):
    # GH#43102
    parser = all_parsers

    case = """row11,row12,row13
row21,row22, row23
row31,row32
"""

    with pytest.raises(
        ParserError, match="Header rows must have an equal number of columns."
    ):
        parser.read_csv(StringIO(case), header=[0, 2])


import numpy as np                      #the numpy library
import matplotlib.pyplot as plt         #Matplot's pyplot

import sys                              #gives access to a c-like sys library
import os                               #gives access to operating system

print(sys.argv)                     #prints any command line arguments, incl program name
print(os.getcwd())                  #prnits current work directory


import cv2
import numpy as np

def find_way_from_maze(image: np.ndarray) -> tuple:
    """
    Найти путь через лабиринт.

    :param image: изображение лабиринта
    :return: координаты пути из лабиринта в виде (x, y), где x и y - это массивы координат
    """
    coords = None
    # Ваш код тут
    pass
    # Ваш код тут

    return coords