#region imports
from AlgorithmImports import *

from QuantConnect.Indicators import IchimokuKinkoHyo
#endregion


class IchimokuCloudCrossOverAlphaModel(AlphaModel):
    """
    This class emits insights to hold a long (short) position after the chikou line of a
    security's Ichimoku Cloud crosses over (under) the top (bottom) of the cloud.
    """
    _symbol_data_by_symbol = {}
    
    def update(self, algorithm, data):
        """
        Called each time our alpha model receives a new data slice.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - data
            A data structure for all of an algorithm's data at a single time step
        
        Returns a list of Insights to the portfolio construction model.
        """
        insights = []

        for symbol, symbol_data in self._symbol_data_by_symbol.items():
            if not data.contains_key(symbol) or data[symbol] is None:
                continue
            
            # Update indicator with the latest TradeBar
            symbol_data.ichimoku.update(data[symbol])
            
            # Determine insight direction
            current_location = symbol_data.get_location()
            if symbol_data.previous_location is not None: # Indicator is ready
                if symbol_data.previous_location != 1 and current_location == 1:
                    symbol_data.direction = InsightDirection.UP
                if symbol_data.previous_location != -1 and current_location == -1:
                    symbol_data.direction = InsightDirection.DOWN
            
            symbol_data.previous_location = current_location
            
            # Emit insight
            if symbol_data.direction:
                insight = Insight.price(symbol, timedelta(days=1), symbol_data.direction)
                insights.append(insight)
        
        return insights
        
    def on_securities_changed(self, algorithm, changes):
        """
        Called each time our universe has changed.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - changes
            The additions and subtractions to the algorithm's security subscriptions
        """
        for security in changes.added_securities:
            symbol = security.symbol
            self._symbol_data_by_symbol[symbol] = SymbolData(symbol, algorithm)
        
        for security in changes.removed_securities:
            self._symbol_data_by_symbol.pop(security.symbol, None)


class SymbolData:
    """
    This class is used to store information on each security in the universe. It is
    responsible for initializing and warming up the Ichimoku indicator and determining
    the position of the chikou line in respect to the cloud.
    """
    previous_location = None
    direction = None
    
    def __init__(self, symbol, algorithm):
        """
        Input:
         - symbol
            Symbol of the security
         - algorithm
            Algorithm instance running the backtest
        """
        # Create Ichimoku indicator
        self.ichimoku = IchimokuKinkoHyo()

        # Warm up indicator
        history = algorithm.history(symbol, self.ichimoku.warm_up_period + 1, Resolution.DAILY).loc[symbol]
        for idx, row in history.iterrows():
            if self.ichimoku.is_ready:
                self.previous_location = self.get_location()
            
            tradebar = TradeBar(idx, symbol, row.open, row.high, row.low, row.close, row.volume)
            self.ichimoku.update(tradebar)
        
    def get_location(self):
        """
        Determines the location of the chikou line in respect to the cloud.
        
        Returns an integer in the interval [-1, 1], representing the location. 
        1 => Above cloud; 0 => Inside cloud; -1 => Below cloud
        """
        chikou = self.ichimoku.chikou.current.value
        
        senkou_span_a = self.ichimoku.senkou_a.current.value
        senkou_span_b = self.ichimoku.senkou_b.current.value
        cloud_top = max(senkou_span_a, senkou_span_b)
        cloud_bottom = min(senkou_span_a, senkou_span_b)
        
        if chikou > cloud_top:
            return 1    # Above cloud
        if chikou < cloud_bottom:
            return -1   # Below cloud
            
        return 0        # Inside cloud
        

#region imports
from AlgorithmImports import *

from alpha import IchimokuCloudCrossOverAlphaModel
from universe import BigEnergyUniverseSelectionModel
#endregion


class VerticalQuantumAtmosphericScrubbers(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2015, 8, 5)
        self.set_end_date(2020, 8, 5)
        self.set_cash(1000000)
        
        self.set_universe_selection(BigEnergyUniverseSelectionModel())
        self.universe_settings.resolution = Resolution.DAILY

        self.set_alpha(IchimokuCloudCrossOverAlphaModel())

        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        
        self.set_execution(ImmediateExecutionModel())
        
        self.set_benchmark("XLE")
        

#region imports
from AlgorithmImports import *

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
#endregion


class BigEnergyUniverseSelectionModel(FundamentalUniverseSelectionModel):
    """
    This universe selection model contain the 10 largest securities in the energy sector.
    """
    def __init__(self, fine_size=10):
        self._fine_size = fine_size
        self._month = -1
        super().__init__(True)

    def select_coarse(self, algorithm, coarse):
        """
        Coarse universe selection is called each day at midnight.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - coarse
            List of CoarseFundamental objects
            
        Returns the symbols that have fundamental data.
        """
        if algorithm.time.month == self._month:
            return Universe.UNCHANGED
        return [ x.symbol for x in coarse if x.has_fundamental_data ]
    
    def select_fine(self, algorithm, fine):
        """
        Fine universe selection is performed each day at midnight after `SelectCoarse`.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - fine
            List of FineFundamental objects that result from `SelectCoarse` processing
        
        Returns a list of symbols that are in the energy sector and have the largest market caps.
        """
        self._month = algorithm.time.month
        energy_stocks = [ f for f in fine if f.asset_classification.morningstar_sector_code == MorningstarSectorCode.ENERGY ]
        sorted_by_market_cap = sorted(energy_stocks, key=lambda x: x.market_cap, reverse=True)
        return [ x.symbol for x in sorted_by_market_cap[:self._fine_size] ]