#region imports
from AlgorithmImports import *
#endregion

## PORTFOLIO MANAGMENT
IVT_PERCENTAGE = 0.2

# The following is a reference to the entry threshold to 
# limit the price were th order should take place
# Check StopLimitOrders for references:
# https://www.quantconnect.com/docs/v2/writing-algorithms/trading-and-orders/order-types/stop-limit-orders
LIMIT_TH = 0.2

## STD Comparisions
# If true the short positions can be opened when detected
ACTIVATE_SHORT = True
# STD is one standard deviation
# Each of the following works as (STD * SHORT_...)
SHORT_TH = 2 # Threshold to open Short Position
SHORT_SL = SHORT_TH + 0.1 # STD to Stop Loss
SHORT_TP = SHORT_TH - 0.1 # STD to take profit


# If true the long positions can be opened when detected
ACTIVATE_LONG = True
# Each of the following works as (STD * LONG_...)
LONG_TH = -2 # Threshold to open Long Position
LONG_SL = LONG_TH + 0.1 # STD to Stop Loss
LONG_TP = LONG_TH - 0.1  # STD to take profit

## UNIVERSE SETTINGS
MY_FUTURES = [Futures.Indices.SP500EMini]#Futures.Indices.NASDAQ100EMini]

# INDICATORS
BB_PERIOD = 30
INDICATORS_RESOLUTION = Resolution.Minute

#region imports
from AlgorithmImports import *

# Custom
from control_parameters import *
## endregion


class RiskWithBollingerBands(RiskManagementModel):
    ''''''
    def __init__(self,bb_period:int, bb_resolution:Resolution, bb_args:dict={},
                    LongProfitPercent:float = 2.5, ShortProfitPercent:float=0.9,
                    LongDrawdownPercent:float = 0.9, ShortDrawdownPercent:float=2.5,
                    ):
        '''
        Inputs:
            - bb_period [int]: Period to apply on BB
            - bb_args [QC BollingerBands parameters]: Should contain extra arguments for BB

            Standar Deviations triggers to use on the bollinger bands:
             - LongProfitPercent: When hits and Long take profit
             - LongDrawdownPercent: When hits and Long Stop Loss
             - ShortProfitPercent: When hits and Short take profit
             - ShortDrawdownPercent: When hits and Short Stop Loss
        '''
        self.bb_period = bb_period 
        self.bb_resolution = bb_resolution
        self.bb_args = bb_args
        self.LongProfitPercent = LongProfitPercent
        self.ShortProfitPercent = ShortProfitPercent
        self.LongDrawdownPercent = LongDrawdownPercent
        self.ShortDrawdownPercent = ShortDrawdownPercent
        self.trailingBBs = dict()

    def ManageRisk(self, algorithm, targets):
        '''
        Manages the algorithm's risk at each time step
        Inputs:
            - algorithm: The algorithm instance of QC.
            - targets: The current portfolio targets are to be assessed for risk
        '''
        riskAdjustedTargets = list()

        for kvp in algorithm.Securities:
            symbol = kvp.Key
            security = kvp.Value

            if security.Type == SecurityType.Future:
                # Get Canonical Object
                trailingBBState = self.trailingBBs.get(symbol.Canonical)
            # Next if not Future
            if not trailingBBState:
                continue
            # Remove if not invested
            if not security.Invested: # For positions closed outside the risk management model
                trailingBBState.position.pop(symbol, None) # remove from dictionary
                continue
            

            # Get position side
            position = PositionSide.Long if security.Holdings.IsLong else PositionSide.Short
            
            # Recorded Holdings Value
            
            stored_position = trailingBBState.position.get(symbol)
            # Add newly invested contract (if doesn't exist) or reset holdings state (if position changed)
            if stored_position is None or position != stored_position:
                # Create a BB State object if not existing or reset it if the position direction changed
                trailingBBState.position[symbol] = position

            # Check take profit trigger
            if ((position == PositionSide.Long and trailingBBState.Price > (trailingBBState.MiddleBand + trailingBBState.GetStandardDeviation(self.LongProfitPercent))) or
                (position == PositionSide.Short and trailingBBState.Price < (trailingBBState.MiddleBand - trailingBBState.GetStandardDeviation(self.ShortProfitPercent)))):
                # Update position
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))
                # Pop the symbol from the dictionary since the holdings state of the security has been changed
                trailingBBState.position.pop(symbol, None) # remove from dictionary
                continue

            elif ((position == PositionSide.Long and trailingBBState.Price < (trailingBBState.MiddleBand - trailingBBState.GetStandardDeviation(self.LongDrawdownPercent))) or
                (position == PositionSide.Short and trailingBBState.Price > (trailingBBState.MiddleBand + trailingBBState.GetStandardDeviation(self.ShortDrawdownPercent)))):
                # liquidate
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))
                # Pop the symbol from the dictionary since the holdings state of the security has been changed
                trailingBBState.position.pop(symbol, None) # remove from dictionary

        return riskAdjustedTargets

## SECURITIES LOGIC: CREATION, INDICATORS, UPDATE, TACKING
    def OnSecuritiesChanged(self, algorithm, changes: SecurityChanges) -> None:
        # Gets an object with the changes in the universe
        # For the added securities we create a SymbolData object that allows 
        # us to track the orders associated and the indicators created for it.
        for security in changes.AddedSecurities:
            if self.trailingBBs.get(security.Symbol) is None: # Create SymbolData object
                self.trailingBBs[security.Symbol] = BBState(algorithm, security.Symbol, 
                                                            self.bb_period, self.bb_resolution, self.bb_args)
        
        # The removed securities are liquidated and removed from the security tracker.
        for security in changes.RemovedSecurities: # Don't track anymore
            self.SubscriptionManager.RemoveConsolidator(security.Symbol, self.trailingBBs[security.Symbol].consolidator)
            self.SecuritiesTracker.pop(security.Symbol, None)

class BBState:
    def __init__(self,algorithm, symbol:Symbol, 
                    period:int,bb_resolution: Resolution,bb_args,
                    position:PositionSide=None):
        if not(position):
            self.position = {}
        self.Symbol = symbol

        # Create
        self.bb_resolution = bb_resolution
        self.bb = BollingerBands('BB '+symbol.Value,period, 1, **bb_args)

        # Create consolidator for Symbol
        self.consolidator = QuoteBarConsolidator(bb_resolution)
        self.consolidator.DataConsolidated += self.UpdateBB
        algorithm.SubscriptionManager.AddConsolidator(self.Symbol, self.consolidator)
    
    def UpdateBB(self, sender: object, consolidated_bar: TradeBar) -> None:
        self.bb.Update(consolidated_bar.EndTime, consolidated_bar.Close)

    @property
    def Price(self):
        return self.bb.Price.Current.Value

    @property
    def StandardDeviation(self):
        return self.bb.StandardDeviation.Current.Value

    @property
    def MiddleBand(self):
        return self.bb.MiddleBand.Current.Value
    
    def GetStandardDeviation(self, decimal:float) -> float:
        return self.bb.StandardDeviation.Current.Value * decimal

# region imports
from AlgorithmImports import *

# Custom
from control_parameters import *
import symbol_data
from custom_risk_management import RiskWithBollingerBands
# endregion

class CalmLightBrownBadger(QCAlgorithm):

## INITIALIZATION
    def Initialize(self):
        self.SetStartDate(2021, 9, 17)  # Set Start Date
        self.SetEndDate(2021, 12, 17)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        # Broker and Account type: Margin allow the use of leverage and shorts
        self.SetBrokerageModel(BrokerageName.QuantConnectBrokerage, AccountType.Margin)
        
        # Universe Settings
        self.InitializeUniverse()
        
        # Add custom modification every time a security is initialized
        # feed data when adding contracts in the middle of OnData method
        self.SetSecurityInitializer(self.CustomSecurityInitializer)

        # Define risk management model: Usin default values
        self.AddRiskManagement(RiskWithBollingerBands(bb_period=BB_PERIOD, bb_resolution=INDICATORS_RESOLUTION,
                                                    ShortProfitPercent=SHORT_TP, ShortDrawdownPercent=SHORT_SL,
                                                    LongProfitPercent=LONG_TP, LongDrawdownPercent=LONG_SL,
                                                    ))
        self.AddRiskManagement(TrailingStopRiskManagementModel())

    def InitializeUniverse(self):
        # Initilize tracker parameters
        self.SecuritiesTracker = {}
        
        # Store desired futures 
        self.__Futures = set()

        # Universe selector for futures
        for f in MY_FUTURES:
            # Minute allows a good enough flow of data for risk management.
            self.__Futures.add(self.AddFuture(f,Resolution.Minute,
                                            dataNormalizationMode = DataNormalizationMode.BackwardsRatio,
                                            dataMappingMode = DataMappingMode.OpenInterest,
                                            contractDepthOffset= 0))

    def CustomSecurityInitializer(self, security):
        # When adding a future we could get 0 as Ask/Bid
        # prices since the data has not been feed. 
        # This solves that issue
        bar = self.GetLastKnownPrice(security)
        security.SetMarketPrice(bar)

    def InitCharts(self):
        
        chart = Chart("BollingerBands")
        self.AddChart(chart)

        series = {}
        Series("<seriesName>")
        chart.AddSeries(series)

## SECURITIES LOGIC: CREATION, INDICATORS, UPDATE, TACKING
    def InitIndicators(self, symbol):
        # Create the indicators related to the input symbol
        # The default MA is Simple, if changed to Exponential in the future
        # remember to change the 
        bb = BollingerBands('BB '+ symbol.Value,BB_PERIOD, 1)
        self.RegisterIndicator(symbol, bb, INDICATORS_RESOLUTION)
        return bb

    def ManageAdded(self, added):
        '''
        Logic for securities added. Create the SymbolData objects that track the added securities.
        Inputs:
            added [list]: List of added securities
        '''
        for security in added:
            if self.SecuritiesTracker.get(security.Symbol) is None: # Create SymbolData object
                self.SecuritiesTracker[security.Symbol] = symbol_data.SymbolData(security, self.get_Time,
                                                                                self.InitIndicators(security.Symbol), short_th=SHORT_TH,long_th=LONG_TH)
    
    def ManageRemoved(self, removed):
        '''
        Logic for securities removed. Remove the SymbolData objects.
        Inputs:
            removed [list]: List of removed securities
        '''
        for security in removed: # Don't track anymore
            if self.Portfolio[security.Symbol].Invested:
                self.Liquidate(security.Symbol) 
            self.SecuritiesTracker.pop(security.Symbol, None)

    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        # Gets an object with the changes in the universe
        # For the added securities we create a SymbolData object that allows 
        # us to track the orders associated and the indicators created for it.
        self.ManageAdded(changes.AddedSecurities)
        
        # The removed securities are liquidated and removed from the security tracker.
        self.ManageRemoved(changes.RemovedSecurities)

## CHECK FOR BUYING POWER

    def CheckBuyingPower(self, symbol, quantity):
        '''
        Check for enough buying power.
        If the buying power for the target quantity is not enough,
        It will return the quantity for which the buying power is enough.
        '''
        if quantity < 0:
            order_direction = OrderDirection.Sell
        else:
            order_direction = OrderDirection.Buy
        # Get the buying power depending of the order direction and symbol
        buy_power = self.Portfolio.GetBuyingPower(symbol, order_direction)
        # Compute possible quantity
        q_t = abs(buy_power) / self.Securities[symbol].Price
        # Select minimum quantity
        return round(min(abs(quantity),q_t),8)*np.sign(quantity)

    def CheckOrdeQuatity(self, symbol, quantity):
        '''Check that the quantity of shares computed meets the minimum requirments'''
        q = abs(quantity)
        # There are requirements for the minimum or maximum that can be purchased per security.
        if q > self.Settings.MinAbsolutePortfolioTargetPercentage and q < self.Settings.MaxAbsolutePortfolioTargetPercentage:
            symbol_properties = self.Securities[symbol].SymbolProperties
            if symbol_properties.MinimumOrderSize is None or q > symbol_properties.MinimumOrderSize:
                return True
        return False

## POSITION MANAGEMENT
    def my_round(self, x, prec=2, base=0.25): 
        return (base * (np. array(x) / base).round()).round(prec)
    
    def OnData(self, data):
        for cont_symbol, tracked in self.SecuritiesTracker.items():
            self.GetView(tracked.bb, cont_symbol.Value)
            if self.Portfolio[cont_symbol].Invested:
                continue
            elif ACTIVATE_SHORT and self.Securities[cont_symbol].Exchange.Hours.IsOpen(self.Time, False) and tracked.IsReady and tracked.OverShort:
                symbol = self.Securities[cont_symbol].Mapped
                # Current asset quantity on portfolio
                hold_value = self.Portfolio[symbol].AbsoluteHoldingsValue
                current_wallet = self.Portfolio.MarginRemaining + hold_value
                # Current quantity
                quantity = self.Portfolio[symbol].Quantity
                target = -(current_wallet * IVT_PERCENTAGE)/self.Securities[symbol].Price
                target -= quantity
                quantity = self.CheckBuyingPower(symbol, int(target)) # Check for buying power
                if self.CheckOrdeQuatity(symbol, quantity): # Check for the minimum quantity
                    stopPrice = self.my_round(tracked.StandardDeviationReferencePrice(tracked.ShortThreshold + LIMIT_TH))
                    limitPrice = self.my_round(tracked.StandardDeviationReferencePrice(tracked.ShortThreshold - LIMIT_TH))
                    buy = self.StopLimitOrder(symbol, quantity, stopPrice, limitPrice)
                    tracked.Order = buy # Save in the security tracker the contract and order
            elif ACTIVATE_LONG and self.Securities[cont_symbol].Exchange.Hours.IsOpen(self.Time, False) and tracked.IsReady and tracked.OverLong:
                symbol = self.Securities[cont_symbol].Mapped
                # Current asset quantity on portfolio
                hold_value = self.Portfolio[symbol].AbsoluteHoldingsValue
                current_wallet = self.Portfolio.MarginRemaining + hold_value
                # Current quantity
                quantity = self.Portfolio[symbol].Quantity
                target = (current_wallet * IVT_PERCENTAGE)/self.Securities[symbol].Price
                target -= quantity
                quantity = self.CheckBuyingPower(symbol, int(target)) # Check for buying power
                if self.CheckOrdeQuatity(symbol, quantity): # Check for the minimum quantity
                    stopPrice = self.my_round(tracked.StandardDeviationReferencePrice(tracked.LongThreshold + LIMIT_TH))
                    limitPrice = self.my_round(tracked.StandardDeviationReferencePrice(tracked.LongThreshold - LIMIT_TH))
                    buy = self.StopLimitOrder(symbol, quantity, stopPrice, limitPrice)
                    tracked.Order = buy # Save in the security tracker the contract and order

    def GetView(self,bb, symbol_value):
        if self.Time.minute % 15 == 0:
            self.Plot("BollingerBands", symbol_value + " middleband", bb.MiddleBand.Current.Value)
            self.Plot("BollingerBands", symbol_value + " upperband", bb.UpperBand.Current.Value)
            self.Plot("BollingerBands", symbol_value + " lowerband", bb.LowerBand.Current.Value)
            self.Plot("BollingerBands", symbol_value + " price", bb.Price.Current.Value)

#region imports
from AlgorithmImports import *
#endregion

class SymbolData:
    # Object to Keep track of the securities

## INITIALIZATION
    def __init__(self, security, time, 
                bb,
                short_th:float=None, long_th:float=None):
        '''
        Inputs:
            - Symbol [QC Symbol]: Reference to the security.
            - time [Main Algo function]: This function returns the time of the main algorithm.
        '''
        self.Security = security
        self.__Symbol = security.Symbol
        self.get_Time = time
        self.__Order = None 
        self.bb = bb

        assert short_th or long_th, 'Both short_th and long_th cannot be null. One or both thresholds shoud be specified.'
        
        if short_th:
            self.ShortThreshold = short_th
        else:
            self.ShortThreshold = long_th
        if long_th:
            self.LongThreshold = long_th
        else:
            self.LongThreshold = short_th

        self.IsOverShort = False
        self.IsOverLong = False
    
    @property
    def Time(self):
        # Allow access to the Time object directly
        return self.get_Time()
    @property
    def IsReady(self):
        return self.bb.IsReady

## MANAGEMENT
    @property
    def Symbol(self):
        return self.__Symbol

    @Symbol.setter
    def Symbol(self,NewSymbol):
        self.__Order = None
        self.__Symbol = NewSymbol
    
    @property
    def Order(self):
        return self.__Order

    @Order.setter
    def Order(self, order):
        self.__Order = order
    
## MANAGEMENT LOGIC
    @property
    def OverShort(self):
        if self.IsReady:
            # We want a short, the threshold is positive so we check that the price is over the middle band + (STD*ShortThreshold)
            # We want a short, the threshold is negative so we check that the price is over the middle band - (STD*ShortThreshold)
            # In both cases we should check that the price is higher because with short we expect the price to raise
            current = self.Security.Price > (self.bb.MiddleBand.Current.Value + (self.bb.StandardDeviation.Current.Value * self.ShortThreshold))
            state = current and (current != self.IsOverShort)
            self.IsOverShort = current
            return state
        return False

    @property
    def OverLong(self):
        if self.IsReady:
            # We want a short, the threshold is positive so we check that the price is lower the middle band + (STD*ShortThreshold)
            # We want a short, the threshold is negative so we check that the price is lower the middle band - (STD*ShortThreshold)
            # In both cases we should check that the price is higher because with long we expect the price to drop.
            current = self.Security.Price < (self.bb.MiddleBand.Current.Value - (self.bb.StandardDeviation.Current.Value * self.LongThreshold))
            state = current and (current != self.IsOverLong)
            self.IsOverLong = current
            return state
        return False

    def StandardDeviationReferencePrice(self, decimal:float=1) -> float:
        return self.bb.MiddleBand.Current.Value + (decimal* self.bb.StandardDeviation.Current.Value)