from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
from datetime import datetime
from datetime import timedelta
import pandas as pd
import numpy as np

class TransdimensionalTachyonCompensator(QCAlgorithm):
    month = 0
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 10, 27)
        self.SetCash(100000)  # Set Strategy Cash
        self.trading_symbols = []
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction) 
        self.lookback_period = 60
        self.spy = self.AddEquity("SPY").Symbol
        # self.Train(self.DateRules.MonthStart(), self.TimeRules.AfterMarketOpen(self.spy, 0), self.train)
        self.Schedule.On(self.DateRules.MonthStart(), self.TimeRules.AfterMarketOpen(self.spy, 0), self.train)

    def CoarseSelectionFunction(self, coarse):
        if self.Time.month == self.month:
            return Universe.Unchanged
        self.month = self.Time.month
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        self.trading_symbols = [ x.Symbol for x in sortedByDollarVolume if x.HasFundamentalData ][:5]
        return self.trading_symbols
        
    def train(self):
        if not self.trading_symbols: return
        
        emas_by_symbol = {}
        
        history = self.History(self.trading_symbols , self.lookback_period , Resolution.Daily)
        for symbol in self.trading_symbols:
            if history.empty or 'close' not in history.columns:
                return
            price_changes = history.close.unstack(0).diff().dropna()
            if symbol not in price_changes.columns:
                continue
            emas_by_symbol[symbol] = {'up': ExponentialMovingAverage(14), 'down': ExponentialMovingAverage(14)}
            for time, price_change in price_changes[symbol].iteritems():
                emas_by_symbol[symbol]['up'].Update(time, price_change if price_change > 0 else 0)
                emas_by_symbol[symbol]['down'].Update(time, abs(price_change) if price_change < 0 else 0)
            self.Plot("EMA-UP", str(symbol), emas_by_symbol[symbol]['up'].Current.Value)
            self.Plot("EMA-Down", str(symbol), emas_by_symbol[symbol]['down'].Current.Value)