lib/blazer/result.rb in blazer-2.2.1 vs lib/blazer/result.rb in blazer-2.2.2

@@ -83,14 +83,48 @@
 
     def forecastable?
       @forecastable ||= Blazer.forecasting && column_types == ["time", "numeric"] && @rows.size >= 10
     end
 
+    # TODO cache it?
+    # don't want to put result data (even hashed version)
+    # into cache without developer opt-in
     def forecast
-      # TODO cache it?
-      # don't want to put result data (even hashed version)
-      # into cache without developer opt-in
-      forecast = Trend.forecast(Hash[@rows], count: 30)
+      count = (@rows.size * 0.25).round.clamp(30, 365)
+
+      case Blazer.forecasting
+      when "prophet"
+        require "prophet"
+
+        df =
+          Daru::DataFrame.new(
+            "ds" => @rows.map { |r| r[0] },
+            "y" => @rows.map { |r| r[1] }
+          )
+
+        # TODO determine frequency
+        freq = "D"
+
+        m = Prophet.new
+        m.fit(df)
+        future = m.make_future_dataframe(periods: count, freq: freq, include_history: false)
+        fcst = m.predict(future)
+        ds = fcst["ds"]
+        if @rows[0][0].is_a?(Date)
+          ds = ds.map { |v| v.to_date }
+        end
+        forecast = ds.zip(fcst["yhat"]).to_h
+      else
+        require "trend"
+
+        forecast = Trend.forecast(Hash[@rows], count: count)
+      end
+
+      # round integers
+      if @rows[0][1].is_a?(Integer)
+        forecast = forecast.map { |k, v| [k, v.round] }.to_h
+      end
+
       @rows.each do |row|
         row[2] = nil
       end
       @rows.unshift(*forecast.map { |k, v| [k, nil, v] })
       @columns << "forecast"