import numpy as np
from scipy import stats
from collections import deque
#region imports
from AlgorithmImports import *
#endregion
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

class NadionTransdimensionalAutosequencers(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 10, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        self.SetBenchmark("SPY")
        self.UniverseSettings.Leverage = 0
        self.Settings.RebalancePortfolioOnInsightChanges = True         
        self.Settings.RebalancePortfolioOnSecurityChanges = True
        self.SetExecution(ImmediateExecutionModel())
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.UniverseSettings.Resolution=Resolution.Daily
        self.SetUniverseSelection( QC500UniverseSelectionModel() )
        self.SetAlpha(MyAlpha())

         
class MyAlpha(AlphaModel):

    def __init__(self):
        self.securities = []
        self.symbol_data_by_symbol = {}
        self.sorted_mom = []
        self.mom_scores = {}
        self.num_stocks = 30

    def Update(self, algorithm, data):
        insight = []
        for k,v in  self.symbol_data_by_symbol.items():   
            #algorithm.Debug(str(repr(v.mom)))
            if not v.mom.IsReady:
                return
            if v.mom.Score >= 40:
                self.mom_scores[k] = v.mom.Score
                # algorithm.Debug(f"{k}: {v.mom.Score}")     
        self.sorted_mom = sorted([k for k,v in self.mom_scores.items()],
            key=lambda x: self.mom_scores[x], reverse=True)
        
        self.selected = self.sorted_mom[:self.num_stocks]
        for security in self.selected:
            insight += [(Insight.Price(security, timedelta(days = 1), InsightDirection.Up))]
        #algorithm.Debug(f"{algorithm.Time}:update insights")

        return insight    
    
    def OnSecuritiesChanged(self, algorithm, changes):
        #algorithm.Debug(f"{algorithm.Time}:security changes called")
        #algorithm.Debug(f"removed: {len(changes.RemovedSecurities)}")
        #algorithm.Debug(f"added: {len(changes.AddedSecurities)}")
        
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(algorithm, security.Symbol)
            #algorithm.Debug(f"{algorithm.Time}: Added {security.Symbol}")

        for security in changes.RemovedSecurities:
            if security.Symbol in self.symbol_data_by_symbol:
                symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
                if symbol_data:
                    symbol_data.dispose()
            #algorithm.Debug(f"{algorithm.Time}: Removed {security.Symbol}")
            if security in self.securities:
                self.securities.remove(security)
                             
        self.securities.extend(changes.AddedSecurities)
        
        for security in self.securities:
            if security.Symbol not in algorithm.ActiveSecurities:
                algorithm.Debug(f"{security} not in active but is in self")
                self.securities.remove(security)
                symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
                if symbol_data:
                    symbol_data.dispose()
                    algorithm.Debug(f"{security.Symbol} data removed")
            
                    
class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        self.mom = CustomMomentum("momentum", 125)
        self.consolidator = TradeBarConsolidator(1)
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        algorithm.RegisterIndicator(symbol, self.mom, self.consolidator)
        algorithm.WarmUpIndicator(self.symbol, self.mom)
        
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)


class CustomMomentum(PythonIndicator):
    def __init__(self, name, period):
        self.Name = name
        self.WarmUpPeriod = period
        self.Time = datetime.min
        self.Value = 0
        self.queue = deque(maxlen=period)

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

    def Update(self, input: BaseData) -> bool:
        self.queue.appendleft(input.Value)  
        x = np.arange(len(self.queue))
        log_ts = np.log(self.queue)
        slope, intercept, r_value, p_value, std_err = stats.linregress(x, log_ts)
        annualized_slope = (np.power(np.exp(slope), 252) - 1) * 100
        self.Time = input.Time
        self.Score = float(annualized_slope * (r_value**2))
        return x == self.queue.maxlen