# 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 AlgorithmImports import *

### <summary>
### Algorithm asserting that the consolidators are removed from the SubscriptionManager.
### This makes sure that we don't lose references to python consolidators when
### they are wrapped in a DataConsolidatorPythonWrapper.
### </summary>
class ManuallyRemovedConsolidatorsAlgorithm(QCAlgorithm):

 def initialize(self) -> None:
 self.set_start_date(2013, 10, 7)
 self.set_end_date(2013, 10, 11)

 self.consolidators = []

 self.spy = self.add_equity("SPY").symbol

 # Case 1: consolidator registered through a RegisterIndicator() call,
 # which will wrap the consolidator instance in a DataConsolidatorPythonWrapper
 consolidator = self.resolve_consolidator(self.spy, Resolution.MINUTE)
 self.consolidators.append(consolidator)

 indicator_name = self.create_indicator_name(self.spy, "close", Resolution.MINUTE)
 identity = Identity(indicator_name)
 self.indicator = self.register_indicator(self.spy, identity, consolidator)

 # Case 2: consolidator registered directly through the SubscriptionManager
 consolidator = self.resolve_consolidator(self.spy, Resolution.MINUTE)
 self.consolidators.append(consolidator)
 self.subscription_manager.add_consolidator(self.spy, consolidator)

 # Case 3: custom python consolidator not derived from IDataConsolidator
 consolidator = CustomQuoteBarConsolidator(timedelta(hours=1))
 self.consolidators.append(consolidator)
 self.subscription_manager.add_consolidator(self.spy, consolidator)

 def on_end_of_algorithm(self) -> None:
 # Remove the first consolidator

 for i in range(3):
 consolidator = self.consolidators[i]
 self.remove_consolidator(consolidator, expected_consolidator_count=2 - i)

 def remove_consolidator(self, consolidator: IDataConsolidator, expected_consolidator_count: int) -> None:
 self.subscription_manager.remove_consolidator(self.spy, consolidator)

 consolidator_count = sum(s.consolidators.count for s in self.subscription_manager.subscriptions)
 if consolidator_count != expected_consolidator_count:
 raise AssertionError(f"Unexpected number of consolidators after removal. "
 f"Expected: {expected_consolidator_count}. Actual: {consolidator_count}")

class CustomQuoteBarConsolidator(PythonConsolidator):
 '''A custom quote bar consolidator that inherits from PythonConsolidator and implements
 the IDataConsolidator interface, it must implement all of IDataConsolidator. Reference
 PythonConsolidator.cs and DataConsolidatorPythonWrapper.PY for more information.

 This class shows how to implement a consolidator from scratch in Python, this gives us more
 freedom to determine the behavior of the consolidator but can't leverage any of the built in
 functions of an inherited class.

 For this example we implemented a Quotebar from scratch'''

 def __init__(self, period):

 #IDataConsolidator required vars for all consolidators
 self.consolidated = None #Most recently consolidated piece of data.
 self.working_data = None #Data being currently consolidated
 self.input_type = QuoteBar #The type consumed by this consolidator
 self.output_type = QuoteBar #The type produced by this consolidator

 #Consolidator Variables
 self.period = period

 def update(self, data):
 '''Updates this consolidator with the specified data'''

 #If we don't have bar yet, create one
 if self.working_data is None:
 self.working_data = QuoteBar(data.time,data.symbol,data.bid,data.last_bid_size,
 data.ask,data.last_ask_size,self.period)

 #Update bar using QuoteBar's update()
 self.working_data.update(data.value, data.bid.close, data.ask.close, 0,
 data.last_bid_size, data.last_ask_size)

 def scan(self, time):
 '''Scans this consolidator to see if it should emit a bar due to time passing'''

 if self.period is not None and self.working_data is not None:
 if time - self.working_data.time >= self.period:

 #Trigger the event handler with a copy of self and the data
 self.on_data_consolidated(self, self.working_data)

 #Set the most recent consolidated piece of data and then clear the working_data
 self.consolidated = self.working_data
 self.working_data = None