import pandas as pd
import numpy as np
import matplotlib.dates as mdates
from matplotlib import pyplot as plt
plt.style.use(['dark_background'])

# + ECHT.UEBERSEE-RUM... (ID = 6)
# + GALLO  ZINFANDEL... (ID = 55)
# - TOILETTENPAPIER 3... (ID = 38)


plt.rcParams.update({'font.size': 18})
artikel_df = pd.read_csv('./data_new/REWE_ARTIKEL.csv')
artikel = 55
markt = 4
artikel_name = artikel_df.loc[artikel_df['ARTIKEL'] == artikel, 'BEZEICHNUNG'].values[0]
where_condition = "ARTIKEL ==" + str(artikel) + " AND MARKT == " + str(markt) + \
                " AND KAL_TAG > '2017-12-31' AND GESCHLOSSEN == 0"

df = prediction.where(where_condition).select("KAL_TAG", "ABVERKAUF_STK", "PREDICTION", "diff").orderBy("KAL_TAG").toPandas()
#.where("KAL_TAG > '2017-31-12'")

y_1 = list(df["ABVERKAUF_STK"])
y_2 = list(df["PREDICTION"])
y_3 = list(df["diff"])

x = list(df["KAL_TAG"])

#fig, ax = plt.subplots(figsize=(16, 9))
#for idx, x_i in enumerate(x):
#    if y_3[idx] >= 0:
#        col = 'green'
#    else:
#        col = 'red'
#    ax.plot((x_i, x_i), (y_1[idx], y_2[idx]), c=col)

fig, ax = plt.subplots(figsize=(16, 9))
ax.plot(x, y_1, alpha=0.8)
ax.fill_between(x,y_1, alpha=0.5)
ax.plot(x, y_2, alpha=0.8)
#ax.scatter(x, y_3)
#ax.scatter(x, y_3, color='cyan')
#ax.set_xlim(0, 500)
#ax.set_ylim(0, 1)
ax.grid(True, axis='y')
ax.set_axisbelow(True)
#ax.set_xticks(np.arange(0, 1.001, 0.1))
#ax.set_yticks(np.arange(0, 1.001, 0.1))
month_fmt = mdates.DateFormatter('%b')
months = mdates.MonthLocator()
ax.xaxis.set_major_locator(months)
ax.xaxis.set_major_formatter(month_fmt)
#for tick in ax.get_xticklabels():
#    tick.set_rotation(65)
#myFmt = mdates.DateFormatter('%y')
#ax.xaxis.set_major_formatter(myFmt)
#fig.autofmt_xdate()
#ax.set_xlabel(x[2].year)
ax.set_ylabel('Sale in quantity')
ax.set_title(artikel_name)
ax.legend(['Actual sale', 'Predicted sale'], loc='upper right')
plt.savefig(str(artikel) + '_' + str(markt) + '_sales_2018.png', transparent=True)
plt.show()