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.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        self.UniverseSettings.Leverage = 0
        self.UniverseSettings.Resolution=Resolution.Daily
        self.SetUniverseSelection( QC500UniverseSelectionModel() )

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: Expiry.EndOfWeek(time)))

        self.SetExecution(ImmediateExecutionModel())
        
        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 not v.mom.IsReady:
                continue
            if v.mom.Score >= 40:
                self.mom_scores[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, Expiry.EndOfWeek, InsightDirection.Up))]
        self.rebalance = Expiry.EndOfWeek(self.Time)

        self.EmitInsights(insight)



    def OnSecuritiesChanged(self, changes):        
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(self, 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()
            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.securities.remove(security)
                symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
                if symbol_data:
                    symbol_data.dispose()
            
                    
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