from collections import deque
from collections.abc import Iterable
from datetime import datetime
from typing import Deque, Dict, List, Union

import numpy as np
import pandas as pd

from indicators import Drawdown
dev_mode = False
if dev_mode:
    from AlgorithmImports import *

class InAndOut:
    def __init__(self, algo:QCAlgorithm, resolution, symbols:List[str], period:int, iniWaitDays, minWaitDays, waitDaysConst, bull=True):
        self.Time = datetime.min
        
        self.period = period
        self.iniWaitDays = iniWaitDays
        self.minWaitDays = minWaitDays
        self.waitDaysConst = waitDaysConst
        self.bull = bull

        (self.market, self.silver, self.gold, self.utility, 
            self.industrial, self.safe, self.risk, self.debt_short, 
            self.debt_inflation, self.metal, self.input, self.cash) = symbols

        self.bull_signal_indices = [self.industrial, self.metal, self.input]

        self.history: dict[Symbol, Deque[float]] = {}
        
        history_df = algo.History(symbols, period, resolution)

        for symbol in symbols:
            if symbol in history_df and len(history_df[symbol]) > 0:
                self.history[symbol] = deque(history_df[symbol]['close'], maxlen=period)
            else:
                self.history[symbol] = deque(maxlen=period)

        self.wait_days = 0

    def Update(self, input:Slice):
        self.Time = input.Time

        for symbol, history in self.history.items():
            if input.Bars.ContainsKey(symbol):
                history.append(input[symbol].Close)
        
    def is_bullish(self):
        '''
        returns (true iff "bullish", wait_days)
        '''
        history_dict = {symbol: pd.Series(data) for symbol, data in self.history.items()}

        # (100 day returns / 11 day centered sma shifted by 60 days) - 1
        cust_returns_dict: dict[Union[Symbol, str], pd.Series] = {}
        for symbol in history_dict:
            if len(history_dict[symbol]) != self.period:
                return False, 0
        
            hist_series = history_dict[symbol]

            cust_returns_dict[symbol] = (hist_series / hist_series.rolling(11, center=True).mean().shift(60)).dropna() - 1
            history_dict[symbol] = history_dict[symbol][-len(cust_returns_dict[symbol]):]  # make all series the same length

        gold_min_silver = 'gold_min_silver'
        industrial_min_utility = 'industrial_min_utility'
        risk_min_safe = 'risk_min_safe'
        cash_inverse = 'cash_inverse'

        cust_returns_dict[gold_min_silver] = cust_returns_dict[self.gold] - cust_returns_dict[self.silver]
        cust_returns_dict[industrial_min_utility] = cust_returns_dict[self.industrial] - cust_returns_dict[self.utility]
        cust_returns_dict[risk_min_safe] = cust_returns_dict[self.risk] - cust_returns_dict[self.safe]
        cust_returns_dict[cash_inverse] = -1 * cust_returns_dict[self.cash]

        # values are true if last return is < 1 percentile
        is_extreme_returns_dict: dict[Union[Symbol, str], bool] = {}
        for symbol, returns in cust_returns_dict.items():
            is_extreme_returns_dict[symbol] = returns.iloc[-1] < np.percentile(returns, 1)

        inflation = 'inflation'
        history_dict[inflation] = cust_returns_dict[self.debt_short] - cust_returns_dict[self.debt_inflation]
        
        isabovemedian_dict = {
            symbol: (series.iloc[-1] > series.median()) for symbol, series in history_dict.items()
        }

        interest_expected = 'interest_expected'

        if is_extreme_returns_dict[self.debt_short] and isabovemedian_dict[self.metal] and isabovemedian_dict[self.input]:
            is_extreme_returns_dict[interest_expected] = False
        else:
            is_extreme_returns_dict[interest_expected] = is_extreme_returns_dict[self.debt_short]

        gold_min_silver_adj = 'gold_min_silver_adj'
    
        if is_extreme_returns_dict[gold_min_silver] and isabovemedian_dict[inflation]:
            is_extreme_returns_dict[gold_min_silver_adj] = False
        else:
            is_extreme_returns_dict[gold_min_silver_adj] = is_extreme_returns_dict[gold_min_silver]
        
        
        
        def wait_days_helper(symbol0, symbol1):
            series0 = cust_returns_dict[symbol0]
            series1 = cust_returns_dict[symbol1]
            if series0.iloc[-1] > 0 and series1.iloc[-1] < 0 and series1.iloc[-2] > 0:
                return self.iniWaitDays
            else:
                return 1

        self.wait_days = int(
            max(
                self.wait_days/2,
                self.iniWaitDays * max (
                    1, 
                    wait_days_helper(self.gold, self.silver), 
                    wait_days_helper(self.utility, self.industrial), 
                    wait_days_helper(self.safe, self.risk)
                )
            )
        )

        signals = self.bull_signal_indices + [gold_min_silver_adj, industrial_min_utility, risk_min_safe, cash_inverse]
        
        bullish = any([is_extreme_returns_dict[signal] for signal in signals])
        return bullish, min(self.minWaitDays, self.wait_days)

    def is_bearish(self):
        '''
        returns (true iff "bearish", wait_days)
        '''
        market_returns = pd.Series(self.history[self.market]).pct_change().dropna()
        volatililty = .6 * np.sqrt(252) * np.log1p(market_returns).std()
        returns_lookback = int(min(
            (1-volatililty)*self.waitDaysConst,
            self.period
        ))
        wait_days = int(volatililty * self.waitDaysConst)
        
        signals = [self.silver, self.gold, self.industrial, self.utility, self.metal, self.cash]

        returns = {}
        for signal in signals:
            data = self.history[signal]
            if len(data) < returns_lookback:
                return False, 0
            returns[signal] = pd.Series(data).pct_change(returns_lookback).iloc[-1]

        def compare(symbol0, symbol1):
            return returns[symbol0] < returns[symbol1]
        
        compares = compare(self.silver, self.gold) and compare(self.industrial, self.utility) and compare(self.metal, self.cash)
        return compares, wait_days            

    def minmax(self, n1, min_val, max_val):
        return max(min(n1, min_val), max_val)

    def get_signal(self):
        '''
        returns (true iff "bullish" and bull==True, wait_days) 
        returns (true iff "bearish" and bull==False, wait_days)
        '''
        if self.bull:
            return self.is_bullish()
        else:
            return self.is_bearish()

class ReturnsManager:
    def __init__(self, algo, period, resolution, max_drawdown, drawdown_lookback, max_alloc):
        '''
        Manages asset weighting. Get weights with the `GetWeights()` method.
        Also tells us if the `max_drawdown` has been reached with `IsSell()` method
        '''
        self.algo = algo
        self.period = period
        self.drawdown_lookback = drawdown_lookback
        self.resolution = resolution
    
        # one day returns
        self.daily_returns_dict: Dict[Symbol, RateOfChange] = {}
        # `period` day returns
        self.returns_dict: Dict[Symbol, RateOfChange] = {}
        self.dd = Drawdown(self.drawdown_lookback)

        self.weights_dict: Dict[Symbol, float] = {}

        self.max_drawdown = max_drawdown
        self.max_alloc = max_alloc
        
    def Update(self, input:Slice):
        portfolio_returns_cross_section = 0
        
        # update daily returns and period returns
        for symbol in self.returns_dict: 
            if input.Bars.ContainsKey(symbol):
                daily_returns = self.daily_returns_dict[symbol]
                daily_returns.Update(input.Time, input[symbol].Close)
                self.returns_dict[symbol].Update(input.Time, input[symbol].Close)
                
                if symbol in self.weights_dict and daily_returns.IsReady:
                    # weighted returns of one item in the portfolio
                    item_returns = self.weights_dict[symbol] * daily_returns.Current.Value
                    portfolio_returns_cross_section += item_returns

        if portfolio_returns_cross_section:
            self.dd.Update(portfolio_returns_cross_section)

    def UpdateWeights(self):
        # used to divide weights to make them sum to 1
        total_weight = sum(
            returns.Current.Value 
            for returns in self.returns_dict.values() 
            if returns.IsReady
        )

        if not total_weight:
            return {}

        self.weights_dict = {
            symbol: min(returns.Current.Value / total_weight, self.max_alloc)
                for symbol, returns in self.returns_dict.items()
        }
        return self.weights_dict

    def GetWeights(self):
        return self.weights_dict

    def IsSell(self):
        return self.dd.Value > self.max_drawdown

    @property
    def IsReady(self):
        return np.any([returns.IsReady for symbol, returns in self.returns_dict.items()]) and len(self.weights_dict) > 0

    def WarmUp(self, symbols:Union[Symbol, None]):
        '''
        warmups the symbol/symbols indicators
        '''
        if not isinstance(symbols, Iterable):
            symbols = [symbols] 

        hist = self.algo.History(symbols, self.period, self.resolution)
        for symbol in symbols:
            if symbol not in hist.index or not len(hist.loc[symbol]):
                continue
            closes = hist.loc[symbol]['close']
            for dt, close in closes.iteritems():
                self.daily_returns_dict[symbol].Update(dt, close)
                self.returns_dict[symbol].Update(dt, close)

    def AddSecurity(self, symbol:Symbol):
        '''
        registers the new symbol to the ReturnsManager
        '''
        if symbol not in self.returns_dict:
            self.daily_returns_dict[symbol] = RateOfChange(1)
            self.returns_dict[symbol] = RateOfChange(self.period)
            # self.WarmUp(symbol)

dev_mode = False # for Shile's use, keep False
if dev_mode:
    from AlgorithmImports import *
resolution = Resolution.Hour

# ---indicies---
market = 'SPY'
silver = 'SLV'
gold = 'GLD'
utility = 'XLU'
industrial = 'XLI'
safe = 'FXF'  # safe currency 
risk = 'FXA'  # risk currency
debt_short = 'SHY'
debt_inflation = 'TIP'
metal = 'DBB'
inp = 'IGE' # input
cash = 'UUP'

# ---equities
equities = ['SPY', 'QQQ', 'AAPL', 'AMD', 'DIA', 'AMZN']
#"ATVI", "ADBE", "AMD", "ALGN", "ALXN", "AMZN", "AMGN", "AAL", "ADI", "AAPL", "AMAT", "ASML", "ADSK", "ADP", "AVGO", "BIDU", "BIIB", "BMRN", "CDNS", "CERN", "CHKP", "CHTR", "TCOM", "CTAS", "CSCO", "CTXS", "CMCSA", "COST", "CSX", "CTSH", "DLTR", "EA", "EBAY", "EXC", "EXPE", "FAST", "FB", "FISV", "GILD", "GOOG", "GOOGL", "HAS", "HSIC", "ILMN", "INCY", "INTC", "INTU", "ISRG", "IDXX", "JBHT", "JD", "KLAC", "KHC", "LRCX", "LBTYA", "LBTYK", "LULU", "MELI", "MAR", "MCHP", "MDLZ", "MNST", "MSFT", "MU", "MXIM", "MYL", "NTAP", "NFLX", "NTES", "NVDA", "NXPI", "ORLY", "PAYX", "PCAR", "BKNG", "PYPL", "PEP", "QCOM", "REGN", "ROST", "SIRI", "SWKS", "SBUX", "NLOK", "SNPS", "TTWO", "TSLA", "TXN", "TMUS", "ULTA", "UAL", "VRSN", "VRSK", "VRTX", "WBA", "WDC", "WDAY", "WYNN", "XEL", "XLNX"]

# ---safeties
safeties = ['GLD', 'TLT']

# ---in and out parameters
# parameters found from file from you
bull = True # set False for bear
inOutLookbackBull = 252
inOutLookbackBear = 126

waitDaysConstant = 1 # WAITD_CONSTANT from your file
iniWaitDays = 1 # INI_WAIT_DAYS from your file
minWaitDays = 1 # 60 from the `min(60, self.WDadjvar)` from your file

# ---supertrend parameters
superTrendPeriod = 20
superTrendMultiple = 1
superTrendUseHA = True  # use Heiken-Ashii for superTrend

# ---squeeze parameters
squeezeTrendPeriod = 20
squeezeBBMultiple = 2 # BollingerBands
squeezeKeltMultiple = 3 # Kelter Channel (originally 1.5, increased to increase bullish trades)
# the lower the BBMultiple and the higher the KeltMultiple - the more likely squeeze will allow trades

# ---portfolio parameters
# for the returns based portfolio allocation
max_drawdown = .02 # max drawdown allowed before liquidation is signaled
drawdown_lookback = 1 # lookback for drawdown
max_alloc = .3 # max allocation to any given stock
returns_lookback = 1 # lookback for returns

# ---general parameters
rebalance = 1  # how often you want to update weights for rebalancing, refresh in&out signal

# disable select indicators
disableSupertrend = False
disableSqueeze = False
disableInAndOut = False  # setting this to True -> go long always except 10% drawdown

options_weight = 0 # what % of portfolio for options
equities_weight = 1 - options_weight

# specific options parameters
optionright = OptionRight.Call
price_strike_ratio = 0.9
'''
price_strike_ratio = (current equity price) / (strike price)
If call option and price_strike_ratio = .9
    that means we want calls whose underlying is 90% of the strike price
    which means it is .1 (10%) OTM
Similary, call option and price_strike_ratio = 1.1
    that means we want calls 10% ITM
Reverse the logic for puts
'''
expiry_liquidation = 10 # how many days before expiry to sell option
min_expiry = 20 # option will expire in at least this many days before buying
options_refresh_period = 20 # how many days

debug = False # show debug messages

dev_mode = False
if dev_mode:
    from AlgorithmImports import *
    
from collections import deque
from datetime import datetime 
from typing import Union
import pandas as pd

class MySuperTrend:
    def __init__(self, period, multiple, movingAverageType = MovingAverageType.Exponential):
        self.Name = "Custom Indicator"
        self.Time = datetime.min
        self.Value = 0
        
        self.multiplier = multiple
        self.atr = AverageTrueRange(period, movingAverageType)
        self.values = deque(maxlen=period)

        self.previousTrailingLowerBand = 0
        self.previousTrailingUpperBand = 0
        self.previousClose = 0
        self.previousTrend = 0

        
    def __repr__(self):
        return "{0} -> IsReady: {1}. Time: {2}. Value: {3}".format(self.Name, self.IsReady, self.Time, self.Value)

    def Update(self, input:TradeBar):
        self.Time = input.EndTime
        self.atr.Update(input)
        superTrend = 0

        currentClose = input.Close
        currentBasicLowerBand = (input.Low + input.High) / 2 - self.multiplier * self.atr.Current.Value
        currentBasicUpperBand = (input.Low + input.High) / 2 + self.multiplier * self.atr.Current.Value

        if self.previousClose > self.previousTrailingLowerBand:
            currentTrailingLowerBand = max(currentBasicLowerBand, self.previousTrailingLowerBand)
        else:
            currentTrailingLowerBand = currentBasicLowerBand

        if self.previousClose < self.previousTrailingUpperBand:
            currentTrailingUpperBand = min(currentBasicUpperBand, self.previousTrailingUpperBand)
        else:
            currentTrailingUpperBand = currentBasicUpperBand

        if currentClose > currentTrailingUpperBand:
            currentTrend = 1
        elif currentClose < currentTrailingLowerBand:
            currentTrend = -1
        else:
            currentTrend = self.previousTrend

        if currentTrend == 1:
            superTrend = currentTrailingLowerBand
        elif currentTrend == -1:
            superTrend = currentTrailingUpperBand
        
        self.previousTrailingLowerBand = currentTrailingLowerBand
        self.previousTrailingUpperBand = currentTrailingUpperBand
        self.previousClose = currentClose
        self.previousTrend = currentTrend

        if not self.atr.IsReady:
            return 0
        
        self.Value = superTrend
        return self.IsReady
    
    @property
    def IsReady(self):
        return self.atr.IsReady and self.Value != 0

class Squeeze:
    '''
    .Value = 1 iff "squeezed" else .Value = 0
    
    Tells us if we are in or out of squeeze

    Is Squeeze: lower BB > lower Keltner and upper BB < upper Keltner
    '''

    def __init__(self, period, bollinger_multiple, kelt_multiple, movingAverageType = MovingAverageType.Exponential):
        '''
        .Value = 1 iff "squeezed" else .Value = 0
        '''
        self.Name = "SuperTrend"
        self.Time = datetime.min
        self.Value = 0
        
        self.bb = BollingerBands(period, bollinger_multiple, movingAverageType)
        self.kelt = KeltnerChannels(period, kelt_multiple, movingAverageType)

    def __repr__(self):
        return "{0} -> IsReady: {1}. Time: {2}. Value: {3}".format(self.Name, self.IsReady, self.Time, self.Value)

    def Update(self, input:TradeBar):
        self.Time = input.EndTime
        self.kelt.Update(input)
        self.bb.Update(input.EndTime, input.Close)

        isSqueeze = self.bb.LowerBand.Current.Value > self.kelt.LowerBand.Current.Value and self.bb.UpperBand.Current.Value <  self.kelt.UpperBand.Current.Value
        self.Value = int(isSqueeze)
        return self.IsReady 

    @property
    def IsReady(self):
        return self.kelt.IsReady and self.bb.IsReady

class Drawdown:
    def __init__(self, period:int):
        '''
        drawdown indicator for past `period` values
        Call Update with floats that represent returns
        '''
        self.values = deque(maxlen=period)
        self.Value = 0

    def Update(self, input:Union[TradeBar, float]):
        if isinstance(input, float):
            self.values.append(input)
        else:
            self.values.append(input.Close)
        
        if self.IsReady:
            x=5
        
        # https://stackoverflow.com/questions/36750571/calculate-max-draw-down-with-a-vectorized-solution-in-python
        cum_returns = (1 + pd.Series(self.values)).cumprod()
        self.Value = 1 - cum_returns.div(cum_returns.cummax()).iloc[-1]  # drawdown
        return self.IsReady
        
    @property
    def IsReady(self):
        return len(self.values) == self.values.maxlen

import configs as cfg
dev_mode = False
if dev_mode:
    from AlgorithmImports import *
    from datetime import timedelta

from typing import Dict, List
from indicators import MySuperTrend, Squeeze
from aggregate_indicators import InAndOut, ReturnsManager
import operator

class AdaptableRedSnake(QCAlgorithm):
    def load_configs_and_indicators(self):
        # OVERRIDE CONFIGS HERE
        # EXAMPLE: cfg.returns_lookback = self.GetParameter('returns_lookback')
        
        index_tickers = [cfg.market, cfg.silver, cfg.gold, cfg.utility, cfg.industrial, cfg.safe, cfg.risk, cfg.debt_short, cfg.debt_inflation, cfg.metal, cfg.inp, cfg.cash]
        self.indices = [
            self.AddEquity(ticker, cfg.resolution).Symbol for ticker in index_tickers
        ]

        inOutLookback = cfg.inOutLookbackBull if cfg.bull else cfg.inOutLookbackBear
        self.inandout = InAndOut(self, cfg.resolution, self.indices, inOutLookback, cfg.iniWaitDays, cfg.minWaitDays, cfg.waitDaysConstant, cfg.bull)
        self.returnsmanager = ReturnsManager(self, cfg.returns_lookback, cfg.resolution, cfg.max_drawdown, cfg.drawdown_lookback, cfg.max_alloc)

        self.equities = self.safeties = [
            self.AddEquity(ticker, cfg.resolution).Symbol for ticker in cfg.equities
        ]
        
        self.symbolData = {
            symbol: SymbolData(self, symbol, cfg.resolution,
                cfg.superTrendPeriod, cfg.superTrendMultiple, 
                cfg.squeezeTrendPeriod, cfg.squeezeBBMultiple, 
                cfg.squeezeKeltMultiple, cfg.superTrendUseHA)
                for symbol in self.equities
        }

        for symbol in self.equities:
            self.returnsmanager.AddSecurity(symbol)

        self.safeties = [
            self.AddEquity(ticker, cfg.resolution).Symbol for ticker in cfg.safeties
        ]
            
        self.next_reentry = self.Time
        self.was_bull = False
        self.inout_signal = False
        self.inout_waitdays = 0

        self.SetWarmUp(max(
            cfg.inOutLookbackBear, cfg.inOutLookbackBull, cfg.superTrendPeriod, cfg.squeezeTrendPeriod, cfg.returns_lookback
        ), Resolution.Hour)

    def Initialize(self):
        self.SetSecurityInitializer(
            lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw) 
                if x.Type == SecurityType.Equity else None
        )
        self.load_configs_and_indicators()

        self.SetStartDate(2020, 1, 1)
        self.SetCash(100000) 

        # self.Schedule.On(self.DateRules.WeekStart(), self.TimeRules.Midnight, self.WeekleyFn)  
        self.days_count = 0
        self.curr_day = -1

        # equity Symbol: option Symbol
        self.options: Dict[Symbol, Symbol] = {}

    def Print(self, msg):
        if cfg.debug:
            self.Debug(msg)

    def OnData(self, data:Slice):
        if self.curr_day == self.Time.day:
            return 
        self.curr_day = self.Time.day

        self.CheckOptions()

        if self.days_count % cfg.options_refresh_period:
            self.SetOptions()
        else:
        # OPTIONS LOGIC
            options: List[Symbol] = []
            for symbol in self.options.values():
                if symbol is None or not data.ContainsKey(symbol):
                    continue
                options.append(symbol)

            invested_options = {
                kvp.Key.Underlying: kvp.Key for kvp in self.Portfolio 
                    if kvp.Value.Invested and kvp.Key.SecurityType == SecurityType.Option 
            }

            for option in options:
                if not option.Underlying in invested_options.keys() and not self.Securities[option].Invested:
                    self.SetHoldings(option, cfg.options_weight / len(options))

        # EQUITIES LOGIC
        if not any([data.Bars.ContainsKey(symbol) for symbol in self.equities]):
            return

        self.days_count += 1
        self.inandout.Update(data)
        self.returnsmanager.Update(data)
        if self.days_count % cfg.rebalance == 0:
            self.returnsmanager.UpdateWeights()
            self.inout_signal, self.inout_waitdays = self.inandout.get_signal()
        
        if self.IsWarmingUp:
            return

        if not self.returnsmanager.IsSell() and (self.inout_signal or cfg.disableInAndOut):
            cfg.debug and self.Print('Bull Condition Reached')
            if not self.was_bull and self.Time >= self.next_reentry:
                self.Print('Going Bull')
                self.go_bull()
                self.was_bull = True
        elif self.was_bull:
            self.Print(f'Going Bear:')
            self.go_bear()
            self.was_bull = False
            self.next_reentry = self.Time + timedelta(days=self.inout_waitdays)
        
        self.PlotSeries()
    
    def CheckOptions(self):
        '''
        Sell options near expiry, set new options for liquidated ones
        '''
        for equity, option in self.options.items():
            if option.ID.Date - self.Time < timedelta(days=cfg.expiry_liquidation):
                self.Liquidate(option)
                self.SetOptions(equity)


    def SetOptions(self, equity=None):
        if self.IsWarmingUp:
            return
        self.Print('Setting Options')
        if equity is None:
            equities = self.equities
        else:
            equities = [equity]
        
        for equity in equities:
            contracts = self.OptionChainProvider.GetOptionContractList(equity, self.Time)
            equity_price = self.Securities[equity].Price

            if cfg.optionright == OptionRight.Call:
                compare = operator.lt if cfg.price_strike_ratio < 1 else operator.gt
            else:
                compare = operator.lt if cfg.price_strike_ratio >= 1 else operator.gt
            
            def filter(id: SecurityIdentifier):
                return (
                    id.OptionRight == cfg.optionright 
                    and compare(equity_price / id.StrikePrice, cfg.price_strike_ratio)
                    and id.Date - self.Time > timedelta(days=cfg.min_expiry)
                )

            # filter
            contracts = [contract for contract in contracts if filter(contract.ID)]
            
            # sort by date
            contracts = sorted(contracts, key=lambda x: x.ID.Date)
            
            # sort by closest strike.
            contracts = sorted(contracts, key=lambda x: abs(equity_price - x.ID.StrikePrice))

            if contracts:
                contract = contracts[0]
                
                self.AddOptionContract(contract, Resolution.Minute)
                self.options[equity] = contract
                
        return self.options

    def PlotSeries(self): 
        self.Plot('BullBear', 'bull=1,bear=0', int(self.was_bull))
 
    def go_bear(self):
        if self.IsWarmingUp:
            return
        for equity in self.equities:
            self.Liquidate(equity)
        for symbol in self.safeties:
            self.SetHoldings(symbol, cfg.equities_weight/len(self.safeties))
        self.was_bull = False

    def go_bull(self):
        if self.IsWarmingUp:
            return
        for equity in self.safeties:
            self.Liquidate(equity)
        for symbol, symbolData in self.symbolData.items():
            if symbolData.IsBuy(self.Securities[symbol].Price):
                weight = self.returnsmanager.GetWeights().get(symbol, 0)
                if weight:
                    self.SetHoldings(symbol, weight * cfg.equities_weight)

class SymbolData:
    def __init__(self, algo:QCAlgorithm, symbol, resolution, periodST, multipleST, periodSQ, BBmultipleSQ, KELTmultipleSQ, useHA):
        self.symbol = symbol
        self.supertrend = MySuperTrend(periodST, multipleST)
        self.squeeze = Squeeze(periodSQ, BBmultipleSQ, KELTmultipleSQ)
        
        self.algo = algo
        algo.RegisterIndicator(symbol, self.supertrend, Resolution.Hour)
        algo.RegisterIndicator(symbol, self.squeeze, Resolution.Hour)

        self.useHA = useHA
        if useHA:
            self.ha = HeikinAshi(symbol, Resolution.Minute)
            algo.RegisterIndicator(symbol, self.ha, Resolution.Hour)
               
    def Update(self, input:TradeBar):
        if self.useHA and self.ha.IsReady:
            haBar = TradeBar(self.algo.Time, self.symbol, 
                self.ha.Open.Current.Value, self.ha.High.Current.Value, 
                self.ha.Low.Current.Value, self.ha.Close.Current.Value, self.ha.Volume.Current.Value)
            self.supertrend.Update(haBar)
        else:
            self.supertrend.Update(input)
        self.squeeze.Update(input)

    def IsBuy(self, price):
        return (self.squeeze.Value or cfg.disableSqueeze) or (price > self.supertrend.Value or cfg.disableSupertrend)