# Imports
from datetime import timedelta
import pickle
import traceback

# QuantConnect specific imports
from AlgorithmImports import *
import QuantConnect as qc

# Import from files
# n/a 

"""
AE Revision Notes
(09/08/2023) - Updated to read the target portfolio allocation from the Object Store.
(09/27/2023) - Allowing orders on the 1minute mark, since there is a 1min delay
                from reading update from the signal algo.
(09/28/2023) - Skipping CheckTargetAllocation() if not during live trading hours.
(12/06/2023) - 
"""

# Global variables
OS_KEY = "LiveAlgo"
DEBUG = False

###############################################################################
class CustomAlgorithm(QCAlgorithm):
    def Initialize(self):
        # Set starting date, cash and ending date of the backtest
        
        # self.SetEndDate(2017, 3, 31)
        self.SetCash(100000)
        self.SetTimeZone('US/Eastern')
# AE 12/6/23 update
        if not self.LiveMode:
            self.SetSecurityInitializer(
                CustomSecurityInitializer(
                    self.BrokerageModel, 
                    FuncSecuritySeeder(self.GetLastKnownPrices)
                )
            )
        # Add market symbol and create market condition instance
        self.bm = self.AddEquity("SPY", Resolution.Hour).Symbol
        self.bm_hours = self.Securities[self.bm].Exchange.Hours

        # Universe selection
        self.AddUniverse(self.CoarseSelectionFunction)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.ExtendedMarketHours = False # not using this for client accounts
        self.UniverseSettings.MinimumTimeInUniverse = 63
        self.symbol_objects = []

        # Schedule function to check the target allocation
        if self.LiveMode:
            self.Schedule.On(
                self.DateRules.EveryDay(self.bm),
                self.TimeRules.Every(timedelta(minutes=1)),
                self.CheckTargetAllocation
            )
            
        # Other settings
        self.Settings.FreePortfolioValuePercentage = 0.05
        # Required variables
        self.target_allocation = {}

#-------------------------------------------------------------------------------
    def CoarseSelectionFunction(self, coarse):
        stocks = list(filter(lambda x: x.HasFundamentalData, coarse))
        sortedByDollarVolume = sorted(
            stocks, key=lambda x: x.DollarVolume, reverse=True
        )
        symbols = [x.Symbol for x in sortedByDollarVolume[:50]]
        # Print universe details when live mode
        if self.LiveMode:
            self.MyLog(f"Coarse filter returned {len(symbols)} stocks.")
        return symbols

#-------------------------------------------------------------------------------
    def CheckTargetAllocation(self):
        """Check Target Allocation."""
# AE added 9/28
        # Skip until we are in live trading hours
        if not self.bm_hours.IsOpen(self.Time, extendedMarketHours=False):
            return
        # Get the desired target allocation from the Object Store
        if DEBUG:
            self.MyLog('CheckTargetAllocation')
        # Always clear the cache before getting the target allocation 
        self.ObjectStore.Clear()
        if self.ObjectStore.ContainsKey(OS_KEY):
            deserialized = self.ObjectStore.ReadBytes(OS_KEY)
            target_allocation = pickle.loads(deserialized)
            if DEBUG:
                self.MyLog(f"Target allocation: {target_allocation}")
            # Log change in target allocation
            if target_allocation != self.target_allocation:
                self.MyLog(f"Change in target allocation: {target_allocation}")
                self.target_allocation = target_allocation
            # Update the portfolio based on the target allocation
            self.UpdatePortfolio(target_allocation)

#-------------------------------------------------------------------------------
    def UpdatePortfolio(self, target_allocation):
        """Update the current portfolio based on the target allocation."""
        # Get the current portfolio value
        nav = self.Portfolio.TotalPortfolioValue
        if DEBUG:
            self.MyLog(f'UpdatePortfolio(), current nav=${nav:.2f}')

        # Get list of tuples (symbol_objects, order shares) based type of order
        new_positions = []
        decrease_size = []
        increase_size = []
        # Loop through the target allocation
        target_portfolio = {}
        for symbol_object_str, tup in target_allocation.items():
            strategy = tup[0]
            target_weight = tup[1]
            dt = tup[2]
            if DEBUG:
                self.MyLog(
                    f'UpdatePortfolio(), symbol_object_str={symbol_object_str}, ' 
                    f'strategy={strategy}, target weight={target_weight}, '
                    f'dt={dt}'
                )
            # Get the actual QC symbol from the symbol string
            symbol_object = None
            for symbol in self.symbol_objects:
# AE update 12/6/23
                if (str(symbol) == symbol_object_str) or ((str(symbol.ID) == symbol_object_str)):
                    symbol_object = symbol
                    if DEBUG:
                        self.MyLog(
                            f'UpdatePortfolio(), {symbol_object_str} symbol '
                            'object found'
                        )
                    break
            # Check if the symbol object is not found
            if symbol_object is None:
                try:
                    # Create the desired symbol object
                    ticker = symbol_object_str.split(" ")[0]
                    symbol = self.AddEquity(ticker)
                    symbol_object = self.AddEquity(
                        ticker, Resolution.Minute
                    ).Symbol
                    if DEBUG:
                        self.MyLog(
                            f'UpdatePortfolio(), {symbol_object_str} symbol '
                            f'object: {symbol_object}')
                except:
                    self.MyLog(
                        f"UpdatePortfolio(), error trying to get symbol "
                        f"object for {symbol_object_str}"
                    )
                    continue

            # Get the target number of shares
            price = self.Securities[symbol_object].Price
            if price != 0:
                target_value = nav*target_weight
# How to account for the margin requirement?
                target_shares = int(target_value/price)
                current_shares = self.Portfolio[symbol_object].Quantity
                # Check for new position 
                if current_shares == 0:
                    tag = 'Initial entry'
                    new_positions.append((symbol_object, target_shares, tag))
                # Check for increasing the position size
                elif abs(target_shares) > abs(current_shares):
                    tag = 'Rebalance increase'
                    order_shares = target_shares-current_shares
                    increase_size.append((symbol_object, order_shares, tag))
                # Check for decreasing the position size
                elif abs(target_shares) < abs(current_shares):
                    tag = 'Rebalance decrease'
                    order_shares = target_shares-current_shares
                    decrease_size.append((symbol_object, order_shares, tag))

        # First handle positions that we are completely exiting
        positions = [
            (symbol, self.Portfolio[symbol].Quantity) \
            for symbol in self.Portfolio.Keys if self.Portfolio[symbol].Invested
        ]
        for symbol, current_qty in positions:
# AE update 12/6/23
            if str(symbol) not in target_allocation.keys() and str(symbol.ID) not in target_allocation.keys():
                # Liquidate the position
                self.MyLog(
                    f"{symbol} not in target allocation, so liquidating the "
                    f"open position of {current_qty} shares"
                )
                self.Liquidate(symbol, tag='exit')

# Skip if we are not on a new hour
        if self.Time.minute == 0 or (self.Time.minute == 59 and self.Time.second >= 50):
            pass
# AE adding 9/27/23
        elif self.Time.minute == 1:
            pass
        else:
            if DEBUG:
                self.MyLog(
                    f'Skipping rebalance until the top of the hour: {self.Time}'
                ) 
            return

        # Next handle orders where we are decreasing the overall position size
        # Only do this on the top of an hour, or if there are new positions
        for symbol, order_shares, tag in decrease_size:
            self.MyLog(
                f"{symbol} decreasing position size order for {order_shares} "
                "shares"
            )
            self.MarketOrder(symbol, order_shares, tag=tag)


        # Next handle new position orders
        for symbol, order_shares, tag in new_positions:
            self.MyLog(
                f"{symbol} new position order for {order_shares} shares"
            )
            self.MarketOrder(symbol, order_shares, tag=tag)  

        # Last handle orders where we are increasing the overall position size
        for symbol, order_shares, tag in increase_size:
            self.MyLog(
                f"{symbol} increasing position size order for {order_shares} "
                "shares"
            )
            self.MarketOrder(symbol, order_shares, tag=tag)

#-------------------------------------------------------------------------------
    def MyLog(self, message):
        """Add algo time to log if live trading. Otherwise just log message."""
        # Log all messages in live trading mode with local time added
        if self.LiveMode:
            self.Log(f'{self.Time}: {message}')
        else:
            self.Log(message)

#-------------------------------------------------------------------------------
    def ResubmitOrder(self, order, msg):
        """Built-in event handler for orders."""
        if type(order) == qc.Orders.MarketOrder \
        or type(order) == qc.Orders.MarketOnOpenOrder:
            order_type = 'Market'
        elif type(order) == qc.Orders.LimitOrder:
            order_type = 'Limit'
            # Get the limit price
            limit_price = order.LimitPrice
        else:
            self.MyLog(
                f"Invalid Order, but not a market or limit order! Order type="
                f"{type(order)}"
            )
            return

        # Get the order message, symbol, and qty
        self.MyLog(
            f"Invalid {order_type} Order! error: {msg}"
        )
        symbol = order.Symbol
        order_qty = int(order.Quantity)

        # Check for insufficient buying power
        if 'Insufficient buying power' in msg:
            # Get the initial margin and free margin
            initial_margin = float(
                msg.split("Initial Margin: ")[1].split(",")[0]
            )
            free_margin = float(
                msg.split("Free Margin: ")[1].split(",")[0].strip('.')
            )
            # Get the max allowed position size
            margin_per_share = abs(initial_margin/order_qty)
            max_shares = int(abs((0.95*free_margin/margin_per_share)))

        # Check for 'desired position your previous day...' error
        elif 'DESIRED POSITION YOUR PREVIOUS DAY EQUITY WITH LOAN VALUE' in msg:
            # Get the initial margin and previous day equity loan value
            initial_margin = float(
                msg.split("INITIAL MARGIN [")[1].split("USD")[0].replace(' ','')
            )
            loan_value = float(
                msg.split("LOAN VALUE [")[1].split("USD")[0].replace(' ','')
            )
            # Get the max allowed position size
            margin_per_share = abs(initial_margin/order_qty)
            max_shares = int(abs((0.95*loan_value/margin_per_share)))

        else:
            self.MyLog(f"Unrecognized error message: {msg}")
            self.MyLog(f"Will try to reduce order qty by 50%")
            # Try to cut order size in half
            max_shares = int(order_qty*0.5)

        # Get new qty
        if order_qty < 0:
            order_qty = -abs(max_shares)
        else:
            order_qty = abs(max_shares)
        if order_qty == 0:
            self.MyLog(
                f"Initial number of shares exceeds margin requirements! Reducing "
                f"order qty resulting in 0 share order, so it's ignored!"
            )
            return
        # Otherwise valid new order qty
        self.MyLog(
            f"Initial number of shares exceeds margin requirements! Reducing "
            f"order qty to {order_qty}"
        )
        # Resubmit an order with a reduced qty
        if order_type == 'Market':
            self.MarketOrder(symbol, order_qty, asynchronous=True)
        elif order_type == 'Limit':
            self.MyLog(f"Limit price={limit_price}")
            order = self.LimitOrder(symbol, order_qty, limit_price)
            # Add order to the list of limit orders
            # self.limit_orders.append(order)

#-------------------------------------------------------------------------------
    def OnSecuritiesChanged(self, changes):
        """Event handler for changes to our universe."""
        # Loop through securities added to the universe
        for security in changes.AddedSecurities:
            # Get the security symbol object
            symbol_object = security.Symbol
            # Add to our list of symbols
            if symbol_object not in self.symbol_objects:
                self.MyLog(
                    f'OnSecuritiesChanged() adding symbol object to list: '
                    f'{symbol_object.Value}->{symbol_object.ID}'
                ) 
                self.symbol_objects.append(symbol_object)
        # Loop through securities removed from the universe
        for security in changes.RemovedSecurities:
            # Get the security symbol object
            symbol_object = security.Symbol
            # Remove from our list of symbols
            if symbol_object in self.symbol_objects:
                self.symbol_objects.remove(symbol_object)

#-------------------------------------------------------------------------------
    def OnOrderEvent(self, orderEvent):
        """Built-in event handler for orders."""
        # Log message
        if self.LiveMode:
            self.MyLog(
                f"New order event: {orderEvent}, Status={orderEvent.Status}"
            )
        # Catch invalid order
        if orderEvent.Status == OrderStatus.Invalid:
            try:
                # Resubmit a new order
                order = self.Transactions.GetOrderById(orderEvent.OrderId)
                # ticket = self.Transactions.GetOrderTicket(orderEvent.OrderId)
                # response = ticket.GetMostRecentOrderResponse()
                msg = orderEvent.get_Message()
                self.ResubmitOrder(order, msg)
            except:
                if self.LiveMode:
                    self.MyLog(
                        f'OnOrderEvent() exception: {traceback.format_exc()}'
                    )

#-------------------------------------------------------------------------------
    # def OnEndOfAlgorithm(self):
    #     """Built-in event handler for end of the backtest."""
    #     # Save the portfolio values to the object store so we can evaluate
    #     #  them in the Research environment
    #     key = 'ZScore'
    #     d = {
    #         'time': self.times,
    #         'value': self.navs,
    #         'volatility': self.market_volatilities,
    #         'direction': self.market_directions
    #     }
    #     serialized = pickle.dumps(d)
    #     self.ObjectStore.SaveBytes(key, serialized)

###############################################################################
class CustomSecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(
        self, 
        brokerage_model: IBrokerageModel, 
        security_seeder: ISecuritySeeder
        ) -> None:
        super().__init__(brokerage_model, security_seeder)

    def Initialize(self, security: Security) -> None:
        """
        Define models to be used for securities as they are added to the 
        algorithm's universe.
        """
        # First, call the superclass definition
        # Sets the reality models of each security using the default models 
        #  of the brokerage model
        super().Initialize(security)
        # Define the buying power model to use for the security
        security.SetBuyingPowerModel(SecurityMarginModel(1.0))