#%% [markdown]
# Lambda School Data Science, Unit 2: Predictive Modeling
#
# # Kaggle Challenge, Module 1
#
# ## Assignment
# - [ ] Do train/validate/test split with the Tanzania Waterpumps data.
# - [ ] Define a function to wrangle train, validate, and test sets in the same way. Clean outliers and engineer features. (For example, [what other columns have zeros and shouldn't?](https://github.com/Quartz/bad-data-guide#zeros-replace-missing-values) What other columns are duplicates, or nearly duplicates? Can you extract the year from date_recorded? Can you engineer new features, such as the number of years from waterpump construction to waterpump inspection?)
# - [ ] Select features. Use a scikit-learn pipeline to encode categoricals, impute missing values, and fit a decision tree classifier.
# - [ ] Get your validation accuracy score.
# - [ ] Get and plot your feature importances.
# - [ ] Submit your predictions to our Kaggle competition. (Go to our Kaggle InClass competition webpage. Use the blue **Submit Predictions** button to upload your CSV file. Or you can use the Kaggle API to submit your predictions.)
# - [ ] Commit your notebook to your fork of the GitHub repo.
#
#
# ## Stretch Goals
#
# ### Reading
#
# - A Visual Introduction to Machine Learning
#   - [Part 1: A Decision Tree](http://www.r2d3.us/visual-intro-to-machine-learning-part-1/)
#   - [Part 2: Bias and Variance](http://www.r2d3.us/visual-intro-to-machine-learning-part-2/)
# - [Decision Trees: Advantages & Disadvantages](https://christophm.github.io/interpretable-ml-book/tree.html#advantages-2)
# - [How a Russian mathematician constructed a decision tree — by hand — to solve a medical problem](http://fastml.com/how-a-russian-mathematician-constructed-a-decision-tree-by-hand-to-solve-a-medical-problem/)
# - [How decision trees work](https://brohrer.github.io/how_decision_trees_work.html)
# - [Let’s Write a Decision Tree Classifier from Scratch](https://www.youtube.com/watch?v=LDRbO9a6XPU) — _Don’t worry about understanding the code, just get introduced to the concepts. This 10 minute video has excellent diagrams and explanations._
# - [Random Forests for Complete Beginners: The definitive guide to Random Forests and Decision Trees](https://victorzhou.com/blog/intro-to-random-forests/)
#
#
# ### Doing
# - [ ] Add your own stretch goal(s) !
# - [ ] Try other [scikit-learn imputers](https://scikit-learn.org/stable/modules/impute.html).
# - [ ] Make exploratory visualizations and share on Slack.
#
#
# #### Exploratory visualizations
#
# Visualize the relationships between feature(s) and target. I recommend you do this with your training set, after splitting your data.
#
# For this problem, you may want to create a new column to represent the target as a number, 0 or 1. For example:
#
# ```python
# train['functional'] = (train['status_group']=='functional').astype(int)
# ```
#
#
#
# You can try [Seaborn "Categorical estimate" plots](https://seaborn.pydata.org/tutorial/categorical.html) for features with reasonably few unique values. (With too many unique values, the plot is unreadable.)
#
# - Categorical features. (If there are too many unique values, you can replace less frequent values with "OTHER.")
# - Numeric features. (If there are too many unique values, you can [bin with pandas cut / qcut functions](https://pandas.pydata.org/pandas-docs/stable/getting_started/basics.html?highlight=qcut#discretization-and-quantiling).)
#
# You can try [Seaborn linear model plots](https://seaborn.pydata.org/tutorial/regression.html) with numeric features. For this classification problem, you may want to use the parameter `logistic=True`, but it can be slow.
#
# You do _not_ need to use Seaborn, but it's nice because it includes confidence intervals to visualize uncertainty.
#
# #### High-cardinality categoricals
#
# This code from a previous assignment demonstrates how to replace less frequent values with 'OTHER'
#
# ```python
# # Reduce cardinality for NEIGHBORHOOD feature ...
#
# # Get a list of the top 10 neighborhoods
# top10 = train['NEIGHBORHOOD'].value_counts()[:10].index
#
# # At locations where the neighborhood is NOT in the top 10,
# # replace the neighborhood with 'OTHER'
# train.loc[~train['NEIGHBORHOOD'].isin(top10), 'NEIGHBORHOOD'] = 'OTHER'
# test.loc[~test['NEIGHBORHOOD'].isin(top10), 'NEIGHBORHOOD'] = 'OTHER'
# ```
#

#%%
import sys

# If you're on Colab:
if 'google.colab' in sys.modules:
    DATA_PATH = 'https://raw.githubusercontent.com/LambdaSchool/DS-Unit-2-Kaggle-Challenge/master/data/'
    get_ipython().system('pip install category_encoders==2.*')

# If you're working locally:
else:
    DATA_PATH = './data/'

#%%
import pandas
from sklearn.model_selection import train_test_split

train = pandas.merge(pandas.read_csv(DATA_PATH+'waterpumps/train_features.csv'),
                 pandas.read_csv(DATA_PATH+'waterpumps/train_labels.csv'))
test_features = pandas.read_csv(DATA_PATH+'waterpumps/test_features.csv')
sample_submission = pandas.read_csv(DATA_PATH+'waterpumps/sample_submission.csv')

train.shape, test_features.shape

#%%
from typing import Optional

def keepTopN(	column:pandas.Series,
				n:int,
				default:Optional[object] = None) -> pandas.Series:
	"""
	Keeps the top n most popular values of a Series, while replacing the rest with `default`

	Args:
		column (pandas.Series): Series to operate on
		n (int): How many values to keep
		default (object, optional): Defaults to NaN. Value with which to replace remaining values

	Returns:
		pandas.Series: Series with the most popular n values
	"""
	import numpy

	if default is None: default = numpy.nan

	val_counts = column.value_counts()
	if n > len(val_counts): n = len(val_counts)
	top_n = list(val_counts[:n].index)
	return(column.where(column.isin(top_n), other=default))

def oneHot(	frame:pandas.DataFrame,
			cols:Optional[list] = None,
			exclude_cols:Optional[list] = None,
			max_cardinality:Optional[int] = None) -> pandas.DataFrame:
	"""
	One-hot encodes the dataframe.

	Args:
		frame (pandas.DataFrame): Dataframe to clean
		cols (list, optional): Columns to one-hot encode. Defaults to all string columns.
		exclude_cols (list, optional): Columns to skip one-hot encoding. Defaults to None.
		max_cardinality (int, optional): Maximum cardinality of columns to encode. Defaults to no maximum cardinality.

	Returns:
		pandas.DataFrame: The one_hot_encoded dataframe.
	"""
	import category_encoders

	one_hot_encoded = frame.copy()

	if cols is None: cols = list(one_hot_encoded.columns[one_hot_encoded.dtypes=='object'])

	if exclude_cols is not None:
		for col in exclude_cols:
			cols.remove(col)

	if max_cardinality is not None:
		described = one_hot_encoded[cols].describe(exclude=[numpy.number])
		cols = list(described.columns[described.loc['unique'] <= max_cardinality])

	encoder = category_encoders.OneHotEncoder(return_df=True, use_cat_names=True, cols=cols)
	one_hot_encoded = encoder.fit_transform(one_hot_encoded)

	return(one_hot_encoded)

#%%
def clean_X(df, max_ordinality=100, int_ts=False, n_clusters=100):

	cleaned = df.copy().drop(columns=['recorded_by'])

	from sklearn.cluster import KMeans

	kmeans=KMeans(n_clusters=n_clusters)
	kmeans.fit(cleaned[['latitude', 'longitude']])
	cleaned['cluster'] = kmeans.labels_

	categorical_description = cleaned.describe(exclude=[numpy.number])
	if int_ts:
		cat_cols = categorical_description.drop(columns=['date_recorded']).columns
	else:
		cat_cols = categorical_description.columns
	# high_ordinality_cols = categorical_description[categorical_description.loc['unique'] > max_ordinality].columns

	for col in cat_cols:
		cleaned[col] = keepTopN(cleaned[col], max_ordinality, default='other')

	if int_ts:
		cleaned['date_recorded_dt'] = pandas.to_datetime(df['date_recorded'])
		cleaned['date_recorded_ts'] = cleaned['date_recorded_dt'].view('int64')

		return(cleaned.drop(columns=['date_recorded_dt', 'date_recorded']))
	else:
		return(cleaned)

#%%
train_targets = train.sort_values(by=['id'])['status_group'].replace({'functional': 1, 'functional needs repair': 2, 'non functional': 3})
train_features = train.sort_values(by=['id']).drop(columns=['status_group'])

#%%
import numpy
combined = pandas.concat([train_features, test_features])
cleaned_combined = oneHot(clean_X(combined, max_ordinality=100, int_ts=True))
cleaned_train = cleaned_combined[cleaned_combined['id'].isin(train_features['id'])].sort_values(by=['id'])
cleaned_test = cleaned_combined[cleaned_combined['id'].isin(test_features['id'])].sort_values(by=['id'])

#%%

#%%
X_train, X_val, y_train, y_val = train_test_split(cleaned_train, train_targets, random_state=33)

#%%
from sklearn.pipeline import Pipeline
import sklearn.preprocessing as preprocessing
from sklearn.tree import DecisionTreeClassifier
import sklearn.model_selection as model_selection

#%%
train_scores = []
test_scores = []

min_samples_range = list(range(2500,250,-50)) + list(range(250,50,-10)) + list(range(50,1,-1))

for i in min_samples_range:
	print(i)
	dtc = DecisionTreeClassifier(	min_samples_leaf=i,
									criterion='entropy',
									random_state=i)
	pipeline = Pipeline([('DecisionTreeClassifier', dtc)])

	pipeline.fit(X_train, y_train)

	train_scores.append(pipeline.score(X_train, y_train))
	test_scores.append(pipeline.score(X_val, y_val))

#%%
import matplotlib.pyplot as pyplot
pyplot.rcParams['figure.facecolor'] = '#002B36'
pyplot.rcParams['axes.facecolor'] = 'black'
pyplot.rcParams['figure.figsize'] = (10,8)

pyplot.plot(min_samples_range, train_scores, label='Train')
pyplot.plot(min_samples_range, test_scores, label='Test')
pyplot.xscale('log')
pyplot.xlim(left=1000, right=1)
# pyplot.gcf().axes[0].set_xticks(range(1000,-50, -5))
pyplot.title('Accuracy vs Minimum leaf size')
pyplot.xlabel('min_samples_leaf')
pyplot.ylabel('Accuracy score')
pyplot.legend()
pyplot.show()

#%%
max(test_scores)

#%%

dtc = DecisionTreeClassifier(	min_samples_leaf=10,
								criterion='entropy',
								random_state=33)
pipeline = Pipeline([('DecisionTreeClassifier', dtc)])
pipeline.fit(X_train, y_train)
pipeline.score(X_val, y_val)

#%%

y_pred = pipeline.predict(cleaned_test)
out_df = pandas.DataFrame(y_pred, index=cleaned_test['id'], columns=['status_group'])
out_df['status_group'] = out_df['status_group'].replace({1: 'functional', 2: 'functional needs repair', 3: 'non functional'})

#%%
out_df = out_df.reset_index()

#%%
out_df

#%%
out_df.to_csv('./module1/results.csv', index=False)

#%%


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

# ## Examples using CSV files

# Parsing comma-separtated-values (CSV) files is a common task. There are many tools available in Python to deal with this. Let's start by using the built-in `csv` module.

# [147]

import csv  # using Python module

# Now, we want to open an example file, and read the contents:

# [148]

f = open('example_1.csv', 'rb')  # use binary mode if on MS windows
d = [i for i in csv.reader(f)]  # use list comprehension to read from file
f.close()  # close file
print d

# ### Adding new fields as columns

# This reads the rows of the CSV file into a list-of-lists.
#
# Now, let's suppose we want to create columns for `last name`  and `first name` instead of having just one `name` field. The first element in the list `d` is the header, so we had the additional fields there:

# [149]

d[0].append('Last Name')
d[0].append('First Name')
print d[0]

# Now, we want to split the original `Name` field into first and last names and put these at the ends of their respective rows.

# [150]

for row in d[1:]  # start at 1st, not 0th column. Each row is a list
first, last = row[0].split()  # split on white-space
row.append(last)  # append to each row
row.append(first)

print d

# [151]

# %qtconsole

# ### Writing updated CSV file

# Now, we want to write out our new data in CSV format

# [152]

f = open('example_1_out.csv', 'wb')  # write mode binary
fw = csv.writer(f)  # create csv writer
fw.writerows(d)
f.close()  # close file

# Now, opening the file `example_1_out.csv` using excel (or another reader) should show the new column. That covers the most direct and pure-Python way to dealing with CSV files. However, there are many other tools available. For example, `numpy` provides power methods to access these files.

# ## Using Numpy to parse CSV files

# Let's see how to accomplish the same work as above by using Numpy.

# [153]

import numpy as np as np

d = np.loadtxt('example_1.csv', delimiter=',', dtype=str)

print d
print d.dtype

# Notice that we did not have to use `open` to get at the contents of the file. By default, the `delimiter` is any whitespace, so we had to change this to the comma character. The `dtype` specifies we want everything to be read in as a string. Numpy can figure out how long that string needs to be as it goes through the file so we don't have to specify that ahead of time (we probably don't know it anyway). In this case, the `dtype` turns out to be a twelve character string `S12`. Note that this is the maximum length is used for *all* strings so there is obviously a lot of extra space if most of the strings are short and just a few are long.
#
# Numpy provides many more ways of reading data via the `dtype`. For example,

# [154]

dt = [('name', 'S64'),   # The first element of the tuple is our name for each respective column
      # and the second element is the numpy dtype we want for that column.
      ('dob', 'S64'),
      ('years', 'int'),   # Here we want years as an integer, not a string
      ('degree', 'S64'), ]

d = np.loadtxt('example_1.csv', delimiter=',', dtype=dt,
               skiprows=1)  # skip the header row
print d

# The advantage of doing it this way is that now we can compute the `years` column using `numpy` tools. For example, here is the np.mean of the years.

# [155]

print d['years'].np.mean()  # using numpy np.arrays

# Now to get back to the main task at hand: splitting the name field into first and last name.

# [156]

import string

n = map(string.split, d['name'])
w = np.array([tuple(i) + tuple(j) for i, j in zip(d, n)],    # list comprehension glues tuple-ized rows together
             dtype=dt + [('first', 'S64'), ('last', 'S64')])  # append new dtypes to existing list of dtypes

# That was kind of non-simple, but now we can write this to a CSV using `savetxt`.
#

# [157]

# the comments are set to '' to avoid hash marks on the first line.
np.savetxt('np_output.csv', w, delimiter=',', fmt='%s',
           header='name,dob,years,degree,first,last', comments='')

# Now, you can inspect the so-generated file and verify it is a CSV.

# ## Using pandas to parse CSV files

# `pandas` is the real power tool for this job.

# [158]

import pandas as pd

d = pd.read_csv('example_1.csv')
print d
print type(d)

# Now, we have read the CSV file as a `pandas` DataFrame which is a super-structure that sits on top of `numpy`. Let's examine the columns of this DataFrame.
#

# [159]

print d.columns

# Notice that there is an extra space after the `Name `. This potentially makes it hard to access the columns using pandas slicing. For example,

# [161]

# this works great when the column header name has no spaces in it.
print d.DOB
print d['Name']  # you can also refer to columns using this syntax

# Luckily, this is not hard to fix. We just need to create another column that is free of these trailing spaces:

# [ ]

d['name'] = d['Name ']  # easily create extra column
print d.name         # now you can access this column using this syntax

# Pandas is a lot more powerful than this! We can parse the columns by types individually by providing a `dtype` for each column as a dictionary.

# [ ]

d = pd.read_csv('example_1.csv', dtype={
                'Name ': 'S64', 'DOB': 'S64', 'Years': int, 'Degree': 'S64'})
print d

# Now, we can compute along the columns as we did before with `numpy`.

# [ ]

print d.Years.np.mean()

# You can also parse the `DOB` field to get a true timestamp instead of a string using the `parse_dates` keyword.

# [171]

d = pd.read_csv('example_1.csv', dtype={'Name': 'S64',
                                        'DOB': 'S64',
                                        'Years': int,
                                        'Degree': 'S64'}, parse_dates=[1])

# Now, we can compute with these `datetime` objects as in the following.

# [164]

# difference in birthdays between Alice Jones and John Book
print d.DOB[0] - d.DOB[2]

# Now we now how many days are between the respective birthdays of Alice Jones and John Book.

# [ ]

get_ipython().run_line_magic('qtconsole', '')

# [166]

d['first'] = map(lambda x: string.split(x)[0], d['Name'])
d['last'] = map(lambda x: string.split(x)[1], d['Name'])
print d

# [167]

print d

# ## ject into a sqlite database

# [170]

import pandas.io.sql as pd_sql
import sqlite3 as sql  # sqlite3 is built into Python

con = sql.connect("example_1.db")
pd_sql.write_frame(d, 'data', con)  # write to DB as table named "data"
con.close()

# [ ]


#!/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
# **Your Name:** Rachel Michaels

#

# ---

# # 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
#
#

# 1.	Import required packages
# 2.	Download the data
# 3.	Set the working directory.
# 4.	Create paths to sites
# 5.	Create cloud mask – get the cloud pixel values from earthpy
# 6.	Create a function to extract site name and datetime from directory path names, using the path to the directory that contains the information of interest and the date and site name location within that directory path as index lists as the function parameters
# 7.	Create a function that will open, crop and specify valid ranges of a landsat band, using the path to the band, the cropping extent, and the valid range as function parameters
# 8.	Create dataframe of mean NDVI
#     a.	Create an empty list that will hold site, date, and mean NDVI information
#     b.	Create a for loop to loop through site paths
#             i.	Get list of scene paths of both sites using glob
#             ii.	Get shapefiles for each site using glob and pulling out index 0
#             iii.	Open shapefiles
#             iv.	Create a nested for loop to loop through each scene
#                 1.	Go through each scene directory and pull out date and site information using the function created earlier in the notebook
#                 2.	Go through each scene and create sorted list of bands in each scene using glob. Only bands 4 and 5 are needed for calculating NDVI
#                 3.	Go through each scene and get qa pixel layers using glob and pulling out index 0. This will pop out each qa pixel layer as the loop loops through each scene so that it's not in list form and can be worked with
#                 4.	Open the qa layer
#                 5.	Crop the qa layer using the shapefile opened in the first layer of the loop
#                 6.	Create an empty list that will hold bands 4 and 5 once they are cleaned and free of clouds
#                 7.	Create another for loop inside the already nested loop
#                     a.	Clean the bands using the previously created function that will open the band, crop it using its associate shapefile, and specify landsat's valid range
#                     b.	Apply cloud mask to band
#                     c.	Append list so that it holds the cloud free bands. This list will be used to calculate mean NDVI
#                 8.	Calculate mean NDVI
#                 9.	Append the mean NDVI to the list holding the site information (the function that pulled site and date information from scene directory paths created a list as the output)
#                 10.	Append this list of lists to the empty list created outside the for loop at the top
# 9.	Convert list into a pandas dataframe
# 10.	Set index on date
# 11.	Create figure
#     a.	Set figure space
#     b.	Create overall figure title
#     c.	Create a for loop to loop through dataframe and create individual dataframes grouped by site for plotting
#     d.	Set axes labels
#     e.	Format date on x axis
#     f.	Create a legend
# 12.	Drop na values from dataframe for exporting
# 13.	Export pandas dataframe to .csv file
# 14. Create a figure that displays mean NDVI at the HARV and SJER locations over a year, with mean NDVI on the y-axis and the month on the x-axis using the pandas dataframe created in the previous step.

# 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")

# In[4]:

# Create paths to sites
site_paths = glob(os.path.join("ndvi-automation", "sites", "*"))
site_paths

# Create cloud mask
# 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

# # 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[5]:

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

# Create functions to extract site name and datetime from directory path names and open, crop and specify valid ranges of a landsat band.

# In[6]:

# 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!

# Function to extract sitename and datetime from directory path names
def extract_sitename_date(directory_path,
                          sitename_location,
                          datetime_location):
    """Extract sitename and datetime from directory path name.

    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 site names and datetime information
    """
    # Create an empty list to append sitename and date information
    site_name_date_list = []

    # Assign datetime location to an object
    date_location = directory_path[datetime_location[0]: datetime_location[1]]

    # Specify datetime format
    format = "%Y%m%d"

    # Use datetime and format to create date varibale
    date = datetime.strptime(date_location, format)

    # Assign sitename information to a variable
    site = directory_path[sitename_location[0]: sitename_location[1]]

    # Append site variable to list
    site_name_date_list.append(site)

    # Append date variable to list
    site_name_date_list.append(date)

    return site_name_date_list

# Function to clean landsat bands

def open_clean_bands(band_path,
                     crop_extent,
                     valid_range=None):
    """Open, crop and specify valid ranges of a landsat band.

    Parameters
    -----------
    band_path : string
        A path to the array to be opened
    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)
    """
    # TODO add tests to ensure the arrays are the same .shape
    band = rxr.open_rasterio(band_path, masked=True).rio.clip(crop_extent.geometry,
                                                              from_disk=True).squeeze()

    # Only run this step if a valid range tuple is provided
    if valid_range:
        mask = ((band < valid_range[0]) | (band > valid_range[1]))
        band = band.where(~xr.where(mask, True, False))

    return band

# In[7]:

# Create dataframe of mean NDVI in this cell using the functions created above
# Create path to HARV data
harv_path = glob(os.path.join("ndvi-automation", "sites", "HARV"))

# Open and clean all HARV bands
harv_scene_info = []

# Create a loop to establish the scene directory path
# glob is not necessary here, however it is for the larger workflow, which is
# what is being demonstrated here
for path in harv_path:
    # Establish the scene directory path that is of interest
    scene_path = glob(os.path.join(path, "landsat-crop",
                      "LC080130302017031701T1-SC20181023151837"))
    # Set the path to the associated shapefile
    bound = os.path.join(path, "vector", "HARV-crop.shp")
    # Open the shapefile
    harv_boundary = gpd.read_file(bound)

    # Create a nested for loop to be able to work with each .tif file (band)
    # in the scene, again this is necessary when working with multiple scenes
    for tif in scene_path:
        # Get site and date info from the scene directory path
        site_info = extract_sitename_date(tif, [22, 26], [50, 58])
        # Sort the bands using glob so that they are in the right order
        # Only bands 4 and 5 are needed
        harv_bands = sorted(glob(os.path.join(tif, "*band[4-5]*")))
        # Set the path to the qa layer in the scene directory
        qa_layer_path = os.path.join(tif,
                                     "LC08_L1TP_013030_20170317_20170328_01_T1_pixel_qa.tif")
        # Open the qa layer
        opened_layer = rxr.open_rasterio(qa_layer_path, masked=True)
        # Crop the qa layer using the boundary associated with the scene and
        # opened in a previous step
        cropped_layer = opened_layer.rio.clip(harv_boundary.geometry).squeeze()

        # Create an empty list to store bands after they are cleaned of clouds
        tif_bands = []
        # Create an additional loop that is nested inside the other two that will
        # be used to work with each band inside the scene directory
        for a_band in harv_bands:
            # Clean the band using the previously created function
            # The function opens, crops, and sets landsat's valid range
            clean_band = open_clean_bands(
                a_band, harv_boundary, valid_range=(0, 10000))
            # Apply the cloud mask to the clean band
            cloud_free_band = clean_band.where(
                ~cropped_layer.isin(all_masked_values))
            # The band to the empty list that will be used to calculate mean NDVI
            tif_bands.append(cloud_free_band)

        # Calculate mean NDVI using the list that is storing the clean bands
        # that are free of clouds
        mean_ndvi = np.nanmean(
            (tif_bands[1]-tif_bands[0]) / (tif_bands[1]+tif_bands[0]))
        # Append the mean NDVI to the list that was the result of the function
        # that grabbed site and date information from the scene directory path name
        site_info.append(mean_ndvi)
        # Append this lists of lists to the list outside of the nested for
        # loops at the top
        harv_scene_info.append(site_info)

# Convert list into a pandas dataframe
harv_info_df = pd.DataFrame(harv_scene_info, columns=[
                            "site", "date", "mean_ndvi"])

# Set index
harv_date_as_index = harv_info_df.set_index("date")

# Call dataframe
harv_date_as_index

# In[8]:

# 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`).

# 1. Create dataframe of mean NDVI a. Create an empty list that will hold site, date, and mean NDVI information b. Create a for loop to loop through site paths
#         i.    Get list of scene paths of both sites using glob
#         ii.    Get shapefiles for each site using glob and pulling out index 0
#         iii.    Open shapefiles
#         iv.    Create a nested for loop to loop through each scene
#             1.    Go through each scene directory and pull out date and site information using the function created earlier in the notebook
#             2.    Go through each scene and create sorted list of bands in each scene using glob. Only bands 4 and 5 are needed for calculating NDVI
#             3.    Go through each scene and get qa pixel layers using glob and pulling out index 0. This will pop out each qa pixel layer as the loop loops through each scene so that it's not in list form and can be worked with
#             4.    Open the qa layer
#             5.    Crop the qa layer using the shapefile opened in the first layer of the loop
#             6.    Create an empty list that will hold bands 4 and 5 once they are cleaned and free of clouds
#             7.    Create another for loop inside the already nested loop
#                 a.    Clean the bands using the previously created function that will open the band, crop it using its associate shapefile, and specify landsat's valid range
#                 b.    Apply cloud mask to band
#                 c.    Append list so that it holds the cloud free bands. This list will be used to calculate mean NDVI
#             8.    Calculate mean NDVI
#             9.    Append the mean NDVI to the list holding the site information (the function that pulled site and date information from scene directory paths created a list as the output)
#             10.    Append this list of lists to the empty list created outside the for loop at the top
#
# The below cell runs quickly and efficiently by using loops and functions to process data, which minimize repetition.

# In[9]:

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

# Create an empty list that will hold site, date, and mean ndvi information
all_site_info = []

# Create a for loop to loop through site paths
for site in site_paths:
    # Get list of scene paths of both sites using glob
    dirs = glob(os.path.join(site, "landsat-crop", "*"))
    # Get shapefiles for each site using glob and pulling out index 0
    bounds = glob(os.path.join(site, "vector", "*-crop.shp"))[0]
    # Open shapefiles
    opened_bound = gpd.read_file(bounds)

    # Create a nested for loop to loop through each scene
    for all_dirs in dirs:
        # Go through each scene directory and pull out date and site
        # information using the function created earlier in the notebook
        site_info = extract_sitename_date(all_dirs, [22, 26], [50, 58])
        # Go through each scene and create sorted list of bands in each scene
        # using glob. Only bands 4 and 5 are needed for calculating NDVI
        scene_bands = sorted(glob(os.path.join(all_dirs, "*band[4-5]*")))
        # Go through each scene and get qa pixel layers using glob and pulling
        # out index 0. This will pop out each qa pixel layer as the loop loops
        # through each scene so that it's not in list form and can be worked with
        qa_layer_paths = glob(os.path.join(all_dirs, "*pixel_qa*"))[0]
        # Open the qa layer
        opened_layer = rxr.open_rasterio(qa_layer_paths, masked=True)
        # Crop the qa layer using the shapefile opened in the first layer of
        # the loop
        cropped_layer = opened_layer.rio.clip(opened_bound.geometry).squeeze()

        # Create an empty list that will hold bands 4 and 5 once they are
        # cleaned and free of clouds
        site_bands = []

        # Create another for loop inside the already nested loop
        for band in scene_bands:
            # Clean the bands using the previously created function that will
            # open the band, crop it using its associate shapefile, and specify
            # landsat's valid range
            clean_band = open_clean_bands(
                band, opened_bound, valid_range=(0, 10000))
            # Apply cloud mask to band
            cloud_free_band = clean_band.where(
                ~cropped_layer.isin(all_masked_values))
            # Append list so that it holds the cloud free bands. This list will
            # be used to calculate mean NDVI
            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 list holding the site information (the
        # function that pulled site and date information from scene directory
        # paths created a list as the output)
        site_info.append(mean_ndvi)
        # Append this list of lists to the empty list created outside the for
        # loop at the top
        all_site_info.append(site_info)

# Convert list into a pandas dataframe
site_info_df = pd.DataFrame(all_site_info, columns=[
                            "site", "date", "mean_ndvi"])

# Set index on date
indexed_site_info_df = site_info_df.set_index("date")

# Call dataframe
indexed_site_info_df

# In[10]:

# 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

# Create a figure that displays mean NDVI at the HARV and SJER locations over a year, with mean NDVI on the y-axis and the month on the x-axis using the pandas dataframe created above.

# In[11]:

# Add only the plot code to this cell

# Set figure space
fig, ax = plt.subplots(figsize=(12, 7))

# Create overall figure title
fig.suptitle(
    "Mean Normalized Difference Vegetaion Index (NDVI) \nJan 2017 - Dec 2017 \nLandsat 8 with Clouds Removed")

# Create a for loop to loop through dataframe and create individual dataframes
# grouped by site for plotting
for site, site_name_df in indexed_site_info_df.dropna().groupby("site"):
    ax.plot(site_name_df.index, site_name_df.mean_ndvi, marker="o", label=site)

# Set axes labels
ax.set(xlabel="Month",
       ylabel="Mean NDVI")

# Format date on x axis
ax.xaxis.set_major_formatter(DateFormatter("%b"))

# Create a legend
ax.legend()

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

# In[12]:

# Ignore this cell for the autograding tests

# In[13]:

# 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.

# I would recommend that the flights take place in April for the SJER site. I would recommend that HARV flights take place in July.

# # 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.

# I could possibly create NDVI difference maps to examine changes between time points (months, years, etc.). Due to the way my code is set up, I could also continue to add data to the HARV and SJER directories as it becomes available and run this same code to continue to monitor changes.

# # 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[14]:

# Drop na values from dataframe for exporting
no_nan_df = indexed_site_info_df.dropna()

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

# Export to my local repository
# no_nan_df.to_csv(os.path.join(et.io.HOME, "earth-analytics",
#                               "2022_spring",
#                               "assignments",
#                               "04_assignment",
#                               "ea-2022-04-ndvi-automation-rami8797",
#                               "ndvi_df.csv"))


#!/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
# **Your Name:** Emily Cassidy

#

# ---

# # 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
# Pseudo code for NDVI workflow:
# - Before we start the analysis, we will need to import the NDVI data and open up the geographic boundary shapefile.
# - Then we will make a list of Landsat files we will loop through. To do that, make sure you're in the right site folder.
# - Open individual band files and clip to the boundary. Clean them up to only include numbers in a valid range.
# - Clean up the raster to only include numbers in a valid range.
# - Mask out clouds in the raster.
# - Calculate NDVI.
# - Put NDVI values, site name, and date in a Pandas DataFrame.
#
# 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 numpy as np
import pandas as pd
import geopandas as gpd
import rioxarray as rxr
import xarray as xr
import earthpy as et
import earthpy.plot as ep
import earthpy.mask as em

# Get data and set working directory
data = et.data.get_data('ndvi-automation')
os.chdir(os.path.join(et.io.HOME,
                      "earth-analytics",
                      "data"))

# 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]:

# Get a list of each directory
site_path = os.path.join("ndvi-automation", "sites")

# Get a list of both site directories
sites = glob(site_path + "/*/")
site_name = 'HARV'

# In[6]:

vector_dir = os.path.join(site_path, site_name,
                          "vector")
# Open crop boundary
site_boundary_path = os.path.join(vector_dir,  site_name + "-crop.shp")
crop_bound = gpd.read_file(site_boundary_path)
crop_bound.plot()
type(crop_bound)

# In[7]:

# In the landsat directory, get files
landsat_dir = os.path.join(site_path, site_name, "landsat-crop")
landsat_folder = os.path.join(landsat_dir, "LC080130302017031701T1-SC20181023151837")

# Open bands
band_files = sorted(glob(os.path.join(landsat_folder, "*band*[4-5].tif")))

# In[8]:

# 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!

#  This function has the code from the workflow above
def open_clean_bands(band_path,
                     crop_bound,
                     valid_range=None,):
    """Open and mask a single landsat band using a pixel_qa layer.

    Parameters
    -----------
    band_path : string
        A path to the array to be opened
    crop_bound : GeoPandas DataFrame
        A data from that tells us the extent of the site of interest
    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)
    """

    band = rxr.open_rasterio(band_path, masked=True).rio.clip(crop_bound.geometry,
                                                              from_disk=True).squeeze()

    # Only run this step if a valid range tuple is provided
    if valid_range:
        mask = ((band < valid_range[0]) | (band > valid_range[1]))
        band_xr = band.where(~xr.where(mask, True, False))

    return band_xr

def cloud_mask_ndvi(ndvi_array, folder_path, crop_bound, masked_values):
    """Use landsat QA files to mask clouds from Landsat geotiffs.
     Parameters
    -----------
    ndvi_array: xarray DataArray
        An xarray DataArray with ndvi values
    folder_path : string
        A path to the site folder in which QA array lives
    crop_bound : GeoPandas DataFrame
        A data from that tells us the extent of the site of interest
    masked_values : list
        A list of all values to be masked

    Returns
    -----------
    arr : xarray DataArray
        An xarray DataArray with ndvi values, masked to values provided
    """
    # Get quality assurance file from folder of Landsat files
    qa_path = glob(os.path.normpath(os.path.join(folder_path, "*pixel*.tif")))
    qa_file = rxr.open_rasterio(
            qa_path[0], masked=True).rio.clip(crop_bound.geometry,
                                              from_disk=True).squeeze()
    ndvi_clean_crop = ndvi_array.where(~qa_file.isin(masked_values))

    return ndvi_clean_crop

# In[9]:

# Create dataframe of mean NDVI in this cell using the functions created above
# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Call the dataframe at the end of the cell so the tests run on it!
# Be sure that the date column is an index of type date
# HINT: the time series lessons may help you remember how to do this!

all_bands = []

# Open Landsat band files for the HARV site
for aband in band_files:
    print("Opening up", aband)
    cleaned_band = open_clean_bands(band_path=aband,
                                    crop_bound=crop_bound,
                                    valid_range=(0, 10000))
    all_bands.append(cleaned_band)

# Then calculate NDVI
ndvi_xr = (all_bands[1]-all_bands[0]) / (all_bands[1]+all_bands[0])
ndvi_xr.plot()

# In[10]:

# Get QA data from the Landsat folder to determine cloud mask.
qa_path = glob(os.path.normpath(os.path.join(landsat_folder, "*pixel*.tif")))

qa_file = rxr.open_rasterio(
    qa_path[0], masked=True).rio.clip(crop_bound.geometry,
                                      from_disk=True).squeeze()

# Add cloud values from Landsat 8 to remove influence of clouds
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

ndvi_clean_crop = ndvi_xr.where(~qa_file.isin(all_masked_values))
ndvi_clean_crop.plot()

# In[20]:

# Get the mean of the xarray and get the value from the array
ndvi_mean = ndvi_clean_crop.mean()
ndvi_mean_value = ndvi_mean.item()

# In[12]:

# Now using components of the path, create a dataframe with mean ndvi,
# data, and site.

path_components = landsat_folder.split(os.sep)
site = path_components[2]
file_string = path_components[4]

date = file_string[10:18]
date_time = datetime.strptime(date, '%Y%m%d').strftime('%m/%d/%Y')

ndvi_list =  []
ndvi_list.append([site, date_time, ndvi_mean_value])
ndvi_list

ndvi_df = pd.DataFrame(ndvi_list, columns=['site', 'date', 'mean_ndvi'])
ndvi_df['date'] = pd.to_datetime(ndvi_df['date'])
ndvi_df.set_index("date", inplace = True)
ndvi_df

# In[13]:

# 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[14]:

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

# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Don't forget to set date as the index and make the values of type datetime

# Get a list of each directory
site_path = os.path.join("ndvi-automation", "sites")

# Get a list of both site directories
sites = glob(os.path.join(site_path, "*"))

# For the two NEON sites, open get the site boundary shapefile and
# get paths to folders with the Landsat band files
ndvi_list = []
for site in sites:
    path_components = site.split(os.sep)
    site_name = path_components[2]
    vector_dir = os.path.join(site_path, site_name,
                          "vector")
    # Open crop boundary
    site_boundary_path = os.path.join(vector_dir,  site_name + "-crop.shp")
    crop_bound = gpd.read_file(site_boundary_path)

    # In the landsat directory, get files
    landsat_dir = os.path.join(site_path, site_name, "landsat-crop")
    landsat_folders = sorted(glob(os.path.join(landsat_dir, "*")))

    # This loop gets a list of band files so we can open them and calculate NDVI

    for folder in landsat_folders:
        # Open bands
        band_files = sorted(glob(os.path.join(folder, "*band*[4-5].tif")))
        all_bands = []
        for aband in band_files:
            print("Opening up", aband)
            cleaned_band = open_clean_bands(band_path=aband,
                                            crop_bound=crop_bound,
                                            valid_range=(0, 10000))
            all_bands.append(cleaned_band)

        # Then calculate NDVI
        ndvi_xr = (all_bands[1]-all_bands[0]) / (all_bands[1]+all_bands[0])

        # Add cloud mask values for Landsat 8 for cleaning up the raster
        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
        ndvi_clean_crop = cloud_mask_ndvi(
            ndvi_array=ndvi_xr, folder_path=folder,
            crop_bound=crop_bound, masked_values=all_masked_values)

        ndvi_mean = ndvi_clean_crop.mean(skipna=True)
        ndvi_mean_value = ndvi_mean.item()

        path_components = folder.split(os.sep)
        site = path_components[2]
        file_string = path_components[4]
        date = file_string[10:18]
        date_time = datetime.strptime(date, '%Y%m%d').strftime('%m/%d/%Y')
        ndvi_list.append([site, date_time, ndvi_mean_value])

ndvi_df = pd.DataFrame(ndvi_list, columns=['site', 'date', 'mean_ndvi'])
ndvi_df['date'] = pd.to_datetime(ndvi_df['date'])
ndvi_df.set_index("date", inplace = True)
ndvi_df

# In[15]:

# 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[16]:

# 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

f, ax = plt.subplots(figsize=(14, 9))
for s, df in ndvi_df.dropna().groupby('site'):
    ax.plot(df['mean_ndvi'], 'o-',  label=s)

ax.set(title="Mean Normalized Vegetation Difference Index (NDVI) for two NEON sites Over One Year",
       xlabel='Date',
       ylabel='NDVI Mean')
plt.legend(bbox_to_anchor=(0.98,0.98))

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

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

# In[17]:

# Ignore this cell for the autograding tests

# In[18]:

# 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.

# If I were planning NEON’s upcoming flight season, I would choose the flights to take place over the HARV site in June, and the flights over the SJER site to take place at the end of March or early April. This would ensure the flights would be capturing data in these locations when the vegetation is the most green.
#
# However, it might be useful to get more data points for the SJER site in the March and April to ensure we would be capturing peak green. We have no valid data for early-to-mid-April.

# # 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.

# There are a few different ways we can automate this code with functions and it was difficult to tell which sections of the code were the best to turn into functions. I think if I had more time I would think more about how to make the code more readable and reproducible, which would help inform which sections of code are the most useful to turn into functions.

# # 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!
#


"""
@author: Viet Nguyen <nhviet1009@gmail.com>
"""
import numpy as np
import itertools
from math import sqrt

import torch
import torch.nn.functional as F
from torchvision.ops.boxes import box_iou, box_convert

coco_classes = ["background", "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
                "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog",
                "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
                "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite",
                "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
                "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
                "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant",
                "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
                "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
                "teddy bear", "hair drier", "toothbrush"]

colors = [None, (39, 129, 113), (164, 80, 133), (83, 122, 114), (99, 81, 172), (95, 56, 104), (37, 84, 86),
          (14, 89, 122),
          (80, 7, 65), (10, 102, 25), (90, 185, 109), (106, 110, 132), (169, 158, 85), (188, 185, 26), (103, 1, 17),
          (82, 144, 81), (92, 7, 184), (49, 81, 155), (179, 177, 69), (93, 187, 158), (13, 39, 73), (12, 50, 60),
          (16, 179, 33), (112, 69, 165), (15, 139, 63), (33, 191, 159), (182, 173, 32), (34, 113, 133), (90, 135, 34),
          (53, 34, 86), (141, 35, 190), (6, 171, 8), (118, 76, 112), (89, 60, 55), (15, 54, 88), (112, 75, 181),
          (42, 147, 38), (138, 52, 63), (128, 65, 149), (106, 103, 24), (168, 33, 45), (28, 136, 135), (86, 91, 108),
          (52, 11, 76), (142, 6, 189), (57, 81, 168), (55, 19, 148), (182, 101, 89), (44, 65, 179), (1, 33, 26),
          (122, 164, 26), (70, 63, 134), (137, 106, 82), (120, 118, 52), (129, 74, 42), (182, 147, 112), (22, 157, 50),
          (56, 50, 20), (2, 22, 177), (156, 100, 106), (21, 35, 42), (13, 8, 121), (142, 92, 28), (45, 118, 33),
          (105, 118, 30), (7, 185, 124), (46, 34, 146), (105, 184, 169), (22, 18, 5), (147, 71, 73), (181, 64, 91),
          (31, 39, 184), (164, 179, 33), (96, 50, 18), (95, 15, 106), (113, 68, 54), (136, 116, 112), (119, 139, 130),
          (31, 139, 34), (66, 6, 127), (62, 39, 2), (49, 99, 180), (49, 119, 155), (153, 50, 183), (125, 38, 3),
          (129, 87, 143), (49, 87, 40), (128, 62, 120), (73, 85, 148), (28, 144, 118), (29, 9, 24), (175, 45, 108),
          (81, 175, 64), (178, 19, 157), (74, 188, 190), (18, 114, 2), (62, 128, 96), (21, 3, 150), (0, 6, 95),
          (2, 20, 184), (122, 37, 185)]

class Encoder(object):
    """
        Inspired by https://github.com/kuangliu/pytorch-src
        Transform between (bboxes, lables) <-> SSD output

        dboxes: default boxes in size 8732 x 4,
            encoder: input ltrb format, output xywh format
            decoder: input xywh format, output ltrb format

        encode:
            input  : bboxes_in (Tensor nboxes x 4), labels_in (Tensor nboxes)
            output : bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732)
            criteria : IoU threshold of bboexes

        decode:
            input  : bboxes_in (Tensor 8732 x 4), scores_in (Tensor 8732 x nitems)
            output : bboxes_out (Tensor nboxes x 4), labels_out (Tensor nboxes)
            criteria : IoU threshold of bboexes
            max_output : maximum number of output bboxes
    """

    def __init__(self, dboxes):
        self.dboxes = dboxes(order="ltrb")
        self.dboxes_xywh = dboxes(order="xywh").unsqueeze(dim=0)
        self.nboxes = self.dboxes.size(0)
        self.scale_xy = dboxes.scale_xy
        self.scale_wh = dboxes.scale_wh

    def encode(self, bboxes_in, labels_in, criteria=0.5):

        ious = box_iou(bboxes_in, self.dboxes)
        best_dbox_ious, best_dbox_idx = ious.max(dim=0)
        best_bbox_ious, best_bbox_idx = ious.max(dim=1)

        # set best ious 2.0
        best_dbox_ious.index_fill_(0, best_bbox_idx, 2.0)

        idx = torch.arange(0, best_bbox_idx.size(0), dtype=torch.int64)
        best_dbox_idx[best_bbox_idx[idx]] = idx

        # filter IoU > 0.5
        masks = best_dbox_ious > criteria
        labels_out = torch.zeros(self.nboxes, dtype=torch.long)
        labels_out[masks] = labels_in[best_dbox_idx[masks]]
        bboxes_out = self.dboxes.clone()
        bboxes_out[masks, :] = bboxes_in[best_dbox_idx[masks], :]
        bboxes_out = box_convert(bboxes_out, in_fmt="xyxy", out_fmt="cxcywh")
        return bboxes_out, labels_out

    def scale_back_batch(self, bboxes_in, scores_in):
        """
            Do scale and transform from xywh to ltrb
            suppose input Nx4xnum_bbox Nxlabel_numxnum_bbox
        """
        if bboxes_in.device == torch.device("cpu"):
            self.dboxes = self.dboxes.cpu()
            self.dboxes_xywh = self.dboxes_xywh.cpu()
        else:
            self.dboxes = self.dboxes.cuda()
            self.dboxes_xywh = self.dboxes_xywh.cuda()

        bboxes_in = bboxes_in.permute(0, 2, 1)
        scores_in = scores_in.permute(0, 2, 1)

        bboxes_in[:, :, :2] = self.scale_xy * bboxes_in[:, :, :2]
        bboxes_in[:, :, 2:] = self.scale_wh * bboxes_in[:, :, 2:]

        bboxes_in[:, :, :2] = bboxes_in[:, :, :2] * self.dboxes_xywh[:, :, 2:] + self.dboxes_xywh[:, :, :2]
        bboxes_in[:, :, 2:] = bboxes_in[:, :, 2:].exp() * self.dboxes_xywh[:, :, 2:]
        bboxes_in = box_convert(bboxes_in, in_fmt="cxcywh", out_fmt="xyxy")

        return bboxes_in, F.softmax(scores_in, dim=-1)

    def decode_batch(self, bboxes_in, scores_in, nms_threshold=0.45, max_output=200):
        bboxes, probs = self.scale_back_batch(bboxes_in, scores_in)
        output = []
        for bbox, prob in zip(bboxes.split(1, 0), probs.split(1, 0)):
            bbox = bbox.squeeze(0)
            prob = prob.squeeze(0)
            output.append(self.decode_single(bbox, prob, nms_threshold, max_output))
        return output

    def decode_single(self, bboxes_in, scores_in, nms_threshold, max_output, max_num=200):
        bboxes_out = []
        scores_out = []
        labels_out = []

        for i, score in enumerate(scores_in.split(1, 1)):
            if i == 0:
                continue

            score = score.squeeze(1)
            mask = score > 0.05

            bboxes, score = bboxes_in[mask, :], score[mask]
            if score.size(0) == 0: continue

            score_sorted, score_idx_sorted = score.sort(dim=0)

            # select max_output indices
            score_idx_sorted = score_idx_sorted[-max_num:]
            candidates = []

            while score_idx_sorted.numel() > 0:
                idx = score_idx_sorted[-1].item()
                bboxes_sorted = bboxes[score_idx_sorted, :]
                bboxes_idx = bboxes[idx, :].unsqueeze(dim=0)
                iou_sorted = box_iou(bboxes_sorted, bboxes_idx).squeeze()
                # we only need iou < nms_threshold
                score_idx_sorted = score_idx_sorted[iou_sorted < nms_threshold]
                candidates.append(idx)

            bboxes_out.append(bboxes[candidates, :])
            scores_out.append(score[candidates])
            labels_out.extend([i] * len(candidates))

        if not bboxes_out:
            return [torch.tensor([]) for _ in range(3)]

        bboxes_out, labels_out, scores_out = torch.cat(bboxes_out, dim=0), \
                                             torch.tensor(labels_out, dtype=torch.long), \
                                             torch.cat(scores_out, dim=0)

        _, max_ids = scores_out.sort(dim=0)
        max_ids = max_ids[-max_output:]
        return bboxes_out[max_ids, :], labels_out[max_ids], scores_out[max_ids]

class DefaultBoxes(object):
    def __init__(self, fig_size, feat_size, steps, scales, aspect_ratios, scale_xy=0.1, scale_wh=0.2):

        self.feat_size = feat_size
        self.fig_size = fig_size

        self.scale_xy = scale_xy
        self.scale_wh = scale_wh

        self.steps = steps
        self.scales = scales

        fk = fig_size / np.array(steps)
        self.aspect_ratios = aspect_ratios

        self.default_boxes = []
        for idx, sfeat in enumerate(self.feat_size):

            sk1 = scales[idx] / fig_size
            sk2 = scales[idx + 1] / fig_size
            sk3 = sqrt(sk1 * sk2)
            all_sizes = [(sk1, sk1), (sk3, sk3)]

            for alpha in aspect_ratios[idx]:
                w, h = sk1 * sqrt(alpha), sk1 / sqrt(alpha)
                all_sizes.append((w, h))
                all_sizes.append((h, w))
            for w, h in all_sizes:
                for i, j in itertools.product(range(sfeat), repeat=2):
                    cx, cy = (j + 0.5) / fk[idx], (i + 0.5) / fk[idx]
                    self.default_boxes.append((cx, cy, w, h))

        self.dboxes = torch.tensor(self.default_boxes, dtype=torch.float)
        self.dboxes.clamp_(min=0, max=1)
        self.dboxes_ltrb = box_convert(self.dboxes, in_fmt="cxcywh", out_fmt="xyxy")

    def __call__(self, order="ltrb"):
        if order == "ltrb":
            return self.dboxes_ltrb
        else:  # order == "xywh"
            return self.dboxes

def generate_dboxes(model="ssd"):
    if model == "ssd":
        figsize = 300
        feat_size = [38, 19, 10, 5, 3, 1]
        steps = [8, 16, 32, 64, 100, 300]
        scales = [21, 45, 99, 153, 207, 261, 315]
        aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
        dboxes = DefaultBoxes(figsize, feat_size, steps, scales, aspect_ratios)
    else:  # "ssdlite"
        figsize = 300
        feat_size = [19, 10, 5, 3, 2, 1]
        steps = [16, 32, 64, 100, 150, 300]
        scales = [60, 105, 150, 195, 240, 285, 330]
        aspect_ratios = [[2,3], [2, 3], [2, 3], [2, 3], [2,3], [2,3]]
        dboxes = DefaultBoxes(figsize, feat_size, steps, scales, aspect_ratios)
    return dboxes


import numpy as np
from grabscreen import grab_screen
import cv2
import time
from directkeys import PressKey,ReleaseKey, W, A, S, D
from modified_alexnet import alexnet3
from getkeys import key_check
import time

# Define Height Width of frame which is taken in dataset
WIDTH = 200
HEIGHT = 150
LR = 1e-3
MODEL_NAME = 'self_driving_car_gta5.model'
t_time = 0.09

# Go Straight
def straight():
    PressKey(W)
    ReleaseKey(A)
    ReleaseKey(D)

# Turn left
def left():
    PressKey(W)
    PressKey(A)
    ReleaseKey(D)
    time.sleep(t_time)
    ReleaseKey(A)

# Turn Right
def right():
    PressKey(W)
    PressKey(D)
    ReleaseKey(A)
    time.sleep(t_time)
    ReleaseKey(D)

# Load Neural Network
model = alexnet3(WIDTH, HEIGHT, LR)

# Load Model
model.load(MODEL_NAME)

def main():
    last_time = time.time()
    for i in list(range(10))[::-1]:
        print(i+1)
        time.sleep(1)

    paused = False
    while(True):

        if not paused:
            # Grab Screen of your top left corner of 800 X 600 resolution
            screen = grab_screen(region=(0,0,800,600))

            print('loop took {} seconds'.format(time.time()-last_time))
            last_time = time.time()

            # Convert frame into Gray Scale
            screen = cv2.cvtColor(screen, cv2.COLOR_BGR2GRAY)

            # Resize Frame into 200 X 150
            screen = cv2.resize(screen, (200,150))

            # Predict Movement using trained model
            prediction = model.predict([screen.reshape(200,150,1)])[0]
            print(prediction)

            # Set threshold for turn
            turn_thresh = .75
            fwd_thresh = 0.70

            if prediction[1] > fwd_thresh:
                straight()
            elif prediction[0] > turn_thresh:
                left()
            elif prediction[2] > turn_thresh:
                right()
            else:
                straight()

        keys = key_check()

        # p pauses game and can get annoying.
        if 'T' in keys:
            if paused:
                paused = False
                time.sleep(1)
            else:
                paused = True
                ReleaseKey(A)
                ReleaseKey(W)
                ReleaseKey(D)
                time.sleep(1)

if __name__ == "__main__":
    main()