From cafbb058b69489690b189809b97076e67f624e84 Mon Sep 17 00:00:00 2001
From: Jamie <jmaslek11@gmail.com>
Date: Mon, 10 May 2021 12:40:52 -0400
Subject: [PATCH] Finish first refactor for LSTM model including loops and
 double checking the set end date.

---
 .../neural_networks_view.py                   |  6 ++-
 .../prediction_techniques/pred_helper.py      | 50 ++++++++++++++-----
 2 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/gamestonk_terminal/prediction_techniques/neural_networks_view.py b/gamestonk_terminal/prediction_techniques/neural_networks_view.py
index 57a6122ee41..a4afc318bec 100644
--- a/gamestonk_terminal/prediction_techniques/neural_networks_view.py
+++ b/gamestonk_terminal/prediction_techniques/neural_networks_view.py
@@ -1054,7 +1054,10 @@ def lstm(other_args: List[str], s_ticker: str, df_stock: pd.DataFrame):
             forecast_data_input,
             dates_forecast_input,
             scaler,
+            is_error
         ) = prepare_scale_train_valid_test(df_stock["5. adjusted close"], ns_parser)
+        if is_error:
+            return
         print(
             f"Training on {X_train.shape[0]} samples.  Using {X_valid.shape[0]} samples for validation."
             f" {ns_parser.n_loops} loops"
@@ -1094,12 +1097,13 @@ def lstm(other_args: List[str], s_ticker: str, df_stock: pd.DataFrame):
         print(forecast_data_df)
         plot_data_predictions(
             df_stock,
-            preds.median(axis=0),
+            np.median(preds,axis=0),
             y_valid,
             y_dates_valid,
             scaler,
             f"LSTM Model on {s_ticker}",
             forecast_data_df,
+            ns_parser.n_loops
         )
         print("")
 
diff --git a/gamestonk_terminal/prediction_techniques/pred_helper.py b/gamestonk_terminal/prediction_techniques/pred_helper.py
index e7f42abc582..e10c5625e9a 100644
--- a/gamestonk_terminal/prediction_techniques/pred_helper.py
+++ b/gamestonk_terminal/prediction_techniques/pred_helper.py
@@ -83,9 +83,11 @@ def prepare_scale_train_valid_test(
 
     if ns_parser.s_end_date:
         df_stock = df_stock[df_stock.index <= ns_parser.s_end_date]
-        if n_input_days + n_predict_days < df_stock.shape[0]:
-            print("Cannot train enough input days to predict with loaded dataframe")
-            return
+        print(df_stock.shape)
+        print(n_input_days + n_predict_days)
+        if n_input_days + n_predict_days > df_stock.shape[0]:
+            print("Cannot train enough input days to predict with loaded dataframe\n")
+            return None,None,None,None,None,None,None,None,None,None,None,True
 
     test_data = df_stock.iloc[-n_input_days:]
     train_data = df_stock.iloc[:-n_input_days]
@@ -144,6 +146,7 @@ def prepare_scale_train_valid_test(
         test_data,
         dates_test,
         scaler,
+        False
     )
 
 
@@ -159,7 +162,7 @@ def forecast(input_values, future_dates, model, scaler):
 
 
 def plot_data_predictions(
-    df_stock, preds, y_valid, y_dates_valid, scaler, title, forecast_data
+    df_stock, preds, y_valid, y_dates_valid, scaler, title, forecast_data, n_loops
 ):
 
     plt.figure(figsize=plot_autoscale(), dpi=PLOT_DPI)
@@ -172,19 +175,34 @@ def plot_data_predictions(
         label="Prices",
     )
     for i in range(len(y_valid) - 1):
+
+        y_pred = scaler.inverse_transform(preds[i].reshape(-1, 1)).ravel()
+        y_act = scaler.inverse_transform(y_valid[i].reshape(-1, 1)).ravel()
+
         plt.plot(
             y_dates_valid[i],
-            scaler.inverse_transform(preds[i].reshape(-1, 1)),
+            y_pred,
             "r",
             lw=3,
         )
         plt.fill_between(
             y_dates_valid[i],
-            scaler.inverse_transform(preds[i].reshape(-1, 1)).ravel(),
-            scaler.inverse_transform(y_valid[i].reshape(-1, 1)).ravel(),
-            color="k",
+            y_pred,
+            y_act,
+            where=(y_pred < y_act),
+            color="r",
             alpha=0.2,
         )
+        plt.fill_between(
+            y_dates_valid[i],
+            y_pred,
+            y_act,
+            where=(y_pred > y_act),
+            color="g",
+            alpha=0.2,
+        )
+
+    # Leave this one out of the loop so that the legend doesnt get overpopulated with "Predictions"
     plt.plot(
         y_dates_valid[-1],
         scaler.inverse_transform(preds[-1].reshape(-1, 1)),
@@ -200,11 +218,11 @@ def plot_data_predictions(
         alpha=0.2,
     )
     plt.axvspan(
-        df_stock.index[-1], forecast_data.index[-1], facecolor="tab:orange", alpha=0.2
+        forecast_data.index[0]-timedelta(days=1), forecast_data.index[-1], facecolor="tab:orange", alpha=0.2
     )
     _, _, ymin, ymax = plt.axis()
     plt.vlines(
-        df_stock.index[-1],
+        forecast_data.index[0]-timedelta(days=1),
         ymin,
         ymax,
         colors="k",
@@ -212,7 +230,15 @@ def plot_data_predictions(
         linestyle="--",
         color="k",
     )
-    plt.plot(forecast_data.index, forecast_data.values, "-ok", ms=5)
+    if n_loops == 1:
+        plt.plot(forecast_data.index, forecast_data.values, "-ok", ms=5, label="Forecast")
+    else:
+        plt.plot(forecast_data.index, forecast_data.median(axis=1).values, "-ok", ms=5, label="Forecast")
+        plt.fill_between(forecast_data.index,
+                         forecast_data.quantile(.25,axis=1).values,
+                         forecast_data.quantile(.75,axis=1).values,
+                         color="c",
+                         alpha = .3)
     plt.legend(loc=0)
     plt.xlim(df_stock.index[0], forecast_data.index[-1] + timedelta(days=1))
     plt.xlabel("Time")
@@ -226,7 +252,7 @@ def plot_data_predictions(
     plt.show()
     print("")
 
-
+# Deprecated
 def get_backtesting_data(
     df_stock: pd.DataFrame, parsed_end_date: datetime, pred_days: int
 ) -> Tuple[pd.DataFrame, pd.DataFrame, bool]: