# region imports
from AlgorithmImports import *
from gplearn.genetic import SymbolicRegressor, SymbolicTransformer
import joblib
# endregion

class GPlearnExampleAlgorithm(QCAlgorithm):
    
    def initialize(self):
        self.set_start_date(2022, 7, 4)
        self.set_end_date(2022, 7, 8)
        self.set_cash(100000)
        self.symbol = self.add_equity("SPY", Resolution.DAILY).symbol
        
        training_length = 252*2
        self.training_data = RollingWindow[float](training_length)
        history = self.history[TradeBar](self.symbol, training_length, Resolution.DAILY)
        for trade_bar in history:
            self.training_data.add(trade_bar.close)

        transformer_model_key = "Transformer"
        regressor_model_key = "Regressor"
        if self.object_store.contains_key(transformer_model_key) and self.object_store.contains_key(regressor_model_key):
            transformer_file_name = self.object_store.get_file_path(transformer_model_key)
            regressor_file_name = self.object_store.get_file_path(regressor_model_key)

            self.gp_transformer = joblib.load(transformer_file_name)
            self.model = joblib.load(regressor_file_name)

        else:
            function_set = ['add', 'sub', 'mul', 'div',
                'sqrt', 'log', 'abs', 'neg', 'inv',
                'max', 'min']
            self.gp_transformer = SymbolicTransformer(function_set=function_set)
            self.model = SymbolicRegressor()
            self.train(self.my_training_method)

        self.train(self.date_rules.every(DayOfWeek.SUNDAY), self.time_rules.at(8,0), self.my_training_method)
        
    def get_features_and_labels(self, n_steps=5):
        training_df = list(self.training_data)[::-1]
        daily_pct_change = ((np.roll(training_df, -1) - training_df) / training_df)[:-1]

        features = []
        labels = []
        for i in range(len(daily_pct_change)-n_steps):
            features.append(daily_pct_change[i:i+n_steps])
            labels.append(daily_pct_change[i+n_steps])
        features = np.array(features)
        labels = np.array(labels)

        return features, labels

    def my_training_method(self):
        features, labels = self.get_features_and_labels()

        # Feature engineering
        self.gp_transformer.fit(features, labels)
        gp_features = self.gp_transformer.transform(features)
        new_features = np.hstack((features, gp_features))

        # Fit the regression model with transformed and raw features.
        self.model.fit(new_features, labels)

    def on_data(self, slice):
        features, _ = self.get_features_and_labels()

        # Get transformed features
        gp_features = self.gp_transformer.transform(features)
        new_features = np.hstack((features, gp_features))

        # Get next prediction
        prediction = self.model.predict(new_features)
        prediction = float(prediction.flatten()[-1])

        if prediction > 0:
            self.set_holdings(self.symbol, 1)
        elif prediction < 0:            
            self.set_holdings(self.symbol, -1)

    def on_end_of_algorithm(self):
        transformer_model_key = "Transformer"
        regressor_model_key = "Regressor"
        transformer_file_name = self.object_store.get_file_path(transformer_model_key)
        regressor_file_name = self.object_store.get_file_path(regressor_model_key)
        joblib.dump(self.gp_transformer, transformer_file_name)
        joblib.dump(self.model, regressor_file_name)
        self.object_store.save(transformer_model_key)
        self.object_store.save(regressor_model_key)