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.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: Expiry.EndOfWeek(time)))
        
        self.SetExecution(ImmediateExecutionModel())
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        #self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.UniverseSettings.Resolution=Resolution.Daily
        
        self.SetUniverseSelection( QC500UniverseSelectionModel() )
        #self.SetAlpha(MyAlpha())
        self.rebalance = self.Time
        self.securities = []
        self.symbol_data_by_symbol = {}
        self.sorted_mom = []
        self.mom_scores = {}
        self.num_stocks = 30

    def OnData(self, data):

        insight = []

        if self.rebalance >= self.Time:
            return
        self.mom_scores = {}
        for k,v in  self.symbol_data_by_symbol.items():

            if v.mom.Score >= 40:
                self.mom_scores[k] = v.mom.Score
                if k.Value == "RIDE":
                    self.Debug(f"{k}: slope: {v.mom.slope} r: {v.mom.rvalue} annualized_slope: {v.mom.annualized_slope}")
                #self.Debug(f"{k}: {len(v.mom.queue)}")
                #self.Debug(f"{k}: {self.mom_scores[k]}")  
        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, Expiry.EndOfWeek, InsightDirection.Up))]
        #self.Debug(f"{algorithm.Time}:update insights")
        self.rebalance = Expiry.EndOfWeek(self.Time)

        self.EmitInsights(insight)



    def OnSecuritiesChanged(self, 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(self, 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 self.ActiveSecurities:
                self.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()
                    self.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.Score = 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)
        self.log_ts = np.log(self.queue)
        x = np.arange(len(self.log_ts))
        result = stats.linregress(x, self.log_ts)
        self.annualized_slope = (np.power(np.exp(result.slope), 252) - 1) * 100
        self.slope = result.slope
        self.rvalue = result.rvalue
        self.Time = input.Time
        self.Score = self.annualized_slope * (result.rvalue**2)
        return len(self.queue) == self.queue.maxlen

'''
         
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 = []

        if self.rebalance >= self.Time:
            return

        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")
        self.rebalance = Expiry.EndOfMonth(self.Time)

        return insight  

'''