#!/usr/bin/env python
# coding: utf-8

#
#
# # Earth Analytics Education - EA  Python Course Spring 2021

# ## Important  - Assignment Guidelines
#
# 1. Before you submit your assignment to GitHub, make sure to run the entire notebook with a fresh kernel. To do this first, **restart the kernel** (in the menubar, select Kernel$\rightarrow$Restart & Run All)
# 2. Always replace the `raise NotImplementedError()` code with your code that addresses the activity challenge. If you don't replace that code, your notebook will not run.
#
# ```
# # YOUR CODE HERE
# raise NotImplementedError()
# ```
#
# 3. Any open ended questions will have a "YOUR ANSWER HERE" within a markdown cell. Replace that text with your answer also formatted using Markdown.
# 4. **DO NOT RENAME THIS NOTEBOOK File!** If the file name changes, the autograder will not grade your assignment properly.
# 6. When you create a figure, comment out `plt.show()` to ensure the autograder can grade your plots. For figure cells, DO NOT DELETE the code that says `DO NOT REMOVE LINE BELOW`.
#
# ```
# ### DO NOT REMOVE LINE BELOW ###
# student_plot1_ax = nb.convert_axes(plt)
# ```
#
# * Only include the package imports, code, and outputs that are required to run your homework assignment.
# * Be sure that your code can be run on any operating system. This means that:
#    1. the data should be downloaded in the notebook to ensure it's reproducible
#    2. all paths should be created dynamically using the `os.path.join`
#
# ## Follow to PEP 8 Syntax Guidelines & Documentation
#
# * Run the `autopep8` tool on all cells prior to submitting (HINT: hit shift + the tool to run it on all cells at once!
# * Use clear and expressive names for variables.
# * Organize your code to support readability.
# * Check for code line length
# * Use comments and white space sparingly where it is needed
# * Make sure all python imports are at the top of your notebook and follow PEP 8 order conventions
# * Spell check your Notebook before submitting it.
#
# For all of the plots below, be sure to do the following:
#
# * Make sure each plot has a clear TITLE and, where appropriate, label the x and y axes. Be sure to include UNITS in your labels.
#

# ### Add Your Name Below
# Leah Manak
#

#

# ---

# # Week 04 and 05 Homework - Automate NDVI Workflow
#
# For this assignment, you will write code to generate a plot of the mean normalized difference vegetation index (NDVI) for two different sites in the United States across one year of data:
#
# * San Joaquin Experimental Range (SJER) in Southern California, United States
# * Harvard Forest (HARV) in the Northeastern United States
#
# The data that you will use for this week is available from **earthpy** using the following download:
#
# `et.data.get_data('ndvi-automation')`
#
# ## Assignment Goals
#
# Your goal in this assignment is to create the most efficient and concise workflow that you can that allows for:
#
# 1. The code to scale if you added new sites or more time periods to the analysis.
# 2. Someone else to understand your workflow.
# 3. The LEAST and most efficient (i.e. runs fast, minimize repetition) amount of code that completes the task.
#
# ### HINTS
#
# * Remove values outside of the landsat valid range of values as specified in the metadata, as needed.
# * Keep any output files SEPARATE FROM input files. Outputs should be created in an outputs directory that is created in the code (if needed) and/or tested for.
# * Use the functions that we demonstrated during class to make your workflow more efficient.
# * BONUS - if you  chose - you can export your data as a csv file. You will get bonus points for doing this.
#
#
# ## Assignment Requirements
#
# Your submission to the GitHub repository should include:
# * This Jupyter Notebook file (.ipynb) with:
#     * The code to create a plot of mean NDVI across a year for  2 NEON Field Sites:
#         * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object
#     * The **data should be cleaned to remove the influence of clouds**. See the [earthdatascience website for an example of what your plot might look like with and without removal of clouds](https://www.earthdatascience.org/courses/earth-analytics-python/create-efficient-data-workflows/).
# * BONUS: Create one output `.csv` file that has 3 columns - NDVI, Date and Site Name - with values for SJER and HARV.
#
# Your notebook should:
# * Have *at least* 2 well documented and well named functions with docstrings.
# * Include a Markdown cell at the top of the notebook that outlines the overall workflow using pseudocode (i.e. plain language, not code)
# * Include additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Replace this cell with your pseudocode  for this workflow
#
# If you happen to be a diagram person a diagram is ok too
#
#

# # Pseudocode for assignment workflow
#
# - Import libraries that will be used in this notebook
# - Make sure unused libraries are not in the list
# - Get the home directory using "et.data.get_data('ndvi-automation')"
# - Create data paths to our data from 'ndvi-automation' downloaded from earthpy
# - Create cloud mask and get cloud pixel values from earthpy
# - Create two functions:
#     1. extract the sitename and datetime from pathnames
#     2. open, crop, and identify ranges of a landsat band
# - Clip extent of bands to match study sites boundaries
# - Clip extent of qa layers to match study sites boundaries
# - Calculate NDVI
# - Add a cloud mask to the NDVI values
# - Calculate Mean NDVI, get site names and the dates, and create a list of lists
# - Make a pandas dataframe including the site names, mean NDVIs, and dates
# - Make sure the dataframe index is in datetime and that the date column is the index
# - Make sure the NA data is not included
# - Plot the data with date on the x-axis and mean NDVI values on the y axis
# - Download the data to a CSV file

# In[1]:

# Autograding imports - do not modify this cell
import matplotcheck.autograde as ag
import matplotcheck.notebook as nb
import matplotcheck.timeseries as ts
from datetime import datetime

# In[2]:

# Import needed packages in PEP 8 order
# and no unused imports listed (10 points total)
import os
from glob import glob

import matplotlib.pyplot as plt
import pandas as pd
import rioxarray as rxr
import xarray as xr
import geopandas as gpd
import earthpy as et
import earthpy.mask as em
from datetime import datetime
import numpy as np
from matplotlib.dates import DateFormatter

# Download the data
et.data.get_data('ndvi-automation')

# Create a path to the directory
directory_path = os.path.join(et.io.HOME,
                              "earth-analytics",
                              "data")

# Set working directory
os.chdir(directory_path)

# In[3]:

# DO NOT MODIFY THIS CELL
# Tests that the working directory is set to earth-analytics/data

path = os.path.normpath(os.getcwd())
student_wd_parts = path.split(os.sep)

if student_wd_parts[-2:] == ['earth-analytics', 'data']:
    print("\u2705 Great - it looks like your working directory is set correctly to ~/earth-analytics/data")
else:
    print("\u274C Oops, the autograder will not run unless your working directory is set to earth-analytics/data")

# # Figure 1: Plot 1 - Mean NDVI For Each Site Across the Year (50 points)
#
# Create a plot of the mean normalized difference vegetation index (NDVI) for the two different sites in the United States across the year:
#
# * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object.
# * Each site should be identified with a different color in the plot and legend.
# * The final plot **data should be cleaned to remove the influence of clouds**.
# * Be sure to include appropriate title and axes labels.
#
# Add additional cells as needed for processing data (e.g. defining functions, etc), but be sure to:
# * follow the instructions in the code cells that have been provided to ensure that you are able to use the sanity check tests that are provided.
# * include only the plot code in the cell identified for the final plot code below

# ## Task 1:
#
# In the cell below, create a single dataframe containing MEAN NDVI, the site name,
# and the date of the data for the HARV site
# scene `HARV/landsat-crop/LC080130302017031701T1-SC20181023151837`.  The column names for the  final
# DataFrame should be`mean_ndvi`, and `site`, and the data should be **indexed on the date**.
#
# Use the functions that we reviewed in class (or create your own versions of them) to implement your code
#
# ### In the Cell below Place  All Functions Needed to Run this Notebook (20 points)

# In[4]:

### DO NOT REMOVE THIS LINE OR EDIT / MOVE THIS CELL ###
start_time = datetime.now()

# In[5]:

# In this cell place all of the functions needed to run your notebook
# You will be graded here on function application, docstrings, efficiency so ensure
# All functions are placed here!

# A function to extract the datetime and sitename from the directory paths
def extract_date_sitename(directory_path,
                          sitename_location,
                          datetime_location):
    """Extract datetime and sitename from directory path names.

    Parameters
    -----------
    directory_path : string
        A path to the directory name
    sitename_location : index list
        Index of sitename location in directory path name
    datetime_location : index list
        Index of datetime location in directory path name

    Returns
    -----------
    list : list of the datetime location and sitename information
    """
    # Create an empty list to append both sitename and date information
    sitename_date = []

    # Assign datetime location to an object and specify datetime format
    date_location = directory_path[datetime_location[0]: datetime_location[1]]
    format = "%Y%m%d"

    # Create a date varaiable using new object and the datetime format
    date = datetime.strptime(date_location, format)

    # Create a location variable called "site"
    site = directory_path[sitename_location[0]: sitename_location[1]]

    # Append site and date variables to list
    sitename_date.append(site)
    sitename_date.append(date)

    # Return the populated sitename_date list
    return sitename_date

# A function to open clean landsat bands
def open_clean_bands(band_path,
                     crop_extent,
                     valid_range=None):
    """Open, crop, and identify the range of the bands.

    Parameters
    -----------
    band_path : string
        A path to the array that we will open
    valid_range : tuple (optional)
        A tuple of min and max range of values for the data. Default = None

    Returns
    -----------
    arr : xarray DataArray
        An xarray DataArray with values that should be masked
        set to 1 for True (Boolean)
    """
    # tests to ensure the arrays are clipped to the same .shape
    band = rxr.open_rasterio(
        band_path, masked=True).rio.clip(
        crop_extent.geometry,from_disk=True).squeeze()

    # This last step is only for a valid tuple
    if valid_range:
        mask = ((band < valid_range[0]) | (band > valid_range[1]))
        band = band.where(~xr.where(mask, True, False))

    return band

# In[6]:

# Create dataframe of mean NDVI in this cell using the functions created above

# Create path to the two sites "SJER" and "HARV"
site_paths = glob(os.path.join("ndvi-automation",
                               "sites",
                               "*"))
site_paths

# Create cloud mask and get the cloud pixel values from earthpy
high_cloud_confidence = (
    em.pixel_flags["pixel_qa"]["L8"]["High Cloud Confidence"])
cloud = em.pixel_flags["pixel_qa"]["L8"]["Cloud"]
cloud_shadow = em.pixel_flags["pixel_qa"]["L8"]["Cloud Shadow"]

all_masked_values = cloud_shadow + cloud + high_cloud_confidence
harv_path = glob(os.path.join("ndvi-automation", "sites", "HARV"))

# Open and clean all HARV: first create empty HARV list
harv_info = []

# Create a for loop for the HARV path
for path in harv_path:
    # Establish the scene directory path
    scene_path = glob(os.path.join(path, "landsat-crop",
                      "LC080130302017031701T1-SC20181023151837"))
    # Set the path to the cropped shapefile and open as the HARV boundary
    bound = os.path.join(path, "vector", "HARV-crop.shp")
    harv_boundary = gpd.read_file(bound)

    # Create a nested for loop associated with each .tif file (band 4-5)
    for tifs in scene_path:
        # Get site and date info from the scene directory path
        site_info = extract_date_sitename(tifs, [22, 26], [50, 58])
        # Order the bands 4-5 with glob
        harv_bands = sorted(glob(os.path.join(tifs, "*band[4-5]*")))
        # Set the path to the qa layer in the scene directory and open it
        qa_layer_path = os.path.join(tifs,
                                     "LC08_L1TP_013030_20170317_20170328_01_T1_pixel_qa.tif")
        qa_layer = rxr.open_rasterio(qa_layer_path, masked=True)
        # Crop the qa layer using the harv_boundary
        cropped_qa = qa_layer.rio.clip(harv_boundary.geometry).squeeze()

        # New empty list for bands without cloud interference
        tif_bands = []
        # Create an additional loop for the bands in harv
        for a_band in harv_bands:
            # Clean the band using the open_clean_bands function
            clean_band = open_clean_bands(
                a_band, harv_boundary, valid_range=(0, 10000))
            # Apply the cloud mask to the clean band
            band_cloud_mask = clean_band.where(
                ~qa_layer.isin(all_masked_values))
            # Add clean bands to empty for calculating the mean NDVI
            tif_bands.append(band_cloud_mask)

        # Calculate mean NDVI using tif_bands list
        mean_ndvi = np.nanmean(
            (tif_bands[1]-tif_bands[0]) / (tif_bands[1]+tif_bands[0]))
        # Append the mean NDVI to the site_info list
        site_info.append(mean_ndvi)
        # Append site_info list to the initial list prior to the for loop
        #called "harv_info"
        harv_info.append(site_info)

# Create a pandas dataframe with the harv_info list
harv_df = pd.DataFrame(harv_info, columns=[
                            "site", "date", "mean_ndvi"])

# Set index from the date
harv_final = harv_df.set_index("date")

# Call dataframe
harv_final

# In[7]:

# This cell  is testing your data output above

student_ndvi_ts_single_site = _

single_scene_points = 0

# Ensure the data is stored in a dataframe.
if isinstance(student_ndvi_ts_single_site, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    single_scene_points += 1
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Ensure that the date column is the index
if isinstance(student_ndvi_ts_single_site.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    single_scene_points += 2
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_ts_single_site.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    single_scene_points += 2
else:
    print('\u274C The data in your date column is not datetime.')

# Ensure the site name is correct
if student_ndvi_ts_single_site.site.values[0] == 'HARV':
    print('\u2705 You have the correct site name!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct site name.')

if np.allclose(0.281131628228094, student_ndvi_ts_single_site.mean_ndvi.values[0]):
    print('\u2705 You have the correct mean NDVI value!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct mean ndvi value.')

print("\n \u27A1 You received {} out of 15 points for creating a dataframe.".format(
    single_scene_points))
single_scene_points

# ## Task 2:
#
# In the cell below, process all of the landsat scenes. Create a DataFrame that contains the following
# information for each scene
#
#
# |   | index  | site  | mean_ndvi  |
# |---|---|---|---|
# | Date  |   |   |   |
# |  2017-01-07  | 0  | SJER  | .4  |
#
# Be sure to call your dataframe at the end of the cell to ensure autograding works.
# HINT: FOR THIS STEP, leave any rows containing missing values (`NAN`).

# In[8]:

# Create dataframe of NDVI including the cleaning data to deal with clouds

# YOUR CODE HERE
# Create an empty list with site, date, and mean ndvi information
location_info = []

# make a for loop for the location paths
for site in site_paths:
    # Get list of the location path with glob
    locations = glob(os.path.join(site, "landsat-crop", "*"))
    # grab the shapefiles from the locations and make the index 0
    bounds = glob(os.path.join(site, "vector", "*-crop.shp"))[0]
    # Open the shapefiles
    opened_bound = gpd.read_file(bounds)

    # Create a nested for loop for all locations
    for all_locations in locations:
        # Extract date and site info using "extract_date_sitename" function
        site_info = extract_date_sitename(all_locations, [22, 26], [50, 58])
        # Create sorted list of bands 5 & 5 in each location using glob
        scene_bands = sorted(glob(os.path.join(all_locations, "*band[4-5]*")))
        # Extract qa pixel layers using glob and pulling out index 0.
        qa_layer_paths = glob(os.path.join(all_locations, "*pixel_qa*"))[0]
        # Open the qa layer
        opened_layer = rxr.open_rasterio(qa_layer_paths, masked=True)
        # Crop the qa layer using the 'opened_layer' shapefile
        cropped_layer = opened_layer.rio.clip(opened_bound.geometry).squeeze()

        # Create an empty list for cleaned bands 4 and 5
        site_bands = []

        # Create a for loop to clean bands 4&5
        for band in scene_bands:
            # Clean the bands using 'open_clean_bands' function
            clean_band = open_clean_bands(
                band, opened_bound, valid_range=(0, 10000))
            # Apply cloud mask
            cloud_free_band = clean_band.where(
                ~cropped_layer.isin(all_masked_values))
            # Append list with the cloud free bands
            site_bands.append(cloud_free_band)

        # Calculate mean NDVI
        mean_ndvi = np.nanmean(
            (site_bands[1]-site_bands[0]) / (site_bands[1]+site_bands[0]))
        # Append the mean NDVI to the empty site_info list
        site_info.append(mean_ndvi)
        # Append this list of lists the empty location_info list
        location_info.append(site_info)

# Create a pandas dataframe
location_info_df = pd.DataFrame(location_info, columns=[
                            "site", "date", "mean_ndvi"])

# Set index on date
indexed_location_df = location_info_df.set_index("date")

final_NDVI_df = indexed_location_df.sort_values(by="date")

# Call dataframe
final_NDVI_df

# In[9]:

# Last sanity check before creating your plot (10 points)

# Ensure that you call your dataframe at the bottom of the cell above
# and that it has columns called: mean_ndvi and site

# Ensure the data is stored in a dataframe.
student_ndvi_df = _

df_points = 0

if isinstance(student_ndvi_df, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    df_points +=2
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Check that dataframe contains the appropriate number of NAN values
if student_ndvi_df.mean_ndvi.isna().sum() == 15:
    print('\u2705 Correct number of masked data values!')
    df_points +=2
else:
    print('\u274C The amount of null data in your dataframe is incorrect.')

# Ensure that the date column is the index
if isinstance(student_ndvi_df.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    df_points +=3
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_df.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    df_points +=3
else:
    print('\u274C The data in your date column is not datetime.')

# Output for timer, # DO NOT MODIFY
end_time = datetime.now()
total_time = end_time - start_time
print(
    "Your total run time for processing the data was {0}.".format(total_time))

print("\n \u27A1 You received {} out of 10 points for creating a dataframe.".format(
    df_points))

df_points

# In[10]:

# Add only the plot code to this cell

# This is the final figure of mean NDVI
# for both sites across the year
# with data cleaned to deal with clouds

fig, ax = plt.subplots(figsize = (12, 7))

for site, df in final_NDVI_df.dropna().groupby("site"):
    if site in ["HARV"]:
        site_name = 'HARV'
        color = 'goldenrod'
    else:
        site_name = 'SJER'
        color = 'purple'

    ax.plot(df.index, df.mean_ndvi, label = site_name, marker = 'o',
            color = color)

ax.set(title = "Mean Normalized Difference Vegetation Index (NDVI) for two sites (HARV & SJER) \n Mar 2017 - Dec 2017 (cloud-free data)",
       xlabel = "Month", ylabel = "Mean NDVI")

ax.xaxis.set_major_formatter(DateFormatter("%b"))

ax.legend(title = "Site")

### DO NOT REMOVE LINES BELOW ###
final_masked_solution = nb.convert_axes(plt, which_axes="current")

# In[11]:

# Ignore this cell for the autograding tests

# In[12]:

# Ignore this cell for the autograding tests

# # Question 1 (10 points)
#
# Imagine that you are planning NEON’s upcoming flight season to capture remote sensing data in these locations and want to ensure that you fly the area when the vegetation is the most green.
#
# When would you recommend the flights take place for each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# Based on the plot, I would recommend the flights to take place in different times for the different locations. I would say April for the SJER site, and July for the HARV site.

# # Question 2 (10 points)
#
# How could you modify your workflow to look at vegetation changes over time in each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# Instead of comparing the two sites, I would create a separate plot for each site. Each location plot would have a line for each year as a separate color, showing how the vegetation might change each month over a span of a certain amount of years.

# # Do not edit this cell! (10 points)
#
# The notebook includes:
# * additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Do not edit this cell! (20 points)
#
# The notebook will also be checked for overall clean code requirements as specified at the **top** of this notebook. Some of these requirements include (review the top cells for more specifics):
#
# * Notebook begins at cell [1] and runs on any machine in its entirety.
# * PEP 8 format is applied throughout (including lengths of comment and code lines).
# * No additional code or imports in the notebook that is not needed for the workflow.
# * Notebook is fully reproducible. This means:
#    * reproducible paths using the os module.
#    * data downloaded using code in the notebook.
#    * all imports at top of notebook.

# ## BONUS - Export a  .CSV File to Share (10 points possible)
#
# This is optional - if you export a **.csv** file with the columns specified above: Site, Date and NDVI Value you can get an additional 10 points.
#
# * FULL CREDIT: File exists in csv format and contains the columns specified.
# We will check your github repo for this file!
#

# In[13]:

# CSV needs to have no nan values... drop them with .dropna()
final_NDVI_df_csv = final_NDVI_df.dropna()

# Export pandas dataframe to csv file
final_NDVI_df_csv.to_csv(os.path.join(
    directory_path,
    "ndvi-automation",
    "outputs",
    "ndvi_df.csv"))

# Export to personal
final_NDVI_df_csv.to_csv(os.path.join(et.io.HOME,
                                      "earth-analytics",
                                      "earth-analytics-python-env",
                                      "ea-2022-04-ndvi-automation-LManak",
                                      "ndvi_df.csv"))

# In[ ]:


# -*- coding: utf-8 -*-
"""Part2_Music_Generation_Solution.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1zVISpFRO6la4MCdg2wyKXj-YppLNjQAl

        Run in Google Colab

# Lab 1: Intro to TensorFlow and Music Generation with RNNs

# Part 2: Music Generation with RNNs

In this portion of the lab, we will explore building a Recurrent Neural Network (RNN) for music generation. We will train a model to learn the patterns in raw sheet music in [ABC notation](https://en.wikipedia.org/wiki/ABC_notation) and then use this model to generate new music.

## 2.1 Dependencies
First, let's download the course repository, install dependencies, and import the relevant packages we'll need for this lab.
"""

# Commented out IPython magic to ensure Python compatibility.
# Import Tensorflow 2.0
# %tensorflow_version 2.x
import tensorflow as tf

# Download and import the MIT 6.S191 package
!pip install mitdeeplearning
import mitdeeplearning as mdl

# Import all remaining packages
import numpy as np
import os
import time
import functools
from IPython import display as ipythondisplay
from tqdm import tqdm
!apt-get install abcmidi timidity > /dev/null 2>&1

# Check that we are using a GPU, if not switch runtimes
#   using Runtime > Change Runtime Type > GPU
assert len(tf.config.list_physical_devices('GPU')) > 0

"""## 2.2 Dataset

![Let's Dance!](http://33.media.tumblr.com/3d223954ad0a77f4e98a7b87136aa395/tumblr_nlct5lFVbF1qhu7oio1_500.gif)

We've gathered a dataset of thousands of Irish folk songs, represented in the ABC notation. Let's download the dataset and inspect it:

"""

# Download the dataset
songs = mdl.lab1.load_training_data()

# Print one of the songs to inspect it in greater detail!
example_song = songs[0]
print("\nExample song: ")
print(example_song)

"""We can easily convert a song in ABC notation to an audio waveform and play it back. Be patient for this conversion to run, it can take some time."""

# Convert the ABC notation to audio file and listen to it
mdl.lab1.play_song(example_song)

"""One important thing to think about is that this notation of music does not simply contain information on the notes being played, but additionally there is meta information such as the song title, key, and tempo. How does the number of different characters that are present in the text file impact the complexity of the learning problem? This will become important soon, when we generate a numerical representation for the text data."""

# Join our list of song strings into a single string containing all songs
songs_joined = "\n\n".join(songs)

# Find all unique characters in the joined string
vocab = sorted(set(songs_joined))
print("There are", len(vocab), "unique characters in the dataset")

"""## 2.3 Process the dataset for the learning task

Let's take a step back and consider our prediction task. We're trying to train a RNN model to learn patterns in ABC music, and then use this model to generate (i.e., predict) a new piece of music based on this learned information.

Breaking this down, what we're really asking the model is: given a character, or a sequence of characters, what is the most probable next character? We'll train the model to perform this task.

To achieve this, we will input a sequence of characters to the model, and train the model to predict the output, that is, the following character at each time step. RNNs maintain an internal state that depends on previously seen elements, so information about all characters seen up until a given moment will be taken into account in generating the prediction.

### Vectorize the text

Before we begin training our RNN model, we'll need to create a numerical representation of our text-based dataset. To do this, we'll generate two lookup tables: one that maps characters to numbers, and a second that maps numbers back to characters. Recall that we just identified the unique characters present in the text.
"""

### Define numerical representation of text ###

char2idx = {u:i for i, u in enumerate(vocab)}

idx2char = np.array(vocab)

"""This gives us an integer representation for each character. Observe that the unique characters (i.e., our vocabulary) in the text are mapped as indices from 0 to `len(unique)`. Let's take a peek at this numerical representation of our dataset:"""

print('{')
for char,_ in zip(char2idx, range(20)):
    print('  {:4s}: {:3d},'.format(repr(char), char2idx[char]))
print('  ...\n}')

### Vectorize the songs string ###

'''TODO: Write a function to convert the all songs string to a vectorized
    (i.e., numeric) representation. Use the appropriate mapping
    above to convert from vocab characters to the corresponding indices.

  NOTE: the output of the `vectorize_string` function
  should be a np.array with `N` elements, where `N` is
  the number of characters in the input string
'''
def vectorize_string(string):
  vectorized_output = np.array([char2idx[char] for char in string])
  return vectorized_output

# def vectorize_string(string):
  # TODO

vectorized_songs = vectorize_string(songs_joined)

"""We can also look at how the first part of the text is mapped to an integer representation:"""

print ('{} ---- characters mapped to int ----> {}'.format(repr(songs_joined[:10]), vectorized_songs[:10]))
# check that vectorized_songs is a numpy array
assert isinstance(vectorized_songs, np.ndarray), "returned result should be a numpy array"

"""### Create training examples and targets

Our next step is to actually divide the text into example sequences that we'll use during training. Each input sequence that we feed into our RNN will contain `seq_length` characters from the text. We'll also need to define a target sequence for each input sequence, which will be used in training the RNN to predict the next character. For each input, the corresponding target will contain the same length of text, except shifted one character to the right.

To do this, we'll break the text into chunks of `seq_length+1`. Suppose `seq_length` is 4 and our text is "Hello". Then, our input sequence is "Hell" and the target sequence is "ello".

The batch method will then let us convert this stream of character indices to sequences of the desired size.
"""

### Batch definition to create training examples ###

def get_batch(vectorized_songs, seq_length, batch_size):
  # the length of the vectorized songs string
  n = vectorized_songs.shape[0] - 1
  # randomly choose the starting indices for the examples in the training batch
  idx = np.random.choice(n-seq_length, batch_size)

  '''TODO: construct a list of input sequences for the training batch'''
  input_batch = [vectorized_songs[i : i+seq_length] for i in idx]
  # input_batch = # TODO
  '''TODO: construct a list of output sequences for the training batch'''
  output_batch = [vectorized_songs[i+1 : i+seq_length+1] for i in idx]
  # output_batch = # TODO

  # x_batch, y_batch provide the true inputs and targets for network training
  x_batch = np.reshape(input_batch, [batch_size, seq_length])
  y_batch = np.reshape(output_batch, [batch_size, seq_length])
  return x_batch, y_batch

# Perform some simple tests to make sure your batch function is working properly!
test_args = (vectorized_songs, 10, 2)
if not mdl.lab1.test_batch_func_types(get_batch, test_args) or \
   not mdl.lab1.test_batch_func_shapes(get_batch, test_args) or \
   not mdl.lab1.test_batch_func_next_step(get_batch, test_args):
   print("======\n[FAIL] could not pass tests")
else:
   print("======\n[PASS] passed all tests!")

"""For each of these vectors, each index is processed at a single time step. So, for the input at time step 0, the model receives the index for the first character in the sequence, and tries to predict the index of the next character. At the next timestep, it does the same thing, but the RNN considers the information from the previous step, i.e., its updated state, in addition to the current input.

We can make this concrete by taking a look at how this works over the first several characters in our text:
"""

x_batch, y_batch = get_batch(vectorized_songs, seq_length=5, batch_size=1)

for i, (input_idx, target_idx) in enumerate(zip(np.squeeze(x_batch), np.squeeze(y_batch))):
    print("Step {:3d}".format(i))
    print("  input: {} ({:s})".format(input_idx, repr(idx2char[input_idx])))
    print("  expected output: {} ({:s})".format(target_idx, repr(idx2char[target_idx])))

"""## 2.4 The Recurrent Neural Network (RNN) model

Now we're ready to define and train a RNN model on our ABC music dataset, and then use that trained model to generate a new song. We'll train our RNN using batches of song snippets from our dataset, which we generated in the previous section.

The model is based off the LSTM architecture, where we use a state vector to maintain information about the temporal relationships between consecutive characters. The final output of the LSTM is then fed into a fully connected [`Dense`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense) layer where we'll output a softmax over each character in the vocabulary, and then sample from this distribution to predict the next character.

As we introduced in the first portion of this lab, we'll be using the Keras API, specifically, [`tf.keras.Sequential`](https://www.tensorflow.org/api_docs/python/tf/keras/models/Sequential), to define the model. Three layers are used to define the model:

* [`tf.keras.layers.Embedding`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding): This is the input layer, consisting of a trainable lookup table that maps the numbers of each character to a vector with `embedding_dim` dimensions.
* [`tf.keras.layers.LSTM`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/LSTM): Our LSTM network, with size `units=rnn_units`.
* [`tf.keras.layers.Dense`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense): The output layer, with `vocab_size` outputs.

### Define the RNN model

Now, we will define a function that we will use to actually build the model.
"""

def LSTM(rnn_units):
  return tf.keras.layers.LSTM(
    rnn_units,
    return_sequences=True,
    recurrent_initializer='glorot_uniform',
    recurrent_activation='sigmoid',
    stateful=True,
  )

"""The time has come! Fill in the `TODOs` to define the RNN model within the `build_model` function, and then call the function you just defined to instantiate the model!"""

### Defining the RNN Model ###

'''TODO: Add LSTM and Dense layers to define the RNN model using the Sequential API.'''
def build_model(vocab_size, embedding_dim, rnn_units, batch_size):
  model = tf.keras.Sequential([

    tf.keras.layers.Embedding(vocab_size, embedding_dim, batch_input_shape=[batch_size, None]),

    LSTM(rnn_units),
    # LSTM('''TODO'''),

    tf.keras.layers.Dense(vocab_size)
    # '''TODO: DENSE LAYER HERE'''
  ])

  return model

model = build_model(len(vocab), embedding_dim=256, rnn_units=1024, batch_size=32)

"""### Test out the RNN model

It's always a good idea to run a few simple checks on our model to see that it behaves as expected.

First, we can use the `Model.summary` function to print out a summary of our model's internal workings. Here we can check the layers in the model, the shape of the output of each of the layers, the batch size, etc.
"""

model.summary()

"""We can also quickly check the dimensionality of our output, using a sequence length of 100. Note that the model can be run on inputs of any length."""

x, y = get_batch(vectorized_songs, seq_length=100, batch_size=32)
pred = model(x)
print("Input shape:      ", x.shape, " # (batch_size, sequence_length)")
print("Prediction shape: ", pred.shape, "# (batch_size, sequence_length, vocab_size)")

"""### Predictions from the untrained model

Let's take a look at what our untrained model is predicting.

To get actual predictions from the model, we sample from the output distribution, which is defined by a `softmax` over our character vocabulary. This will give us actual character indices. This means we are using a [categorical distribution](https://en.wikipedia.org/wiki/Categorical_distribution) to sample over the example prediction. This gives a prediction of the next character (specifically its index) at each timestep.

Note here that we sample from this probability distribution, as opposed to simply taking the `argmax`, which can cause the model to get stuck in a loop.

Let's try this sampling out for the first example in the batch.
"""

sampled_indices = tf.random.categorical(pred[0], num_samples=1)
sampled_indices = tf.squeeze(sampled_indices,axis=-1).numpy()
sampled_indices

"""We can now decode these to see the text predicted by the untrained model:"""

print("Input: \n", repr("".join(idx2char[x[0]])))
print()
print("Next Char Predictions: \n", repr("".join(idx2char[sampled_indices])))

"""As you can see, the text predicted by the untrained model is pretty nonsensical! How can we do better? We can train the network!

## 2.5 Training the model: loss and training operations

Now it's time to train the model!

At this point, we can think of our next character prediction problem as a standard classification problem. Given the previous state of the RNN, as well as the input at a given time step, we want to predict the class of the next character -- that is, to actually predict the next character.

To train our model on this classification task, we can use a form of the `crossentropy` loss (negative log likelihood loss). Specifically, we will use the [`sparse_categorical_crossentropy`](https://www.tensorflow.org/api_docs/python/tf/keras/losses/sparse_categorical_crossentropy) loss, as it utilizes integer targets for categorical classification tasks. We will want to compute the loss using the true targets -- the `labels` -- and the predicted targets -- the `logits`.

Let's first compute the loss using our example predictions from the untrained model:
"""

### Defining the loss function ###

'''TODO: define the loss function to compute and return the loss between
    the true labels and predictions (logits). Set the argument from_logits=True.'''
def compute_loss(labels, logits):
  loss = tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
  # loss = tf.keras.losses.sparse_categorical_crossentropy('''TODO''', '''TODO''', from_logits=True) # TODO
  return loss

'''TODO: compute the loss using the true next characters from the example batch
    and the predictions from the untrained model several cells above'''
example_batch_loss = compute_loss(y, pred)
# example_batch_loss = compute_loss('''TODO''', '''TODO''') # TODO

print("Prediction shape: ", pred.shape, " # (batch_size, sequence_length, vocab_size)")
print("scalar_loss:      ", example_batch_loss.numpy().mean())

"""Let's start by defining some hyperparameters for training the model. To start, we have provided some reasonable values for some of the parameters. It is up to you to use what we've learned in class to help optimize the parameter selection here!"""

### Hyperparameter setting and optimization ###

# Optimization parameters:
num_training_iterations = 2000  # Increase this to train longer
batch_size = 4  # Experiment between 1 and 64
seq_length = 100  # Experiment between 50 and 500
learning_rate = 5e-3  # Experiment between 1e-5 and 1e-1

# Model parameters:
vocab_size = len(vocab)
embedding_dim = 256
rnn_units = 1024  # Experiment between 1 and 2048

# Checkpoint location:
checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "my_ckpt")

"""Now, we are ready to define our training operation -- the optimizer and duration of training -- and use this function to train the model. You will experiment with the choice of optimizer and the duration for which you train your models, and see how these changes affect the network's output. Some optimizers you may like to try are [`Adam`](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam?version=stable) and [`Adagrad`](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad?version=stable).

First, we will instantiate a new model and an optimizer. Then, we will use the [`tf.GradientTape`](https://www.tensorflow.org/api_docs/python/tf/GradientTape) method to perform the backpropagation operations.

We will also generate a print-out of the model's progress through training, which will help us easily visualize whether or not we are minimizing the loss.
"""

### Define optimizer and training operation ###

'''TODO: instantiate a new model for training using the `build_model`
  function and the hyperparameters created above.'''
model = build_model(vocab_size, embedding_dim, rnn_units, batch_size)
# model = build_model('''TODO: arguments''')

'''TODO: instantiate an optimizer with its learning rate.
  https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/
  Using Adam optimizer to start.'''
optimizer = tf.keras.optimizers.Adam(learning_rate)
# optimizer = # TODO

@tf.function
def train_step(x, y):
  # Use tf.GradientTape()
  with tf.GradientTape() as tape:

    '''TODO: feed the current input into the model and generate predictions'''
    y_hat = model(x) # TODO
    # y_hat = model('''TODO''')

    '''TODO: compute the loss!'''
    loss = compute_loss(y, y_hat) # TODO
    # loss = compute_loss('''TODO''', '''TODO''')

  grads = tape.gradient(loss, model.trainable_variables) # TODO
  # grads = tape.gradient('''TODO''', '''TODO''')

  # Apply the gradients to the optimizer so it can update the model accordingly
  optimizer.apply_gradients(zip(grads, model.trainable_variables))
  return loss

##################
# Begin training!#
##################

history = []
plotter = mdl.util.PeriodicPlotter(sec=2, xlabel='Iterations', ylabel='Loss')
if hasattr(tqdm, '_instances'): tqdm._instances.clear() # clear if it exists

for iter in tqdm(range(num_training_iterations)):

  # Grab a batch and propagate it through the network
  x_batch, y_batch = get_batch(vectorized_songs, seq_length, batch_size)
  loss = train_step(x_batch, y_batch)

  # Update the progress bar
  history.append(loss.numpy().mean())
  plotter.plot(history)

  # Update the model with the changed weights!
  if iter % 100 == 0:
    model.save_weights(checkpoint_prefix)

# Save the trained model and the weights
model.save_weights(checkpoint_prefix)

"""## 2.6 Generate music using the RNN model

Now, we can use our trained RNN model to generate some music! When generating music, we'll have to feed the model some sort of seed to get it started (because it can't predict anything without something to start with!).

Once we have a generated seed, we can then iteratively predict each successive character (remember, we are using the ABC representation for our music) using our trained RNN. More specifically, recall that our RNN outputs a `softmax` over possible successive characters. For inference, we iteratively sample from these distributions, and then use our samples to encode a generated song in the ABC format.

Then, all we have to do is write it to a file and listen!

### Restore the latest checkpoint

To keep this inference step simple, we will use a batch size of 1. Because of how the RNN state is passed from timestep to timestep, the model will only be able to accept a fixed batch size once it is built.

To run the model with a different `batch_size`, we'll need to rebuild the model and restore the weights from the latest checkpoint, i.e., the weights after the last checkpoint during training:
"""

'''TODO: Rebuild the model using a batch_size=1'''
model = build_model(vocab_size, embedding_dim, rnn_units, batch_size=1) # TODO
# model = build_model('''TODO''', '''TODO''', '''TODO''', batch_size=1)

# Restore the model weights for the last checkpoint after training
model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))
model.build(tf.TensorShape([1, None]))

model.summary()

"""Notice that we have fed in a fixed `batch_size` of 1 for inference.

### The prediction procedure

Now, we're ready to write the code to generate text in the ABC music format:

* Initialize a "seed" start string and the RNN state, and set the number of characters we want to generate.

* Use the start string and the RNN state to obtain the probability distribution over the next predicted character.

* Sample from multinomial distribution to calculate the index of the predicted character. This predicted character is then used as the next input to the model.

* At each time step, the updated RNN state is fed back into the model, so that it now has more context in making the next prediction. After predicting the next character, the updated RNN states are again fed back into the model, which is how it learns sequence dependencies in the data, as it gets more information from the previous predictions.

![LSTM inference](https://raw.githubusercontent.com/aamini/introtodeeplearning/2019/lab1/img/lstm_inference.png)

Complete and experiment with this code block (as well as some of the aspects of network definition and training!), and see how the model performs. How do songs generated after training with a small number of epochs compare to those generated after a longer duration of training?
"""

### Prediction of a generated song ###

def generate_text(model, start_string, generation_length=1000):
  # Evaluation step (generating ABC text using the learned RNN model)

  '''TODO: convert the start string to numbers (vectorize)'''
  input_eval = [char2idx[s] for s in start_string] # TODO
  # input_eval = ['''TODO''']
  input_eval = tf.expand_dims(input_eval, 0)

  # Empty string to store our results
  text_generated = []

  # Here batch size == 1
  model.reset_states()
  tqdm._instances.clear()

  for i in tqdm(range(generation_length)):
      '''TODO: evaluate the inputs and generate the next character predictions'''
      predictions = model(input_eval)
      # predictions = model('''TODO''')

      # Remove the batch dimension
      predictions = tf.squeeze(predictions, 0)

      '''TODO: use a multinomial distribution to sample'''
      predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()
      # predicted_id = tf.random.categorical('''TODO''', num_samples=1)[-1,0].numpy()

      # Pass the prediction along with the previous hidden state
      #   as the next inputs to the model
      input_eval = tf.expand_dims([predicted_id], 0)

      '''TODO: add the predicted character to the generated text!'''
      # Hint: consider what format the prediction is in vs. the output
      text_generated.append(idx2char[predicted_id]) # TODO
      # text_generated.append('''TODO''')

  return (start_string + ''.join(text_generated))

'''TODO: Use the model and the function defined above to generate ABC format text of length 1000!
    As you may notice, ABC files start with "X" - this may be a good start string.'''
generated_text = generate_text(model, start_string="X", generation_length=1000) # TODO
# generated_text = generate_text('''TODO''', start_string="X", generation_length=1000)

"""### Play back the generated music!

We can now call a function to convert the ABC format text to an audio file, and then play that back to check out our generated music! Try training longer if the resulting song is not long enough, or re-generating the song!
"""

### Play back generated songs ###

generated_songs = mdl.lab1.extract_song_snippet(generated_text)

for i, song in enumerate(generated_songs):
  # Synthesize the waveform from a song
  waveform = mdl.lab1.play_song(song)

  # If its a valid song (correct syntax), lets play it!
  if waveform:
    print("Generated song", i)
    ipythondisplay.display(waveform)

"""## 2.7 Experiment and **get awarded for the best songs**!

Congrats on making your first sequence model in TensorFlow! It's a pretty big accomplishment, and hopefully you have some sweet tunes to show for it.

Consider how you may improve your model and what seems to be most important in terms of performance. Here are some ideas to get you started:

*  How does the number of training epochs affect the performance?
*  What if you alter or augment the dataset?
*  Does the choice of start string significantly affect the result?

Have fun and happy listening!

![Let's Dance!](http://33.media.tumblr.com/3d223954ad0a77f4e98a7b87136aa395/tumblr_nlct5lFVbF1qhu7oio1_500.gif)

"""


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Test solutions to chapter 12 exercises.

###############################################################################
# test_chapter12.py
#
# Revision:     1.00
# Date:         7/12/2021
# Author:       Alex
#
# Purpose:      Runs unit tests on all chapter 12 exercises from "Data
#               Structures and Algorithms in Python" by Goodrich et. al.
#
###############################################################################
"""

# %% Imports
# Standard system imports

# Related third party imports
import pytest
import numpy as np

# Local application/library specific imports
import dsa.chapter12_exercises as chap12

# %% Reinforcement Exercises
def test_prop_12p1():
    """Solution to exercise R-12.1.

    Give a complete justification of Proposition 12.1.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    Proposition 12.1 asserts the following:

    The merge-sort tree associated with an execution of merge-sort on a
    sequence of size n has height ceil(logn).

    The merge-sort tree is a binary tree where each node represents a recursive
    call of the merge-sort algorithm.  Each time merge-sort is called, it
    splits the input sequence into two (roughly) equal halves until the input
    sequence is of length 1.  The original sequence S is the root node of the
    binary tree, while the individual elements of S form the leaves of the
    binary tree.

    If the length of the sequence is not a power of 2, then there will be
    splits where one half will equal n//2, and the other half n//2 + 1.
    The length of the sequence at each node is thus either 2^h or 2^h + 1,
    where h is the height of the node.  This means that not all of the leaf
    nodes will be at the same depth, as an extra split operation is necessary
    for nodes containing n//2 + 1 elements.

    Leaf nodes have height 0, and 2^0 = 1 so the sequence will be length of 1.
    This matches our expectation that leaves contain a single element.
    The root node containing the original sequence must be length n, and the
    height of the root node (h_root) is the same as the height of the tree.

    Because n may not be a power of 2, 2^h_root may be greater than or equal to
    the number of elements in sequence S.  However, height must be an integer
    and so any decimal result for h_root should thus take the ceiling of the
    value.

    Therefore:

    2^(h_root) >= n
    h_root >= log(n)
    h_root = ceil(log(n))
    """
    assert chap12.prop_12p1()

def test_merge_sort_arrows():
    """Solution to exercise R-12.2.

    In the merge-sort tree shown in Figures 12.2 through 12.4, some edges are
    drawn as arrows. What is the meaning of a downward arrow? How about
    an upward arrow?

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    A downward arrow represents a recursive call to merge_sort(), splitting the
    sequence into two halves.  An upward arrow represents a call to merge(),
    which merges two sorted sequences back into the parent node that represents
    the original sequence they were split from.
    """
    assert chap12.merge_sort_arrows()

def test_merge_sort_stable():
    """Solution to exercise R-12.4.

    Is our array-based implementation of merge-sort given in Section 12.2.2
    stable? Explain why or why not.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    Let S1 = [2, 3, 5] and S2 = [2, 4, 6].  Both are sorted sequences that
    contain the duplicate key 2.  Recall that S1 is composed of keys from the
    left half of S, and S2 of keys from the right half of S.  If a duplicate
    key exists in both S1 and S2, the duplicate key in S1 must have preceded
    the key in S2.  In order for a sort to be stable, this order must be
    maintained in the sorted list.

    When viewing the source code of merge() on page 543 of the text, the
    comparison between S1 and S2 is written using the < operator:

    S1[i] < S2[j]

    If True, element i from S1 is copied to S.  If False, element j from S2 is
    copied to S instead.  If i = j = 0 then:

    S1[0] < S2[0]
    2 < 2

    This statement is False, 2 is not less than itself.  Therefore S2[0] will
    be copied to S before S1[0].  As mentioned above, duplicates from S1 should
    be copied to S first because they preceded the duplicates in S2 in the
    original sequence.  Therefore this implementation of merge-sort is NOT
    stable.  However, if the comparison operator were changed to <= the merge-
    sort algorithm would be stable.
    """
    assert chap12.merge_sort_stable()

def test_linked_list_stable():
    """Solution to exercise R-12.5.

    Is our linked-list-based implementation of merge-sort given in Code
    Fragment 12.3 stable? Explain why or why not.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    No, it is NOT stable for the same reason as exercise R-12.4.  The merge()
    function again compares S1 to S2 using the < operator:

    S1.first() < S2.first()

    If True, S1's element is copied to S, and if False S2's element is copied
    to S.  Again, in the case that S1.first() == S2.first() then S2's element
    will be copied to S before S1's element.  This will not maintain the
    original order of the duplicate keys in the sorted sequence.
    """
    assert chap12.linked_list_stable()

@pytest.mark.parametrize(('A, B'), [(list(range(10)),
                                     list(range(5, 15))),
                                    (list(range(0, 100, 10)),
                                     list(range(40, 60, 2))),
                                    (list(range(0, 10)),
                                     list(range(10, 20))),
                                    (list(range(0, 20, 2)),
                                    list(range(1, 20, 2)))
                                    ])
def test_sequence_union(A, B):
    """Solution to exercise R-12.7.

    Suppose we are given two n-element sorted sequences A and B each with
    distinct elements, but potentially some elements that are in both
    sequences.  Describe an O(n)-time method for computing a sequence
    representing the union A ∪ B (with no duplicates) as a sorted sequence.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    Important information:
    1. A and B are both length n
    2. A and B are both sorted
    3. The elements of A and B are distinct, meaning no repeats within A or B

    Because the lists are sorted, we don't need to compare every element in
    A to every element in B.  We know that if A[i] < B[j] that we can increment
    i until A[i] >= B[j].  Conversely, if A[i] > B[j] we should increment j
    until A[i] <= B[j].  If a match is found, we append it to the union list
    and increment both i and j.  Because both A and B are distinct there is
    no possibility of getting multiple matches of the same value.

    The while loop will continue so long as both i and j are < n.  This means
    that the maximum number of loops is 2n.  Each loop performs an O(1)
    operation (append is O(1)* amortized), and so the algorithim's run-time
    efficiency is proportional to the number of loops, which is O(n).
    """
    assert chap12.sequence_union(A, B) == sorted(list(set(A) & set(B)))

def test_quick_sort_pivot1():
    """Solution to exercise R-12.8.

    Suppose we modify the deterministic version of the quick-sort algorithm
    so that, instead of selecting the last element in an n-element sequence as
    the pivot, we choose the element at index n//2. What is the running time
    of this version of quick-sort on a sequence that is already sorted?

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    If the pivot is always selected at n // 2, and the sequence is already
    sorted, then the pivot will be the median value of the sequence and split
    it evenly into two halves at every recursive call of quick_sort().
    This is the ideal scenario as the resulting quick-sort binary tree will
    have a height of log(n), and so the quick-sort algorithm's run-time will be
    O(nlogn).
    """
    assert chap12.quick_sort_pivot1()

def test_quick_sort_pivot2():
    """Solution to exercise R-12.9.

    Consider a modification of the deterministic version of the quick-sort
    algorithm where we choose the element at index n//2 as our pivot.
    Describe the kind of sequence that would cause this version of quick-sort
    to run in Ω(n^2) time.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    The sequence would need to ensure that the pivot value selected at index
    n // 2 was always either the smallest or the largest value in the sequence.
    This ensures that all values are moved to either L or G for every split,
    and so the sequence is only decreasing in length by 1 (the pivot value).
    This would cause the height of the tree to be n, and so the quick-sort
    algorithm's run-time would be Ω(n^2).

    To create this sequence, simply sort the values in reverse order, split
    the sequence in two, and then swap the two halves.  An example is shown
    below:

    [0 1 2 3 4 5 6 7 8 9]       # Example sequence
    [9 8 7 6 5 4 3 2 1 0]       # Reverse sort
    [4 3 2 1 0 9 8 7 6 5]       # Swap halves

    Now to demonstrate that the length of the sequence will only decrease by 1
    for each quick-sort call.  Remember, if the pivot value is the smallest or
    greatest value, all of the other values are placed into either L or G:

    [4 3 2 1 0 9 8 7 6 5]       # n = 10, n//2 = 5, pivot = 9, all values in L
    [4 3 2 1 0 8 7 6 5]         # n =  9, n//2 = 4, pivot = 0, all values in G
    [4 3 2 1 8 7 6 5]           # n =  8, n//2 = 4, pivot = 8, all values in L
    [4 3 2 1 7 6 5]             # n =  7, n//2 = 3, pivot = 1, all values in G
    [4 3 2 7 6 5]               # n =  6, n//2 = 3, pivot = 7, all values in L
    [4 3 2 6 5]                 # n =  5, n//2 = 2, pivot = 2, all values in G
    [4 3 6 5]                   # n =  4, n//2 = 2, pivot = 6, all values in L
    [4 3 5]                     # n =  3, n//2 = 1, pivot = 3, all values in G
    [4 5]                       # n =  2, n//2 = 1, pivot = 5, all values in L
    [4]                         # n <  2, Return
    """
    assert chap12.quick_sort_pivot2()

def test_inplace_flaw():
    """Solution to exercise R-12.12.

    If the outermost while loop of our implementation of inplace quick sort
    (line 7 of Code Fragment 12.6) were changed to use condition left < right
    (rather than left <= right), there would be a flaw. Explain the flaw and
    give a specific input sequence on which such an implementation fails.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    This change would prevent the left index from exceeding the right index.
    The value at the left index and the pivot are swapped at the end of the
    quick-sort call.  This means that the wrong value will be swapped with the
    pivot.

    This is an issue whenever the value that the left and right indices meet at
    is less than the pivot value.  This value will end up on the right side of
    the sequence when it is swapped with the pivot (assuming we use the last
    element of the sequence as the pivot), when it should be on the left side
    of the sequence.

    The sequence used in Figure 12.14 works as an example.  The beginning
    sequence is [85, 24, 63, 45, 17, 31, 96, 50].  After the inplace quick-sort
    with the < operator, the result is [17, 31, 24, 50, 45, 63, 85, 96].

    [85, 24, 63, 45, 17, 31, 96, 50]  # Original sequence
                 l,r              p

    [31, 24, 17, 50, 63, 85, 96, 45]  # 50 at index 3
                  p

    [31, 24, 17]    [63, 85, 96, 45]
     l,r      p      l,r          p

    [17, 24, 31]    [45, 85, 96, 63]  # 17 at index 0, 45 at index 4
      p               p

        [24, 31]        [85, 96, 63]
         l,r  p          l,r      p

        [31, 24]        [63, 96, 85]  # 31 at index 1, 63 at index 5
          p               p

            [24]            [96, 85]  # 24 at index 2
              p              l,r  p

                            [85, 96]  # 85 at index 6
                              p

                                [96]  # 96 at index 7
                                  p

    [17, 31, 24, 50, 45, 63, 85, 96]

    The final sorted sequence is clearly incorrect.
    """
    assert chap12.inplace_flaw()

def test_bucket_sort_inplace():
    """Solution to exercise R-12.17.

    Is the bucket-sort algorithm in-place? Why or why not?

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    An algorithm is in-place if uses only a small amount of memory in addition
    to that needed for the original input.  Bucket-sort uses a separate bucket
    array to store entries from the original sequence S.  The bucket array's
    memory usage must be at least as large as S in order to store all of S's
    entries.  The entries are then placed back into S in sorted order.

    Therefore bucket-sort requires O(n) additional memory usage and is *not*
    in-place.
    """
    assert chap12.bucket_sort_inplace()

@pytest.fixture(name="sequences", scope="function", params=[1, 2, 3, 5, 7])
def d_tuples_fixture(request):
    """Fixture to supply sequences of d-tuples for the radix sort function."""

    class GenerateDTuples:
        """Fixture class to store sequences of d-tuples."""

        def __init__(self):
            """Reproducible sequences of d-tuples to test radix sort."""
            rng = np.random.default_rng(4)
            self.N = 100
            N = self.N              # Range of integer keys is [0, N-1]
            d = request.param       # d-tuple, d = 3 is a triplet or 3-tuple
            n = 10                  # Length of sequence
            s = 10                  # Number of sequences to test
            self.d_tuples = (
                [tuple(rng.choice(range(N), size=d)) for _ in range(n)]
                for _ in range(s)
            )
    return GenerateDTuples()

def test_radix_sort_triplets(sequences):
    """Solution to exercise R-12.18.

    Describe a radix-sort method for lexicographically sorting a sequence S of
    triplets (k, l, m), where k, l, and m are integers in the range [0, N − 1],
    for some N ≥ 2. How could this scheme be extended to sequences of d-tuples
    (k1, k2 , ..., kd), where each ki is an integer in the range [0, N − 1]?

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    The radix-sort of a sequence of d-tuples is simply the application of d
    bucket-sorts of the sequence in reverse order of keys.  In a lexicographic
    sort the dth key is the least important, and the first key the most
    important.  By stably sorting by the dth key first, then the d-1 key, d-2
    key, etc. we ensure a lexicographic ordering of the d-tuples.

    The solution below lexicographically sorts a sequence of d-tuples for any
    value of d > 0.  Note that I used a list of lists to represent the bucket
    array.  Popping and appending at the end of a list are O(1) operations,
    while popping and inserting at the beginning of a list are O(n) operations.

    Because of this, I iterate through the sequence of d-tuples in reverse
    order so that I can pop the last d-tuple from the sequence and append it to
    the end of the bucket at the appropriate index.  This inserts the d-tuples
    in the bucket in reverse order, such that d-tuples that appeared later in
    the sequence are stored in the beginning of the bucket.

    Once the sequence is empty, I iterate through the bucket array in order
    from index 0 to N-1.  I again pop d-tuples from the end of each bucket and
    append them to the end of the sequence.  The d-tuples at the end of the
    bucket are the d-tuples that appeared earliest in the sequence, and so
    this preserves the order of the entries in the sequence (stable) while
    allowing for O(1) array operations.

    This is the reverse of the front-to-back method discussed in the text,
    where the sequence is accessed front-to-back and elements are removed from
    the buckets front-to-back.  Because I operate back-to-front on both the
    sequence and the buckets the reversal cancels out, and the result is a
    stable sort.

    This series of operations repeats for all d keys in the tuple.
    """
    for sequence in sequences.d_tuples:
        expected_sort = sorted(sequence)  # Stable, lexicographic sort
        chap12.radix_sort_triplets(sequence, sequences.N)
        assert sequence == expected_sort

def test_how_long_quick_merge_sorts():
    """Solution to exercise R-12.19.

    Suppose S is a sequence of n values, each equal to 0 or 1. How long will
    it take to sort S with the merge-sort algorithm? What about quick-sort?

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    The exercise does not state whether the 0s and 1s are equally likely to
    appear in the sequence, but either way there will be many duplicate values.
    Merge-sort has a worst-case run time of O(nlogn), regardless of how many
    duplicates exists in the sequence.  It also cannot take advantage of "runs"
    in the data, and so I would expect merge-sort to require O(nlogn) time.

    As for quick-sort, each recursive call of quick-sort chooses
    a value from the sequence as a pivot, and splits the sequence into
    sub-sequences of values that are either greater than, less than, or equal
    to the pivot value.  But with the elements of the sequence restricted to
    two values, the pivot can never be the median value of the sequence.
    A pivot value that is consistently close to the median of the sequence is
    typically the ideal case to achieve O(nlogn) performance.

    Instead, the pivot will always be either 0 or 1, and the values in the
    sequence will either all be stored in L and E, or all stored in G and E.
    There are two scenarios that could ensue:

    1. The distribution of 0s and 1s in the sequence are roughly equal
    2. The distribution of 0s and 1s is skewed towards either 0 or 1

    If the sequence is composed of a roughly equal number of 0s and 1s, then
    each recursive call of quick-sort will split the sequence roughly in half:
    half of the elements in E, and the other half in (one of) L or G.  This
    halving of the length of the sequence will result in a quick-sort binary
    tree of height log(n), and so the run-time efficiency of quick-sort would
    thus be O(nlogn).

    If the distribution is skewed towards mostly 0s or mostly 1s, then each
    recursive call of quick-sort will place the majority of values in either E
    or (one of) L or G.  If the sequence is mostly 0s (for instance), most of
    the time the pivot will be a zero and most of the values will go to E.
    Very few values will go to G, and so the size of the sequence will be
    quickly reduced.  In the extreme, if the entire sequence was 0s with no 1s,
    then all of the value would be placed in E and quick-sort would return
    after a single call as L and G are empty.  This corresponds to a
    quick-sort binary tree of height 1, and so the expected run-time is O(n).

    As a final answer, under this scenario we would expect quick-sort's
    run-time efficiency to be bounded between Ω(n) and O(nlogn).
    """
    assert chap12.how_long_quick_merge_sorts()

def test_how_long_bucket_sort():
    """Solution to exercise R-12.20.

    Suppose S is a sequence of n values, each equal to 0 or 1. How long will
    it take to sort S stably with the bucket-sort algorithm?

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    In this scenario N = 2, and all of the values will be inserted into either
    index 0 or index 1 of the bucket array.  Bucket-sort will stably sort the
    sequence regardless of the number of duplicates.

    The run-time efficiency of bucket-sort is O(n + N), and if N is small
    relative to n bucket-sort is O(n).

    A sequence of n < 2 is already sorted, and so it's safe to say that in this
    scenario bucket-sort's run-time efficiency is O(n).
    """
    assert chap12.how_long_bucket_sort()

def test_inplace_sort():
    """Solution to exercise R-12.21.

    Given a sequence S of n values, each equal to 0 or 1, describe an in-place
    method for sorting S.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    I used a solution similar to the in-place quick-sort algorithm, except
    that the pivot is always be the same value and there is no need for
    recursion.
    """
    n = 100         # Length of sequence
    num_tests = 10  # Number of random sequences to test
    rng = np.random.default_rng(27)
    for _ in range(num_tests):
        rints = list(rng.integers(low=0, high=2, size=n))
        expected_sort = sorted(rints)
        chap12.inplace_sort(rints)
        assert rints == expected_sort

def test_insertion_sort():
    """Solution to exercise R-12.22.

    Give an example input list that requires merge-sort and heap-sort to take
    O(n log n) time to sort, but insertion-sort runs in O(n) time. What if you
    reverse this list?

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    Insertion-sort runs in O(n) time when the sequence is already sorted, such
    as [0, 1, 2, 3, 4, 5].  This is because the inner loop of the insertion-
    sort never needs to run, and the outer loop runs for n iterations.

    Merge-sort is guaranteed to run in worst-case O(nlogn) time regardless of
    whether the list is sorted.  The fact that the list is already sorted does
    not improve its run-time, because it will recursively split the sequence
    in halves even if it is already sorted.

    Heap-sort's run-time efficiency is tied to the height of the heap.  As
    items are added to the heap they will have to be "up-heap bubbled" in order
    to preserve the heap-order and complete binary tree properties.  This
    requires O(nlogn) time as each added item i requires O(logi) time to move
    up the tree.  However, if a heap is constructed from a sorted list each
    added element should preserve the heap-order and completeness properties.
    This will reduce the phase 1 run-time to O(n).  But in order to sort the
    list, the remove_min() method must be called n times.  Each time
    remove_min() is called, the root of the heap is replaced by a value from
    the bottom of the heap.  This value must then be "down-heap bubbled" which
    will require time proportional to the height of the tree.  Thus the phase 2
    run-time of heap-sort is O(nlogn), making the overall run-time efficiency
    of heap-sort O(nlogn) even for an already sorted list.

    If the sorted list is reversed, then insertion-sort runs in O(n^2) time
    because the inner loop must perform 0, 1, 2, ..., n-1 swaps.  This is the
    sum of the first n integers, hence the O(n^2) run-time.  Merge-sort and
    heap-sort will still sort the reversed list in O(nlogn) time.
    """
    assert chap12.insertion_sort()

def test_best_algorithm():
    """Solution to exercise R-12.23.

    What is the best algorithm for sorting each of the following: general com-
    parable objects, long character strings, 32-bit integers, double-precision
    floating-point numbers, and bytes? Justify your answer.

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    1. General comparable objects:
        a. If the sequence is small-to-medium-sized, heap-sort is guaranteed to
           be O(nlogn) worst-case.

        b. If the sequence is large and memory is not an issue, merge-sort is
           worst-case O(nlogn) and may outperform heap-sort.

        c. If the sequence is large and memory usage is an issue, quick-sort
           can easily run in-place and may outperform both heap-sort and
           quick-sort with its expected O(nlogn) run-time efficiency.  However,
           worst-case it is O(n^2), and this may not be acceptable for real-
           time applications.

        d. Insertion-sort, radix-sort, bucket-sort:
           Bucket-sort and radix-sort require sequences of integers, and so in
           general can't be recommended.  Insertion-sort performs well on small
           sequences - especially if they are already nearly sorted.  Again,
           this is a special case and in general insertion-sort can't be
           recommended.

    2. Long character strings:
        a. Radix-sort is ideal for character strings, because it will perform a
           lexicographic sort of the strings in O(d(n + N)) time.  If the
           strings are long d might be quite large, but even so other sorting
           methods will not lexicographically sort strings without
           modification.

        b. Quick-sort and heap-sort are not stable sorting methods, and so
           cannot lexicographically sort a string.

        c. Merge-sort is a stable sorting algorithm, and so if d is large
           enough it may be worth modifying merge-sort to handle strings.

    3. 32-bit integers:
        a. If the integers happen to be constrained to a range N that is O(n),
           then the best sorting algorithm is bucket-sort, because bucket-sort
           can sort integers in O(n) time.

        b. If the integers are not constrained to a range N that is O(n) and
           are quite large compared to n, then another sorting method will
           provide better results.  In this case, the same recommendations
           for general comparable objects will apply.

    4. Double-precision floating-point numbers
       a. The same recommendations for general comparable objects apply here.

    5. Bytes
        a. It's possible to represent a byte as a d-tuple of d bits, each bit
           being either 0 or 1.  In this case N = 2, and so radix-sort should
           run in O(n) time.

        b. Bytes can be represented as integers, and so a bucket-sort may be
           appropriate if N is constrained to be O(n).
    """
    assert chap12.best_algorithm()

def test_quick_select():
    """Solution to exercise R-12.24.

    Show that the worst-case running time of quick-select on an n-element
    sequence is Ω(n^2).

    ---------------------------------------------------------------------------
    Solution:
    ---------------------------------------------------------------------------
    In the worst case, the randomly selected pivot value will always be the
    maximum value of the sequence when searching for the minimum value,
    or the pivot will always be the minimum value of the sequence when
    searching for the maximum value.  This will result in all of the
    sequence values being placed in either the L or G list, and the length of
    the sequence will only reduce by one per recursive call to quick-select.
    The algorithm will then run n, n-1, n-2, n-3, ..., 1 comparisons,
    which is the sum of the first n integers and is equal to n(n+1)/2.  In this
    worst-case scenario the lower bound on run-time efficiency is thus Ω(n^2).
    """
    assert chap12.quick_select()


from yolo import *
import os, time
from PIL import Image, ImageDraw, ImageFile
import math
import cv2
import numpy as np
import subprocess

ImageFile.LOAD_TRUNCATED_IMAGES = True

WAIT_AFTER_JUMP = 3000
CAPTURE_FOLDER = '/home/cooli7wa/Desktop/tmp_image/'
CAPTURE_FILE = CAPTURE_FOLDER + 'screen.png'
IMG_PATH = '/home/cooli7wa/project/pycharm/tiaotiao/img/'
PRESS_PARAM = 1.375
CHESS_CENTER_CORRECT = 22
INVALID_BOX_DISTANCE = 20
RESTORE_NUM = 50
RESTART_GAME_POS = [600, 1700]

screen_param = {'width': 0, 'height': 0}
restore_counter = 0

def image_open(path):
    img = Image.open(path)
    return img

def get_screen_parameter():
    screen_capture_download()
    img = image_open(CAPTURE_FILE)
    shape = img.size
    screen_param['width'] = shape[0]
    screen_param['height'] = shape[1]
    print('screen parameter:', screen_param)

def screen_capture_download():
    process = subprocess.Popen('adb shell screencap -p', shell=True, stdout=subprocess.PIPE)
    screenshot = process.stdout.read()
    screenshot = screenshot.replace(b'\r\n', b'\n')
    f = open('{}'.format(CAPTURE_FILE), 'wb')
    f.write(screenshot)
    f.close()

def jump(distance):
    press_time = round(distance * PRESS_PARAM)
    cmd = 'adb shell input swipe {x1} {y1} {x2} {y2} {duration}'.format(
        x1=screen_param['width'] // 2,
        y1=screen_param['height'] // 2,
        x2=screen_param['width'] // 2,
        y2=screen_param['height'] // 2,
        duration=press_time
    )
    print('distance: {}, press_time: {}ms'.format(distance, press_time))
    os.popen(cmd)

def cal_center_point(box):
    x = (box[2] + box[0]) / 2
    y = (box[3] + box[1]) / 2
    return [x, y]

def is_invalid_box(chess_center, box_center):
    if box_center[1] >= chess_center[1]:
        return True
    if cal_distance(box_center, chess_center) <= INVALID_BOX_DISTANCE:
        return True
    return False

def cal_distance(point1, point2):
    x = point2[0] - point1[0]
    y = point2[1] - point1[1]
    return math.sqrt((x ** 2) + (y ** 2))

def draw_points(image, points):
    draw = ImageDraw.Draw(image)
    for point in points:
        draw.ellipse((point[0] - 10, point[1] - 10, point[0] + 10, point[1] + 10), 'seagreen', 'skyblue')
    return image

def change_box_xy(boxes):
    for i, box in enumerate(boxes):
        boxes[i] = [box[j] for j in [1, 0, 3, 2]]

def restore_image(img_a, img):
    global restore_counter
    cv2.imwrite('{}/{}_a.jpg'.format(CAPTURE_FOLDER, restore_counter), img_a)
    cv2.imwrite('{}/{}.jpg'.format(CAPTURE_FOLDER, restore_counter), img)
    restore_counter = (restore_counter + 1) % RESTORE_NUM

def add_mismatch_image(image):
    jpg_list = list(filter(lambda i: i.rsplit('.')[1] == 'jpg', os.listdir(IMG_PATH)))
    next_number = int(sorted(jpg_list)[-1].rsplit('.')[0]) + 1
    cv2.imwrite(IMG_PATH + '%06d.jpg' % next_number, image)

def restart_game():
    jump(2000)
    time.sleep(8)
    os.popen('adb shell input tap {} {}'.format(RESTART_GAME_POS[0], RESTART_GAME_POS[1]))
    time.sleep(2)

def main_loop(model):
    while True:
        screen_capture_download()
        image_origin = image_open(CAPTURE_FILE)
        image = image_origin.copy()
        boxes, classes = model.detect_image(image)
        image_origin = np.array(image_origin)
        image_origin = cv2.cvtColor(image_origin, cv2.COLOR_RGB2BGR)
        boxes = boxes.tolist()
        classes = classes.tolist()
        change_box_xy(boxes)
        target_box, chessman_box = [], []

        for i in range(len(boxes)):
            if classes[i] == 2:
                chessman_box = boxes[i]
            else:
                if not target_box:
                    target_box = boxes[i]
                else:
                    center_cur = cal_center_point(boxes[i])
                    center_tar = cal_center_point(target_box)
                    if center_cur[1] < center_tar[1]:
                        target_box = boxes[i]
        if not chessman_box or not target_box:
            print('no chessman or target box')
            add_mismatch_image(image_origin)
            restart_game()
            continue
        chessman_center = [(chessman_box[2] + chessman_box[0]) / 2, chessman_box[3] - CHESS_CENTER_CORRECT]
        target_center = cal_center_point(target_box)
        image = draw_points(image, [chessman_center, target_center])
        image = image.resize((screen_param['width'] // 3, screen_param['height'] // 3), Image.ANTIALIAS)
        image = np.array(image)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        cv2.imshow('', image)
        if is_invalid_box(chessman_center, target_center):
            print('no possible target box')
            add_mismatch_image(image_origin)
            restart_game()
            continue
        restore_image(image, image_origin)
        distance = cal_distance(target_center, chessman_center)
        jump(distance)
        cv2.waitKey(WAIT_AFTER_JUMP)

if __name__ == '__main__':
    get_screen_parameter()
    restore_counter = 0
    main_loop(YOLO())


import numpy as np
import argparse
import os
import json

from sklearn.metrics.pairwise import cosine_distances, euclidean_distances

def main():
    parser = argparse.ArgumentParser('retrieval eval')
    parser.add_argument('--output-dir', type=str)
    parser.add_argument('--trainsplit', type=str, required=True)
    parser.add_argument('--valsplit', type=str, required=True)
    parser.add_argument('--num_replica', type=int, default=8)
    parser.add_argument('--data-source', type=str)
    args = parser.parse_args()

    for i in range(args.num_replica):
        os.path.exists(os.path.join(args.output_dir, 'feature_{}_{}.npy'.format(args.trainsplit, i)))
        os.path.exists(os.path.join(args.output_dir, 'feature_{}_cls_{}.npy'.format(args.trainsplit, i)))
        os.path.exists(os.path.join(args.output_dir, 'feature_{}_{}.npy'.format(args.valsplit, i)))
        os.path.exists(os.path.join(args.output_dir, 'feature_{}_cls_{}.npy'.format(args.valsplit, i)))
        os.path.exists(os.path.join(args.output_dir, 'vid_num_{}.npy'.format(args.trainsplit)))
        os.path.exists(os.path.join(args.output_dir, 'vid_num_{}.npy'.format(args.valsplit)))

    vid_num_train = np.load(os.path.join(args.output_dir, 'vid_num_{}.npy'.format(args.trainsplit)))
    train_padding_num = vid_num_train[0] % args.num_replica
    vid_num_val = np.load(os.path.join(args.output_dir, 'vid_num_{}.npy'.format(args.valsplit)))
    val_padding_num = vid_num_val[0] % args.num_replica

    feat_train = []
    feat_train_cls = []
    for i in range(args.num_replica):
        feat_train.append(np.load(os.path.join(args.output_dir, 'feature_{}_{}.npy'.format(args.trainsplit, i))))
        feat_train_cls.append(
            np.load(os.path.join(args.output_dir, 'feature_{}_cls_{}.npy'.format(args.trainsplit, i))))
    if train_padding_num > 0:
        for i in range(train_padding_num, args.num_replica):
            feat_train[i] = feat_train[i][:-1, :]
            feat_train_cls[i] = feat_train_cls[i][:-1]
    feat_train = np.concatenate(feat_train, axis=0).squeeze()
    feat_train_cls = np.concatenate(feat_train_cls, axis=0).squeeze()
    print('feat_train: {}'.format(feat_train.shape))
    print('feat_train_cls: {}'.format(feat_train_cls.shape))

    feat_val = []
    feat_val_cls = []
    for i in range(args.num_replica):
        feat_val.append(np.load(os.path.join(args.output_dir, 'feature_{}_{}.npy'.format(args.valsplit, i))))
        feat_val_cls.append(np.load(os.path.join(args.output_dir, 'feature_{}_cls_{}.npy'.format(args.valsplit, i))))
    if val_padding_num > 0:
        for i in range(val_padding_num, args.num_replica):
            feat_val[i] = feat_val[i][:-1, :]
            feat_val_cls[i] = feat_val_cls[i][:-1]
    feat_val = np.concatenate(feat_val, axis=0)
    feat_val_cls = np.concatenate(feat_val_cls, axis=0)
    print('feat_val: {}'.format(feat_val.shape))
    print('feat_val_cls: {}'.format(feat_val_cls.shape))

    # kNN retrieval
    if args.valsplit == 'test':
        ks = [3]
    else:
        ks = [1, 5, 10, 20, 50]
    topk_correct = {k: 0 for k in ks}

    class_top = 1
    if args.data_source == 'ucf':
        class_num = 101
    elif args.data_source == 'hmdb':
        class_num = 51
    else:
        raise Exception('The data-source argument no assigned!')
    class_correct = {cls: 0 for cls in range(0, class_num)}
    class_total = {cls: 0 for cls in range(0, class_num)}

    X_train = feat_train
    y_train = feat_train_cls
    X_test = feat_val
    y_test = feat_val_cls

    distances = cosine_distances(X_test, X_train)
    indices = np.argsort(distances)

    for k in ks:
        # print(k)
        top_k_indices = indices[:, :k]
        if args.valsplit == 'test':
            print(top_k_indices)
            np.save(os.path.join(args.output_dir, 'top_k_indices.npy'), top_k_indices)
        # print(top_k_indices.shape, y_test.shape)
        for ind, test_label in zip(top_k_indices, y_test):
            labels = y_train[ind]
            if test_label in labels:
                # print(test_label, labels)
                topk_correct[k] += 1
                if k == class_top:
                    class_correct[test_label] += 1
            if k == class_top:
                class_total[test_label] += 1

    for k in ks:
        correct = topk_correct[k]
        total = len(X_test)
        print('Top-{}, correct = {:.2f}, total = {}, acc = {:.3f}'.format(k, correct, total, correct / total))

    # save label
    if args.valsplit != 'test':
        label_file = os.path.join(args.output_dir, 'class_retrieval_vclr.txt')
        f = open(label_file, 'w')
        for k in class_correct.keys():
            correct = class_correct[k]
            total = class_total[k]
            info = 'Classs-{}, Top-{}, correct = {:.2f}, total = {}, acc = {:.3f}'.format(
                k, class_top, correct, total, correct / total)
            print(info)
            f.write(info)
            f.write('\n')
        f.close()

    with open(os.path.join(args.output_dir, 'topk_correct.json'), 'w') as fp:
        json.dump(topk_correct, fp)

if __name__ == '__main__':
    main()