#region imports
from AlgorithmImports import *
#endregion
from AlgorithmImports import *
import pytz

# Your New Python File
class DualThrustAlphaModel(AlphaModel):

    def __init__(self,
                 ph,
                 range_period,
                 resolution = Resolution.HOUR,
                 bars_to_consolidate = 1):

        self.ph = ph
        # period the range is calculated over
        self.range_period = range_period
        # initialize with empty dict.
        self._symbol_data_by_symbol = {}
        # time for bars we make the calculations on
        resolution_in_time_span =  Extensions.to_time_span(resolution)
        self.consolidator_time_span = Time.multiply(resolution_in_time_span, bars_to_consolidate)

        self.period = timedelta(minutes=60)

    def update(self, algorithm, data):
        insights = []

        for symbol, symbol_data in self._symbol_data_by_symbol.items():
            if not symbol_data.is_ready: #or symbol_data.insight_close_time_utc > algorithm.utc_time:
                continue

            holding = algorithm.portfolio[symbol]
            #price = algorithm.securities[symbol].price 

            if symbol_data.wvf > symbol_data.rangehigh: #or symbol_data.wvf > (symbol_data.std_wvf * 2):
                symbol_data.insight_close_time_utc = self.period
                insights.append(Insight.price(symbol, symbol_data.insight_close_time_utc, InsightDirection.UP))
                algorithm.debug(f"wvf: {symbol_data.wvf}")
                algorithm.debug(f"rangehigh: {symbol_data.rangehigh}")

        return insights

    def on_securities_changed(self, algorithm, changes):
        # added
        for symbol in [x.symbol for x in changes.added_securities]:
            if symbol not in self._symbol_data_by_symbol:
                # add symbol/symbol_data pair to collection
                symbol_data = self.SymbolData(algorithm, symbol, self.ph, self.range_period, self.consolidator_time_span)
                self._symbol_data_by_symbol[symbol] = symbol_data

        # removed
        for symbol in [x.symbol for x in changes.removed_securities]:
            symbol_data = self._symbol_data_by_symbol.pop(symbol, None)
            if symbol_data is None:
                algorithm.error("Unable to remove data from collection: DualThrustAlphaModel")
            else:
                # unsubscribe consolidator from data updates
                algorithm.subscription_manager.remove_consolidator(symbol, symbol_data.get_consolidator())

    class SymbolData:
    
        def __init__(self, algorithm, symbol, ph, range_period, consolidator_resolution):

            self.symbol = symbol
            self.range_window = RollingWindow[TradeBar](range_period)
            self.consolidator = TradeBarConsolidator(consolidator_resolution)
            algorithm.subscription_manager.add_consolidator(symbol, self.consolidator)
            self.wvf_window_rangehigh = RollingWindow[float](50)
            self.wvf_window_std = RollingWindow[float](20)
            self.rangehigh = 0
            self.wvf = 0
            self.std_wvf = 0
            self.insight_close_time_utc = datetime.min.replace(tzinfo=pytz.UTC)

            def on_data_consolidated(sender, consolidated):

                self.range_window.add(consolidated)

                if self.range_window.is_ready:
                    hh = max([x.high for x in self.range_window])
                    hc = max([x.close for x in self.range_window])
                    lc = min([x.close for x in self.range_window])
                    ll = min([x.low for x in self.range_window])

                    self.wvf = ((hc - consolidated.low)/hc) * 100

                    self.wvf_window_rangehigh.add(self.wvf)
                    self.wvf_window_std.add(self.wvf)
                    
                    self.rangehigh = max([x for x in self.wvf_window_rangehigh]) * ph

            self.consolidator.data_consolidated += on_data_consolidated

            history = algorithm.history[TradeBar](symbol, 100, Resolution.HOUR)
            for bar in history:
                self.consolidator.update(bar)

        # Returns the interior consolidator
        def get_consolidator(self):
            return self.consolidator
        
        @property
        def is_ready(self):
            return self.range_window.is_ready

from AlgorithmImports import *
from VIXAlpha import *

class SELLPUTFRAMEWORK(QCAlgorithm):

    def initialize(self):

        # -- STRATEGY INPUT PARAMETERS --
        self.ph = 0.85
        self.range_period = 22
        #self.bbl = 20
        #self.lb = 50
        
        #self.mult = 2.0
        self.consolidator_bars = 20
        #self.wvf = None

        # set leverage
        self.leverage = 3
        self.yoy_return = 0.6

        # Settings
        self.universe_settings.resolution = Resolution.HOUR
        self.set_start_date(2024,5,1)
        self.set_cash(100000)

        # Warming up
        resolution_in_time_span =  Extensions.to_time_span(self.universe_settings.resolution)
        warm_up_time_span = Time.multiply(resolution_in_time_span, self.consolidator_bars)
        self.set_warm_up(warm_up_time_span)

        # Universe Selection
        tickers = ['TQQQ']
        symbols = [Symbol.create(ticker, SecurityType.EQUITY, Market.USA) for ticker in tickers]
        self.set_universe_selection(ManualUniverseSelectionModel(symbols))

        # Alpha Model
        #self.set_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(hours = 1)))
        self.set_alpha(DualThrustAlphaModel(self.ph, self.range_period, self.universe_settings.resolution, self.consolidator_bars))
        

        ## Portfolio Construction
        self.set_portfolio_construction(NullPortfolioConstructionModel())

        ## Execution
        self.set_execution(ImmediateExecutionModel())

        ## Risk Management
        self.set_risk_management(NullRiskManagementModel())

    def on_order_event(self, order_event):
        if order_event.status == OrderStatus.FILLED:
            self.debug("Purchased Stock: {0}".format(order_event.symbol))