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

class my_strategy(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2025, 1, 1)
        self.set_cash(100000)

        # Use ManualUniverseSelectionModel method to select a static set of equities, similar to traditional add_security methods, for a fixed trading universe.
        tickers = ["SPY", "QQQ"]
        my_symbol_list = []
        my_symbol = {}
        for ticker in tickers:
            my_symbol[ticker] = Symbol.create(ticker, SecurityType.EQUITY, Market.USA)
            my_symbol_list.append(my_symbol[ticker])

        self.add_universe_selection(ManualUniverseSelectionModel(my_symbol_list))
        self.universe_settings.resolution = Resolution.DAILY

        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        self.set_execution(ImmediateExecutionModel())

        self.add_alpha(my_ibs_alpha(symbol=my_symbol['SPY'], ibs_len=3, resolution=Resolution.DAILY))

from AlgorithmImports import *

class my_ibs_alpha(AlphaModel):
    
    
    def __init__(self, symbol, ibs_len = 3, resolution = Resolution.Daily):
        ''' Initializes a new instance of the MacdAlphaModel class
        Args:
            ibs_len: The average ibs period</param>'''
        self.symbol = symbol
        self.ibs_len = ibs_len
        self.resolution = resolution
        self.insightCollection = InsightCollection()
        self.symbolData_dict = {}

        self.buyThreshold = 0.4
        self.sellThreshold = 0.6

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, ibs_len, resolutionString)


    def Update(self, algorithm, data):
        insights = []
        sd = SymbolData(algorithm, self.symbol, self.ibs_len, self.resolution)
        avg_ibs = sd.avg_ibs.current.value

        if avg_ibs < self.buyThreshold:
            direction = InsightDirection.Up
        if avg_ibs > self.sellThreshold:
            direction = InsightDirection.Flat

        # ignore signal for same direction as previous signal
        if direction == sd.PreviousDirection:
            return insights

        sd.PreviousDirection = direction

        if direction == InsightDirection.Flat:
            self.CancelInsights(algorithm, sd.Security)
            return insights

        insight = Insight.Price(sd.Security, 1, direction)
        insights.append(insight)
        self.insightCollection.Add(insight)

        return insights


    def CancelInsights(self, algorithm, symbol):
        if not self.insightCollection.ContainsKey(symbol):
            return
        insights = self.insightCollection[symbol]
        algorithm.Insights.Cancel(insights)
        self.insightCollection.Clear([ symbol ]);


class SymbolData:
    def __init__(self, algorithm, security, ibs_len, resolution):
        self.Security = security
        self.Consolidator = algorithm.ResolveConsolidator(self.Security, resolution)
        algorithm.subscription_manager.add_consolidator(self.Security, self.Consolidator)

        #self.ibs = InternalBarStrength(self.Security)
        self.my_ibs = algorithm.ibs(self.Security, resolution)
        self.avg_ibs = IndicatorExtensions.sma(self.my_ibs, ibs_len)

        algorithm.RegisterIndicator(self.Security, self.my_ibs, self.Consolidator)
        algorithm.WarmUpIndicator(self.Security, self.my_ibs, resolution)
        
        algorithm.RegisterIndicator(self.Security, self.avg_ibs, self.Consolidator)
        algorithm.WarmUpIndicator(self.Security, self.avg_ibs, resolution)

        self.PreviousDirection = None
        algorithm.plot("IBS", 'ibs', self.my_ibs.current.value)
        algorithm.plot("AverageIBS", 'ibs', self.avg_ibs.current.value)