# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *

from IchimokuCloudCrossOverAlphaModel import IchimokuCloudCrossOverAlphaModel
from BigEnergyUniverseSelectionModel import BigEnergyUniverseSelectionModel

class VerticalQuantumAtmosphericScrubbers(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 6, 5)
        self.SetEndDate(2020, 6, 15)
        self.SetCash(1000000)
        
        self.SetUniverseSelection(BigEnergyUniverseSelectionModel())
        self.UniverseSettings.Resolution = Resolution.Hour

        self.SetAlpha(IchimokuCloudCrossOverAlphaModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
#        self.SetExecution(ImmediateExecutionModel())
        
        self.SetBenchmark("QQQ")

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

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 SelectCoarse(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.HasFundamentalData ]
    
        
    def SelectFine(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.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Technology ]
        sorted_by_market_cap = sorted(energy_stocks, key=lambda x: x.MarketCap, reverse=True)
        return [ x.Symbol for x in sorted_by_market_cap[:self.fine_size] ]

from QuantConnect.Indicators import IchimokuKinkoHyo

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.ContainsKey(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(data[symbol].Price)
            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 OnSecuritiesChanged(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.AddedSecurities:
            symbol = security.Symbol
            self.symbol_data_by_symbol[symbol] = SymbolData(symbol, algorithm)
        
        for security in changes.RemovedSecurities:
            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.WarmUpPeriod + 1, Resolution.Hour).loc[symbol]
        for idx, row in history.iterrows():
            if self.ichimoku.IsReady:
                self.previous_location = self.get_location(row.close)
            
            tradebar = TradeBar(idx, symbol, row.open, row.high, row.low, row.close, row.volume)
            self.ichimoku.Update(tradebar)
        
    def get_location(self, price):
        """
        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 = price
        
        senkou_span_a = self.ichimoku.SenkouA.Current.Value
        senkou_span_b = self.ichimoku.SenkouB.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