import numpy as np
import pandas as pd
import seaborn as sns
import statsmodels.formula.api as sm
import statsmodels.imputation.mice as mice
from statsmodels.regression.linear_model import OLS
from collections import defaultdict
import matplotlib.pyplot as plt
from sklearn.feature_selection import VarianceThreshold
from sklearn.preprocessing import PolynomialFeatures
from sklearn.feature_selection import SelectKBest
from sklearn.feature_selection import chi2


df = pd.read_csv('trending track prediction set_2500tracks360days_minavg1000.csv', encoding = "ISO-8859-1")
# COLUMNS: TRACKID, DOWNLOAD_ACTIVITY_DATE, SUM(UNITS)

df.columns = ['TRACKID', 'DATE', 'STREAMS']


print(df[:10])


df['ROLLING_AVG'] = df.groupby('TRACKID')['STREAMS'].rolling(365, min_periods=0).mean().reset_index(0,drop=True)
df['ROLLING_STD'] = df.groupby('TRACKID')['STREAMS'].rolling(365, min_periods=0).std().reset_index(0,drop=True)

df['ZSCORE'] = (df['STREAMS']-df['ROLLING_AVG'])/df['ROLLING_STD']


#set threshold value here
df['TRENDING'] = np.where(df['ZSCORE']>=8, 1,0)

df['PRETRENDING'] = 0

for i, row in df.iterrows():
	if row['TRENDING'] == 1:
		df.loc[i-15:i-1,'PRETRENDING'] = 1

#uncomment this line if using with AWS ML
#df = df.drop(columns=['ROLLING_AVG', 'ROLLING_STD', 'ZSCORE'])	

df.to_csv('training_set.csv', index=False)