from itertools import groupby
from AlgorithmImports import *
import pandas as pd
import numpy as np
import scipy
from scipy import stats

class SectorMomentumTest(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2010, 1, 1)  # Set Start Date
        self.SetEndDate(2022, 9, 21)  # Set End Date
        self.SetCash(100000000)  # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.Leverage = 2.0
        self.RebalanceDays = 0
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(self.SelectCoarse, self.SelectFine))
        self.universe: List[Symbol] = []
        self.dataHistory: dict[Symbol, RollingWindow[float]] = {}
        self.SetBenchmark(self.spy) 
    def SelectCoarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        return [x.Symbol for x in coarse if x.DollarVolume > 0]
    def SelectFine(self, fine: List[FineFundamental]) -> List[Symbol]:
        fine_ = [x for x in fine if x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                    and x.CompanyReference.IsLimitedPartnership == 0
                                    and x.SecurityReference.IsPrimaryShare == 1
                                    and x.MarketCap > 0]
        sectorGroups = {key: list(group) for key, group in groupby(sorted(fine_, key=lambda x: x.AssetClassification.MorningstarSectorCode), key=lambda x: x.AssetClassification.MorningstarSectorCode)}
        sectorValuations = {}
        for sector, group in sectorGroups.items():
            sectorValuations.update(self.MomentumScore(group))
        valuationSorted = sorted(fine_, key=lambda x: sectorValuations[x.Symbol], reverse=True)[:round(len(fine_) / 5)]
        self.universe = [x.Symbol for x in valuationSorted]
        return self.universe
    def MomentumScore(self, fine: List[FineFundamental]) -> Dict[Symbol, float]:
        for x in fine:
            if x.Symbol in self.dataHistory:
                self.dataHistory[x.Symbol].Add(x.Price)
        addedSymbols = [x.Symbol for x in fine if x.Symbol not in self.dataHistory]
        if len(addedSymbols) > 0:
            history: List[TradeBars] = self.History[TradeBar](addedSymbols, 253, Resolution.Daily)
            newData = {symbol: RollingWindow[float](252) for symbol in addedSymbols}
            for slice in history:
                for symbol, bar in slice.items():
                    newData[symbol].Add(bar.Close)
            self.dataHistory.update(newData)
        activeSymbols = {x.Symbol: x for x in fine}
        inactiveSymbols = [symbol for symbol in self.dataHistory.keys() if symbol not in activeSymbols]
        for symbol in inactiveSymbols:
            self.dataHistory.pop(symbol, None)
        twelveMonthScore = stats.rankdata([(self.dataHistory[x.Symbol][0] - self.dataHistory[x.Symbol][251]) / self.dataHistory[x.Symbol][251] if x.Symbol in self.dataHistory and self.dataHistory[x.Symbol].IsReady else -np.inf for x in fine]) / len(fine)
        nineMonthScore = stats.rankdata([(self.dataHistory[x.Symbol][0] - self.dataHistory[x.Symbol][188]) / self.dataHistory[x.Symbol][188] if x.Symbol in self.dataHistory and self.dataHistory[x.Symbol].IsReady else -np.inf for x in fine]) / len(fine)
        sixMonthScore = stats.rankdata([(self.dataHistory[x.Symbol][0] - self.dataHistory[x.Symbol][125]) / self.dataHistory[x.Symbol][125] if x.Symbol in self.dataHistory and self.dataHistory[x.Symbol].IsReady else -np.inf for x in fine]) / len(fine)
        threeMonthScore = stats.rankdata([(self.dataHistory[x.Symbol][0] - self.dataHistory[x.Symbol][62]) / self.dataHistory[x.Symbol][62] if x.Symbol in self.dataHistory and self.dataHistory[x.Symbol].IsReady else -np.inf for x in fine]) / len(fine)
        
        compositeScore = np.sum([twelveMonthScore, nineMonthScore, sixMonthScore, threeMonthScore], axis=0)
        compositeRank = stats.rankdata(compositeScore) / len(fine)        
        result: Dict[Symbol, float] = {x.Symbol: compositeRank[i] for i, x in enumerate(fine)}
        return result      
    def OnSecuritiesChanged(self, changes: SecurityChanges):
        pass       
    def DailyRoutine(self):
        if self.RebalanceDays > 0:
            self.RebalanceDays -= 1
            return
        if self.RebalanceDays == 0:
            numSecurities = len(self.universe)
            targetHoldings = [PortfolioTarget(symbol, 0.995/numSecurities) for symbol in self.universe]
            for key in self.Portfolio.keys():
                if key not in self.universe and self.Portfolio[key].Invested:
                    targetHoldings.append(PortfolioTarget(key, 0.0))
            self.SetHoldings(targetHoldings)
            self.RebalanceDays = 21 
    def OnData(self, data):
        self.DailyRoutine()
        self.Plot('Days', 'Rebalance', self.RebalanceDays)
        self.Plot('Number of Securities', 'Universe', len(self.universe))
        self.Plot("Number of Securities", "Invested", len([x for x in self.Portfolio.Values if x.Invested]))
    def OnOrderEvent(self, orderEvent: OrderEvent) -> None:
        if orderEvent.Status == OrderStatus.Invalid and orderEvent.Direction == OrderDirection.Buy:
            self.MarketOrder(orderEvent.Symbol, orderEvent.Quantity - 1)