#region imports
from AlgorithmImports import *
#endregion

class UncorrelatedAssetsAlphaModel(AlphaModel):

    securities = []
    previous_insight_symbols = []
    symbols_to_liquidate  = []
    month = -1

    def __init__(self, portfolio_size, lookback_days):
        self.portfolio_size = portfolio_size
        self.lookback_days = lookback_days

    def Update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        # Rebalance monthly when there is QuoteBar data available
        if self.month == algorithm.Time.month or data.QuoteBars.Count == 0:
            return []

        # Select the assets we can trade
        tradable_symbols = [security.Symbol for security in self.securities if security.Symbol in data.QuoteBars and security.Price > 0]

        # Fetch history on the tradable assets
        history = algorithm.History(tradable_symbols, self.lookback_days, Resolution.Daily, dataNormalizationMode=DataNormalizationMode.ScaledRaw)
        if history.empty: 
            return []

        self.month = algorithm.Time.month
    
        # Compute daily returns
        returns = history['close'].unstack(level=0).pct_change().iloc[1:]
    
        # Get correlation
        correlation = returns.corr()
        
        # Find the subset of assets with lowest absolute sum correlation
        selected = []
        for index, row in correlation.iteritems():
            corr_rank = row.abs().sum()
            selected.append((index, corr_rank))
        most_uncorrelated_symbols = [x[0] for x in sorted(selected, key = lambda x: x[1])[:self.portfolio_size]]
        
        # Create insights to long the assets with the most uncorrelated symbols
        insights = []
        expiry = self.securities[0].Exchange.Hours.GetNextMarketOpen(Expiry.EndOfMonth(algorithm.Time), extendedMarketHours=False) - timedelta(minutes=1)
        for symbol in most_uncorrelated_symbols:
            insights.append( Insight.Price(symbol, expiry, InsightDirection.Up) )

        return insights

    def OnSecuritiesChanged(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.RemovedSecurities:
            if security in self.securities:
                self.securities.remove(security)
                
        self.securities.extend(changes.AddedSecurities)
                

# region imports
from AlgorithmImports import *
from alpha import UncorrelatedAssetsAlphaModel
from portfolio import ExtendedEqualWeightingPortfolioConstructionModel
from universe import MostLiquidUniverseSelectionModel

# endregion

class UncorrleatedAssetsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2012, 1, 1)
        self.SetEndDate(2023, 12, 1)
        self.SetCash(1_000_000)

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Adjusted
        # self.AddUniverseSelection(MostLiquidUniverseSelectionModel(self.GetParameter("universe_size", 100), self.UniverseSettings))
        self.yinn = self.AddEquity('YINN', Resolution.Daily)
        self.yang = self.AddEquity('YANG', Resolution.Daily)

        # self.AddAlpha(UncorrelatedAssetsAlphaModel(
        #     self.GetParameter("portfolio_size", 10),
        #     self.GetParameter("lookback_days", 252),
        # ))
        self.threshold = self.GetParameter('threshold', 0.2)

        # self.Settings.RebalancePortfolioOnSecurityChanges = True
        # self.SetPortfolioConstruction(ExtendedEqualWeightingPortfolioConstructionModel(lambda time: None))

        #self.AddRiskManagement(NullRiskManagementModel())

        #self.SetExecution(ImmediateExecutionModel()) 

    def OnData(self, slice: Slice):
        if not self.Portfolio.Invested:
            self.SetHoldings("YINN", -0.5)
            self.SetHoldings("YANG", -0.5)
        else:
            port_val = self.Portfolio.TotalHoldingsValue
            yinn_val = self.Portfolio['YINN'].HoldingsValue
            yang_val = self.Portfolio['YANG'].HoldingsValue
            log_exposure = (3)*yinn_val/port_val + (-3)*yang_val/port_val
            if self.Time.year==2016 and self.Time.month==3 and self.Time.day >= 15:
                self.Debug(f'{self.Time} port val: {port_val}, yinn_val: {yinn_val}, yang_val: {yang_val}\
                            exp: {log_exposure}, yang_price: {slice["YANG"].Close}, yinn_price: {slice["YINN"].Close}'  )

            if abs(log_exposure) > self.threshold:
                self.Debug('Rebalancing...')
                self.SetHoldings("YINN", -0.5)
                self.SetHoldings("YANG", -0.5)
        

#region imports
from AlgorithmImports import *
#endregion

class ExtendedEqualWeightingPortfolioConstructionModel(EqualWeightingPortfolioConstructionModel):

    target_collection = PortfolioTargetCollection()
    symbols_to_liquidate = []

    def CreateTargets(self, algorithm: QCAlgorithm, insights: List[Insight]) -> List[PortfolioTarget]:
        portfolio_targets = list(super().CreateTargets(algorithm, insights))
        self.target_collection.AddRange(portfolio_targets)

        # Flag securities as to-be-liquidated if they are no longer targeted
        if len(insights) > 0:
            targeted_symbols = [insight.Symbol for insight in insights]
            for target in self.target_collection:
                if target.Symbol not in targeted_symbols:
                    self.symbols_to_liquidate.append(target.Symbol)
        
        # Liquidate some securities that are flagged as to-be-liquidated
        liquidated_symbols = []
        for symbol in self.symbols_to_liquidate:
            security = algorithm.Securities[symbol]
            if security.IsDelisted:
                liquidated_symbols.append(symbol)

            elif security.Price > 0 and symbol in algorithm.CurrentSlice.QuoteBars:
                liquidated_symbols.append(symbol)
                portfolio_targets.append(PortfolioTarget.Percent(algorithm, symbol, 0))

        # Remove targets from the PortfolioTargetCollection if the security was delisted or liquidated
        for symbol in liquidated_symbols:
            self.target_collection.Remove(symbol)
            self.symbols_to_liquidate.remove(symbol)

        return portfolio_targets

        

#region imports
from AlgorithmImports import *
#endregion

class MostLiquidUniverseSelectionModel(CoarseFundamentalUniverseSelectionModel):
    def __init__(self, universe_size: int, universe_settings: UniverseSettings) -> None:
        self.universe_size = universe_size
        super().__init__(self.SelectCoarse, universe_settings)

    def SelectCoarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        # Select the most liquid stocks
        sorted_by_dollar_volume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)
        return [c.Symbol for c in sorted_by_dollar_volume[:self.universe_size]]