#region imports
from AlgorithmImports import *
#endregion

class PortfolioHelper():
    
    # ==========================================
    # Constructor. Accepts algo Object
    # ==========================================
    
    def __init__(self, algo, maxOpenPositions, debugMode=False):

        # Algo reference
        # ----------------------------------------
        self.algo              = algo
        self.debugMode       = debugMode
        self.maxOpenPositions  = maxOpenPositions
        

    ## How many available spots are available in the portfolio
    ## ------------------------------------------------------------	
    @property
    def PortfolioCapacity(self):
        return int(max(0,(self.maxOpenPositions - self.NumHoldings)))
    

    ## Check if we are already holding the max # of open positions.	
    ## ------------------------------------------------------------	
    @property	
    def PortfolioAtCapacity(self):	        
        numHoldings = self.NumHoldings
        atCapacity  = numHoldings >= self.maxOpenPositions	

        if( self.debugMode):
            self.algo.Plot("Num of Positions", 'Line', numHoldings)	

            # # Plot max held ever	
            # # ------------------	
            # if( not hasattr(self,"maxHeldEver")):	
            #     self.maxHeldEver = numHoldings	
            # else:	
            #     if(self.maxHeldEver < numHoldings):	
            #         debugstop = "here"	
            #         self.algo.Plot("Max Positions", 'Line', numHoldings)	
            # self.maxHeldEver = max(self.maxHeldEver,numHoldings)	
        
        return atCapacity


    # =====================================
    @property
    def NumHoldings(self):
        numHoldings = len([x.Key for x in self.algo.Portfolio if x.Value.Invested])
        return numHoldings 

    # =====================================
    def InvestedInSymbol (self, symbol):
        return (self.algo.Portfolio.ContainsKey(symbol)) and (self.algo.Portfolio[symbol].Invested)
    
    # =====================================
    def ProfitLossMsgForSymbol (self, symbol):
        if(self.InvestedInSymbol(symbol)):            
            profitPct   = self.algo.Portfolio[symbol].UnrealizedProfitPercent
            winlossMsg  =  f"{round(profitPct*100,2)}% {( 'Win' if (profitPct > 0) else 'Loss')}" 
            return winlossMsg
        else:
            return None

#region imports
from AlgorithmImports import *
#endregion
##########################################################################
# The Ichimoku Breakout Testbed
# ----------------------------------------------------
# 
# Entry:
# -------
# Ichimoku cloud breakout 
#  
#
# Exit:
# -------
# When price dips below Tenkan
#
# @shock_and_awful
##########################################################################

from PortfolioHelper import *
class IchimokuUniverse(QCAlgorithm):
    
    # =====================================
    def Initialize(self):
        self.SetStartDate(2020, 6, 1)
        # self.SetEndDate(2020, 12, 30)
        self.SetCash(100000)
        self.PortfolioHelper = PortfolioHelper(self, 5, debugMode=True) # 5 max open positions
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction) 
        self.symDataDict = { }

    # =====================================
    def OnData(self, dataSlice):
        for symbol in dataSlice.Keys:
            if symbol in self.symDataDict:
                symbolData = self.symDataDict[symbol] 
                symbolData.UpdateIndicatorsWithBars(dataSlice[symbol])
                if( symbolData.ExitSignalFired() ):
                    msg = self.PortfolioHelper.ProfitLossMsgForSymbol(symbol)
                    self.Liquidate(symbol, tag=msg)                    
                    
        
    
    # =====================================
    def CoarseSelectionFunction(self, universe):  
        coarseuniverse = sorted(universe, key=lambda c: c.DollarVolume, reverse=True)  
        coarseuniverse = [c for c in coarseuniverse if c.Price > 30][:30]

        return [x.Symbol for x in coarseuniverse]


    # =====================================
    def FineSelectionFunction(self, universe):
        
        # If portfolio is at capacity, skip universe selection
        # ----------------------------------------------------
        if self.PortfolioHelper.PortfolioAtCapacity:
            return []

        fineUniverse = [x for x in universe if x.SecurityReference.IsPrimaryShare
                        and x.SecurityReference.SecurityType == "ST00000001"
                        and x.SecurityReference.IsDepositaryReceipt == 0
                        and x.CompanyReference.IsLimitedPartnership == 0]

        selected     = [] 
        queuedAssetCount = 0       
        for element in fineUniverse:  
            symbol = element.Symbol
            
            # Store data for this symbol, seed it with some history     
            # -----------------------------------------------------            
            if symbol not in self.symDataDict:
                
                # Fetch enough historical bars with which to seed the symbol's indicators
                history = self.History(symbol, 201, Resolution.Daily) 
                self.symDataDict[symbol] = SymbolData(symbol, history, self) 

            # Stop screening If we've queued up more assets than we have capacity for 
            # ------------------------------------------------------------------------
            if(queuedAssetCount >= self.PortfolioHelper.PortfolioCapacity):
                del self.symDataDict[symbol]
            else:
                # If the entry signal fired for this symbol, add to universe     
                # --------------------------------------------------------------       
                if  self.symDataDict[symbol].IsReady() and self.symDataDict[symbol].EntrySignalFired():
                    selected.append(symbol)
                    queuedAssetCount += 1
                    
                else:
                    del self.symDataDict[symbol]
        
        return selected
        
    # 
    # =====================================
    def OnSecuritiesChanged(self, changes):
        
        # The trade actually takes place here, when 
        # the symbol passes coarse and fine filters
        # and is added to our dictionary of securities
        # --------------------------------------------
        for security in changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 0.1)

            
##################################
# 
# SymbolData Class
#
##################################
class SymbolData():
    
    # ==========================================
    # Constructor. Accepts History array
    # ==========================================
    def __init__(self, theSymbol, history, algo):

        # Algo / Symbol / Price reference
        # ----------------------------------------
        self.algo                 = algo
        self.symbol               = theSymbol
        self.lastPrice            = 0

        # Initialize our Ichi and EMA indicators
        # ----------------------------------------
        self.ichimoku             = IchimokuKinkoHyo(9, 26, 26, 52, 26, 26)
        self.emaThirtyFour        = ExponentialMovingAverage(34)
        self.emaTwoHundred        = ExponentialMovingAverage(200)
        
        # Initialize vars for tracking indicator state
        # ----------------------------------------------
        self.IsTKAboveCloud       = False    
        self.IsPriceBelowKJ       = False    
        self.IsBelowCloudWindow   = RollingWindow[Boolean](2)
        self.LastPriceWindow      = RollingWindow[float](2)  
        self.ichimoku.Updated    += self.UpdateIchimokuWindows
    
        # Loop over the history data and update the Ichimoku indicator
        # -------------------------------------------------------------
        if history.empty or 'close' not in history.columns:
            return

        for index, row in history.loc[theSymbol].iterrows():
            tradeBar        = TradeBar()
            tradeBar.Close  = row['close']
            tradeBar.Open   = row['open']
            tradeBar.High   = row['high']
            tradeBar.Low    = row['low']
            tradeBar.Volume = row['volume']
            tradeBar.Time   = index
            tradeBar.Period = timedelta(1)
            tradeBar.Symbol = theSymbol
            
            
            self.UpdateIndicatorsWithBars(tradeBar)
            
            
    # =====================================    
    def UpdateIndicatorsWithBars(self, tradeBar):
        if((tradeBar is not None) and (self.ichimoku is not None)): 
            self.lastPrice = tradeBar.Close
            self.LastPriceWindow.Add( tradeBar.Close )
            self.ichimoku.Update(tradeBar)
            self.emaThirtyFour.Update(tradeBar.Time, tradeBar.Close)
            self.emaTwoHundred.Update(tradeBar.Time, tradeBar.Close)
                

    # =====================================    
    def UpdateIchimokuWindows(self, sender, updated):
        senkouA      = sender.SenkouA.Current.Value
        senkouB      = sender.SenkouB.Current.Value
        cloudTop     = max(senkouA, senkouB)
        belowCloud   = updated.Price <= cloudTop
        
        self.lastPrice      = updated.Price
        self.IsPriceBelowKJ = updated.Price < sender.Kijun.Current.Value
        self.IsTKAboveCloud = (sender.Tenkan.Current.Value > cloudTop)
        self.IsBelowCloudWindow.Add( belowCloud )

    # =====================================            
    def IsReady(self):
        return self.ichimoku.IsReady  and \
                self.emaThirtyFour.IsReady and \
                self.emaTwoHundred.IsReady
        

    # =====================================
    def EntrySignalFired(self):

        # return true if we broke out of the cloud, and TK is above the cloud
        # --------------------------------------------------------------------
        if( ( self.IsBelowCloudWindow.Count > 1 ) and \
                ( self.IsBelowCloudWindow[1] and (not self.IsBelowCloudWindow[0] )) and \
                ( self.IsTKAboveCloud and (self.LastPriceWindow[0] > self.emaThirtyFour.Current.Value > self.emaTwoHundred.Current.Value))):        
                # ( self.IsTKAboveCloud )):
            return True
        return False
        
        
    # =====================================
    def ExitSignalFired(self):
        # return true if price dipped below the TK    
        if self.IsPriceBelowKJ:         
            return True
        return False