#region imports
from AlgorithmImports import *
#endregion
# https://quantpedia.com/Screener/Details/14

import numpy as np
import pandas as pd
from scipy.stats import linregress
from collections import deque
from datetime import timedelta
from datetime import datetime
import math

# TODO: add volatility based position sizes - risk parity - DONE
# TODO: add rebalancing based on risk parity every month
# TODO: remove securities if they fall out of the top momentum ranking AND fall under the sma filter
# TODO: fix margin issue / errors - DONE (mostly)
# TODO: try filtering on momentum, close > sma and rsi(3)  
# TODO: add in the money puts when index < 200 sma based on momentum rank 

class MomentumRanker(PythonIndicator):
    
    def momentum(self, closes):
        returns = np.log(closes)
        x = np.arange(len(returns))
        slope, _, rvalue, _, _ = linregress(x, returns)
        return ((1 + slope) ** 252) * (rvalue ** 2)  # annualize slope and multiply by R^2

    def __init__(self, period):
        self.Time = datetime.min
        self.period = period
        self.closes = deque(maxlen=period)  # queue for storing closes
        self.Value = 0
    
    def __repr__(self):
        return f"Momentum Ranker -> IsReady: {self.IsReady}. Time: {self.Time}. Value: {self.Value}." 

    def IsReady(self):
        return (self.Value > 0) # return(len(self.closes)==self.closes.maxlen)
    
    # TODO: change this to IndicatorDataPoint
    def Update(self, bar:TradeBar):
        self.closes.append(bar.Close)
        self.Time = bar.Time
    
        self.Value = 0

        if len(self.closes) == self.closes.maxlen:
            self.Value = self.momentum(np.array([*self.closes])) # convert deque to list then to np.array
        
        return (self.Value > 0)

class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol

        self.filter_period = 100
        self.volatility_period = 20
        self.momentum_period = 90 

        self.momentum = MomentumRanker(self.momentum_period)
        self.atr = AverageTrueRange(symbol, self.volatility_period)
        self.filter = SimpleMovingAverage(symbol, self.filter_period)

        self.consolidator = TradeBarConsolidator(1)
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        algorithm.RegisterIndicator(self.symbol, self.momentum)
        algorithm.RegisterIndicator(self.symbol, self.filter, self.consolidator) # maybe or maybe like momentum?
        algorithm.RegisterIndicator(self.symbol, self.atr, self.consolidator)
        
        algorithm.WarmUpIndicator(self.symbol, self.atr)
        algorithm.WarmUpIndicator(self.symbol, self.filter)

        # warmup momentum indicators
        trade_bars = self.algorithm.History[TradeBar](symbol, self.momentum_period, Resolution.Daily)
        for bar in trade_bars:
            self.momentum.Update(bar)
    
    def UpdateIndicators(self, data:TradeBar): 
        # update the indicators
        self.momentum.Update(data)
        self.filter.Update(data.Time, data.Close)  # SMA takes time, IndicatorDataPoint, may need to construct and IndicatorDataPoint
        self.atr.Update(data)                   

    def IsReady(self):
        return self.momentum.IsReady and self.atr.IsReady and self.filter.IsReady
    
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)
   

class MomentumEffectAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(1998, 1, 1)  # Set Start Date
        self.SetEndDate(2023, 3, 1)    # Set End Date       
        self.SetCash(100000)           # Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Adjusted
        self.UniverseSettings.MinimumTimeInUniverse = timedelta(days=28)

        self.symbol_dictionary = {} # dictionary for holding SymbolData key by symbol
        self.addedSymbols = []
        self.removedSymbols = []

        # adjust these in line with Clenow - want approx 20 stocks in the portfolio
        self.num_coarse = 500 # Number of symbols selected at Coarse Selection
        self.num_fine = 100    # Number of symbols selected at Fine Selection
        self.num_positions = 20     # Number of symbols with open positions
        self.risk_factor = 0.001    # targeting 10 basis point move per day
        
        self.index_filter_period = 200

        # variables to control the portfolio rebalance and the stock selection reranking
        self.month = -1
        self.rebalance = False
        self.rerank = False

        # create the equities universe
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

        self.index = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.SetBenchmark(self.index)

        self.index_sma = SimpleMovingAverage(self.index, self.index_filter_period)
        self.RegisterIndicator(self.index, self.index_sma, Resolution.Daily)

        # set Brokerage model and Fee Structure
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # set Free Cash to 2.5%
        self.Settings.FreePortfolioValuePercentage = 0.025
        
        """
        # can't use scheduling with coarse / fine universe selection
        # but TODO: can probably schedule the rerank and rebalance methods, to run after OnData
        self.Schedule.On(self.DateRules.MonthEnd(self.index, daysOffset = 0),
                         self.TimeRules.BeforeMarketClose("SPY", 60),
                         self.RerankMomentum)
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday),
                         self.TimeRules.BeforeMarketClose("SPY", 120),
                         self.RebalancePortfolio)
        
        # Selection will run on every week on Monday 2 hours before market close
        self.AddUniverseSelection(ScheduledUniverseSelectionModel(
                                self.DateRules.Every(DayOfWeek.Monday),
                                self.TimeRules.AfterMarketClose("SPY", 15),
                                self.SelectSymbols))
        
        """

    # my methods

    def CalculateRiskParityPositionSize(self, symbol, atr)-> float:
        amountToRisk = self.Portfolio.TotalPortfolioValue * self.risk_factor
        quantity = self.CalculateOrderQuantity(symbol, 0.01) # get the quanity of shares for holding of 1%
        shares = amountToRisk / atr
        rounded_holding = 0.0
        if quantity != 0:
            holding_percent = (shares/quantity)/100
            rounded_holding = round(holding_percent, 3) # round down to 0.001
        return rounded_holding

    # QC methods
    def OnWarmUpFinished(self) -> None:
        self.Log("Equities Momentum Algorithm Ready")

    
    # this article explains when universe selection is done (12:00am each day) and the order of the calls
    # coarse, fine, onsecuritieschanged, then ondata
    # https://www.quantconnect.com/forum/discussion/6485/onsecuritieschanged-questions/p1

    # may eventually switch this out for small or mid cap stocks
    def CoarseSelectionFunction(self, coarse):
        '''Drop securities which have no fundamental data or have too low prices. Select those with highest by dollar volume'''
        
        if not self.Time.weekday() == 0: # monday
            return Universe.Unchanged

        self.rerank = True

        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 5], key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in selected[:self.num_coarse]]

    # may not really want to do anything here - might be that coarse filtering is enough
    def FineSelectionFunction(self, fine):
        '''Select security with highest market cap, and common stock'''
        selected = sorted(fine, key=lambda f: f.MarketCap and 
                                              f.SecurityReference.SecurityType == 'ST00000001', # stock is common stock, ST00000002 is preferred stock, 
                                              reverse=True)

        return [x.Symbol for x in selected[:self.num_fine]]
     

    def OnData(self, data):
        
        liquidated = []
        purchased = []

        # update the indidcators with new bar data
        for symbol, symbolData in self.symbol_dictionary.items():
            if data[symbol] is not None:
                symbolData.UpdateIndicators(data[symbol])

        if data.Time == datetime(1998,2,23): 
            self.Log("Time based breakpoint opportunity")

        # rerank weekly, rebalance monthly
        if not self.rerank:
            return

        if (self.month != self.Time.month):
            self.rebalance = True
            self.month = self.Time.month

        # rerank securities based on momentum and close above sma filter
        mom_list = [sym for sym, symb_data in self.symbol_dictionary[sym] if symb_data.momentum.IsReady]
        sorted_mom = sorted(mom_list, key=lambda x: self.symbol_dictionary[x].SymbolData.momentum.Current.Value, reverse=True)

        # filter selected to chose only those with a price above their sma 100 
        sorted_filter = sorted([sym for sym, symb_data in self.symbol_dictionary[sym] if symb_data.filter.IsReady],
            key=lambda x: data[x].Close > self.symbol_dictionary[x].SymbolData.filter.Current.Value)

        combined = list((set(sorted_mom)) & set(sorted_filter))
        sorted_combined = sorted(combined,key=lambda x: self.symbol_dictionary[x].SymbolData.momentum.Current.Value, reverse=True)
        selected = sorted_combined[:self.num_positions] 

        # Sell
        # Liquidate securities that are in our portfolio, but no longer in the top momentum rankings, and have closed below the filter
        for security in self.Portfolio:
            symbol = security.Symbol
            if symbol not in selected:
                if data[symbol].Close < self.symbol_dictionary[symbol].filter.Current.Value:
                    self.Liquidate(symbol, f'{symbol} {data[symbol].Close} < sma(100): {self.filter[symbol].Current.Value}')
                    liquidated.append(symbol.ID.ToString().split(' ')[0])
                      
        # log which ones were removed
        self.Log(f'{liquidated} removed by sma filter')

        # Buy
        # if we pass our index filter, buy selected securities (check not already in holdings first though)
        if data[self.index].Close > self.index_sma.Current.Value:
            buffer = (self.Portfolio.TotalPortfolioValue / 100) # 1% buffer
            for symbol in selected:
                if (not symbol in self.Portfolio):
                    remaining_margin = self.Portfolio.MarginRemaining
                    cash = self.Portfolio.Cash # might use cash instead of margin
                    if remaining_margin > buffer:
                        holding = self.CalculateRiskParityPositionSize(symbol, self.volatility[symbol].Current.Value)
                        self.SetHoldings(symbol, holding)
                        purchased.append(f"{symbol.ID.ToString().split(' ')[0]} {holding}")

        # log which ones were purchased
        self.Log(f'{purchased} added based on momentum rankings')

        portfolio_update = f'Value: {self.Portfolio.TotalPortfolioValue}, Cash: {self.Portfolio.Cash}, Margin Used: {self.Portfolio.TotalMarginUsed}, Margin Remaining {self.Portfolio.MarginRemaining}'
        self.Log(f'Portfolio: {portfolio_update}')

        if not self.rebalance:
            return

        rebalanced = []

        for symbol in self.Portfolio:
            holding = self.CalculateRiskParityPositionSize(symbol, self.symbol_dictionary[symbol].atr.Current.Value)
            self.SetHoldings(symbol, holding)
            rebalanced.append(f"{symbol.ID.ToString().split(' ')[0]} {holding},")

        self.Log(f'Portfolio Weights: {rebalanced}')

        self.rebalance = False
        self.rerank = False

 
    def OnSecuritiesChanged(self, changes):
        # Clean up securities list and indicator data for removed securities
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.symbol_dictionary:
                self.symbol_dictionary.pop(symbol) 
                self.removedSymbols.append(symbol.ID.ToString())

        # Create indicators and warm them up for securities newly added to the universe
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.symbol_dictionary: 
                self.symbol_dictionary[symbol] = SymbolData(self, symbol)
                self.addedSymbols.append(symbol.ID.ToString())

        self.Log(f'Added: {self.addedSymbols}')
        self.Log(f'Removed: {self.removedSymbols}')