from collections import defaultdict
from dataManager import DataManager
from tradeManager import TradeManager
from strategy import Strategy

class Manager():
    def __init__( self, algo ):
        self._algo = algo
        self.dataManager = DataManager( algo )
        self.tradeManager = TradeManager( algo, self.dataManager )
        self.strategy = Strategy(self.tradeManager, self.dataManager )
        
    def addUniverse( self, ticker ):
        self.dataManager.addNewTicker(ticker)
                
        
    def run( self ):
        '''
        this should run every tick
        '''
        self.dataManager.run()
        self.tradeManager.run()
        self.strategy.run()

from manager import Manager

        
class BasicFW(QCAlgorithm):


    def Initialize(self):
        self.SetStartDate(2020, 9, 1)
        #self.SetEndDate(2020, 8, 1)
        self.SetCash(10000000)
        resolution = Resolution.Hour
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.UniverseSettings.Resolution = resolution
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction )
        self.AddEquity("SPY", resolution)
        # Consolidate 1min SPY -> 1-Day Bars
        self.manager = Manager( self )
        self.universe = []
        
    

    def OnData( self, slice ):
        self.manager.run()

        
    def OnSecuritiesChanged(self, changes):
        #On adding securities, creating SymbolData for each.
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.manager.addUniverse( symbol )

    def CoarseSelectionFunction( self, coarse ):
        top = [x for x in coarse if x.HasFundamentalData
                           and 100 > x.Price > 30
                           and  x.Volume > 1000000
                           and  x.DollarVolume > 1000000]
        return [x.Symbol for x in top]
        
    def FineSelectionFunction( self, fine ):
        top = [x for x in fine if x.OperationRatios.QuickRatio.ThreeMonths >= 1.0
                           and x.OperationRatios.TotalDebtEquityRatioGrowth.OneYear <= 1.0
                           and  x.ValuationRatios.SustainableGrowthRate > 0.07
                           and  x.ValuationRatios.PEGRatio < 1]
        return [x.Symbol for x in top][:10]

class TradeManager():
    def __init__( self, algo, dataManager, minutesForTrade = 60*3, maxOpenTrades = 5):
        self._algo = algo
        self.tradeTime = timedelta( minutes = minutesForTrade )
        self.tradeDetails = {}
        self._dataManager = dataManager
        self._maxOpenTrades = maxOpenTrades
        
    def openTrade( self, ticker, numOfShares, sl = None, tp = None ):
        numOfTrades = sum( 1 for inst in self._algo.Portfolio if inst.Value.Invested )
        if numOfTrades < self._maxOpenTrades:
            self._algo.MarketOrder(ticker, numOfShares  )
            self.tradeDetails[ ticker ] = {
                'time' : self._algo.Time,
                'sl' : sl,
                'tp' : tp
            }
        
    def run( self ):
        '''
        this should run every tick
        '''
        self.manageOpenTrades()
    
    
    def closeTrade( self, ticker ):
        self._algo.Liquidate( ticker )
        self.tradeDetails.pop( ticker )
        
    def rollingWindow( self, ticker ):
        return self._dataManager.getData( ticker )
        
    def manageOpenTrades( self ):
        for inst in self._algo.Portfolio:
            symbol = inst.Value.Symbol
            if inst.Value.Invested and symbol in self.tradeDetails:
                
                #Time limit
                if self.tradeDetails[ symbol ]['time'] + self.tradeTime < self._algo.Time: 
                    self.closeTrade( symbol )
                
                if symbol in self.tradeDetails:
                    # Long position
                    if inst.Value.Quantity > 0:
    
                        if self.tradeDetails[ symbol ]['sl'] is not None and self.rollingWindow( symbol )[0].Close <= self.tradeDetails[ symbol ][ 'sl' ] :
                            self.closeTrade( symbol )
                            
                        elif self.tradeDetails[ symbol ]['tp'] is not None and self.rollingWindow( symbol )[0].Close >= self.tradeDetails[ symbol ][ 'tp' ] :
                            self.closeTrade( symbol )
                            
                    # short position
                    if inst.Value.Quantity < 0:
                        if self.tradeDetails[ symbol ]['sl'] is not None and self.rollingWindow( symbol )[0].Close >= self.tradeDetails[ symbol ][ 'sl' ] :
                            self.closeTrade( symbol )
                            
                        elif self.tradeDetails[ symbol ]['tp'] is not None and self.rollingWindow( symbol )[0].Close <= self.tradeDetails[ symbol ][ 'tp' ] :
                            self.closeTrade( symbol )

from collections import defaultdict
import statistics
DATA_SIZE = 60*10


class DataManager():
    def __init__( self, algo, minutesToStore = 60*10 ):
        
        self._algo = algo
        self._dataTime = timedelta( minutes = minutesToStore ) #for how long to store the data
        self._data = defaultdict( lambda : RollingWindow[TradeBar]( DATA_SIZE ) )
        now = self._algo.Time
        self._dataStartTime = defaultdict( lambda : now )
        
    def run( self ):
        '''
        this should run every tick
        '''
        self.collectSingleBars()
        self.removeStaleData()
        
    def getData( self, ticker ):
        if ticker in self._data:
            return self._data[ ticker ]
        
    def getAllTickers( self ):
        return list( self._data.keys() )
    
    def collectSingleBars( self ):
        for ticker in self._data.keys():
            if not self._algo.CurrentSlice.ContainsKey(  ticker ) or self._algo.CurrentSlice[ ticker ] is None:
                #self._algo.Debug(ticker.Value + ' Not available')
                continue    
            self._data[ ticker ].Add( self._algo.CurrentSlice[ ticker ] )
    
    def addNewTicker( self, ticker ):
        self._dataStartTime[ ticker ] = self._algo.Time
        if self._data[ ticker ].Count == 0:
            df = self._algo.History( ticker, DATA_SIZE, Resolution.Minute)
            for barTuple in df.itertuples():
                data = barTuple[1]
                tradeBar = TradeBar()
                tradeBar.Close = barTuple.close
                tradeBar.Open = barTuple.open
                tradeBar.High = barTuple.high
                tradeBar.Low = barTuple.low
                tradeBar.Volume = getattr( barTuple, 'volume', 0 )
                tradeBar.Time = barTuple.Index[1]
                tradeBar.Symbol = barTuple.Index[0]
                tradeBar.Period = Resolution.Minute
                self._data[ ticker ].Add( tradeBar )
            
    def removeStaleData( self ):
        #return
        toRemove = []
        for ticker, startTime in self._dataStartTime.items():
            if startTime + self._dataTime < self._algo.Time and self._algo.Portfolio[ ticker ].Invested == False :
                toRemove.append( ticker)
        for ticker in toRemove:
            self._dataStartTime.pop( ticker, None )
            self._data.pop( ticker, None )

from dataManager import DataManager
from tradeManager import TradeManager

class Strategy():
    def __init__( self, tradeManager, dataManager ):
        self._tradeManager = tradeManager
        self._dataManager = dataManager
        
    def run( self ):
        '''
        this should run every tick
        '''
        self.openTradeBySignal()
        
    def rollingWindow( self, ticker ):
        return self._dataManager.getData( ticker )
        
    def openTradeBySignal( self ):
        winSize = 15
        for ticker in self._dataManager.getAllTickers():
            lowest =  lambda x : min( self.rollingWindow( ticker )[i].Low for i in range( x * winSize, (x+1) * winSize ) )
            
            lastClose = self.rollingWindow( ticker )[0].Close
            
            if lowest(30) < lastClose < lowest(10) :
                self._tradeManager.openTrade( ticker, 100, sl = lowest(30)  , tp = lowest(10) )