from fbprophet import Prophet

class ARIMADailyAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(datetime.today())
        self.SetCash(1000000)
        self.AddEquity("TLT", Resolution.Daily)
        # self.SetBenchmark("SPY")
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        self.AddAlpha(ProphetModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(lambda time: None))

class ProphetModel(AlphaModel):
    """
    Facebook Prophet time series prediction trial algorithm.
    """

    def __init__(self):

        # Prophet training past:
        N_YEARS = 5
        self.training_period = 365*N_YEARS
        RETRAIN = 6

        # Prediction Period:
        self.prediction_period = 5
        self.holding_period = timedelta(days = self.prediction_period)
        self.retrain_period = RETRAIN*self.prediction_period
        self.operation_count = False
        
        # Alpha Model:
        self.price_model = False
        self.Name = 'Old Man Prophet'

    def Update(self, algorithm, data):

        insights = []
        if data.HasData == False: return []

        if not self.operation_count or self.operation_count%self.retrain_period == 0:
            history = algorithm.History(algorithm.Securities.Keys,
                                        self.training_period,
                                        Resolution.Daily)
            time_series = self.create_time_series(history)
            self.price_model = Prophet()
            time_series.reset_index(inplace=True)
            time_series.columns = ['ds', 'y']
            self.price_model.fit(time_series)
            algorithm.Debug('Training on:' + str(algorithm.Time))

        for security in algorithm.ActiveSecurities.Values:
            if not security.Invested:
                price_forecast = self.price_model.make_future_dataframe(periods=self.prediction_period, freq='d')
                price_forecast = self.price_model.predict(price_forecast)
                target_price = price_forecast.iloc[-1]['yhat']
                
                if data[security.Symbol] is None: return []
                current_price = data[security.Symbol].Close

                if target_price > current_price: direction = InsightDirection.Up
                elif target_price < current_price: direction = InsightDirection.Down
                else: direction = InsightDirection.Flat
                
                price_difference = target_price - current_price

                insight = Insight(security.Symbol,
                                  self.holding_period,
                                  InsightType.Price,
                                  direction,
                                  price_difference, int(True),
                                  self.Name, int(True))
                pred_dict = {InsightDirection.Up: 'Up', InsightDirection.Down: 'Down', InsightDirection.Flat: 'Flat'}
                algorithm.Debug('Predicting: ' + str(pred_dict[direction]))
                insights.append(insight)
        
        self.operation_count += 1
        return insights

    def create_time_series(self, df, column='close', freq='d'):
        df = df.reset_index()
        time_series = df[[column]].set_index(df['time']).asfreq(freq).ffill()
        return time_series