#region imports
from AlgorithmImports import *
#endregion
# 0. LOADING OF THE LIBRARIES

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data import *
from datetime import timedelta

import numpy as np

class ScheduledCAPM(QCAlgorithm):

    # 1. INITIALIZATION
    def Initialize(self):
        self.SetStartDate(2022,9,1)
        self.SetEndDate(2025,4,15)  
        self._cash = 1000000
        self.SetCash(self._cash)                    
        
        # 2. PORTFOLIO CONSTRUCTION MODEL
        self._tickers =  ['AGG',
                        'IWM',
                        'IAU', 
                        'COMT',
                        'USMV',
                        'DGRO', 
                        'QUAL', 
                        'DVY', 
                        'MTUM',
                        'VLUE', 
                        'EFAV',
                        'EEMV', 
                        'IDV',
                        'IQLT', 
                        'IYW',
                        'IGF', 
                        'IYH']
         
       
        # 3. _benchING FOR REGRESSIONS FOR THE ALPHA MODEL
        self.symbols = [self.AddEquity(ticker).Symbol for ticker in self._tickers ] 
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.ivv = self.AddEquity("IVV", Resolution.Daily).Symbol
        self._bench = self.spy
        
        # 4. PARAMETRIZATION OF CLOCKS FOR ALPHA IMPLEMENTATION AND RISK MANAGEMENT
        self.lookback = 30
        self.SetWarmup(31)
        self.counter = 0
        self._secondCounter =0
        
        # 5. REFERENCE FOR PLOTTING
        self.reference = self.History(self.ivv, 5, Resolution.Daily)['close']
        self._initialValue = self.reference.iloc[0]
        
        # 6. EXECUTION MODEL (EVERY MONDAY)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen(self.spy, 10), self.Rebalance)    
        
        
        # 7. PARAMETERS FOR RISK MANAGEMENT
        self._portfolioValue = [self._cash]
        self._drawdown = -0.03
        
        # 8. _benchING FOR PLOTTING
        self.reference = self.History(self.spy, 10, Resolution.Daily)['close']
        self._initialValue = self.reference.iloc[0]
        
      
    def Rebalance(self):
        self.counter += 1
        self.Debug(f"Counter : {self.counter}")
        
        if self.IsWarmingUp: 
            return
        
        if self.counter % 3 == 0:
            
            history = self.History(
            self.symbols + [self._bench], 
            self.lookback,
            Resolution.Daily).close.unstack(level=0)
            
            self.symbols_alpha = self.SelectSymbols_alphas(history)
            self.symbols_beta = self.SelectSymbols_betas(history)
            
            for holdings in self.Portfolio.Values:
                symbol = holdings.Symbol
            
            # 9. ALPHA MODEL    
            if symbol not in self.symbols_alpha and symbol not in self.symbols_beta and holdings.Invested:
                self.Liquidate(symbol)
            # 10. PORTFOLIO CONSTRUCTION MODEL    
            for symbol in self.symbols_alpha:
                self.SetHoldings(symbol, 0.1)
            
            
            for symbol in self.symbols_beta:
                if symbol in self.symbols_alpha:
                    self.SetHoldings(symbol, 0.2)
                else:
                   self.SetHoldings(symbol, 0.1) 
                   
            self.SetHoldings("SPY", 1.0)
       
    def OnData(self,data):
        # 11. RISK MANAGEMENT   
        self._secondCounter +=1
        self._portfolioValue.append(self.Portfolio.TotalPortfolioValue)
        
        if self._secondCounter % 5 == 0:
            if (self._portfolioValue[-1]-self._portfolioValue[-5])/self._portfolioValue[-1] < self._drawdown:
                self.Liquidate()
            
        self.Plot("Relative Performance", "_bench", self._cash*self.Securities["SPY"].Close/self._initialValue)
        self.Plot("Relative Performance", "Total Portfolio Value", self.Portfolio.TotalPortfolioValue)
        
        
   # 12. AUXILIARY METHODS FOR CALCULATING ALPHAS AND BETAS
    def SelectSymbols_alphas(self, history):
        alphas = dict()
        _bench = history[self._bench].pct_change().dropna()

        for symbol in self.symbols:
            returns = history[symbol].pct_change().dropna()
            bla = np.vstack([_bench, np.ones(len(returns))]).T
            result = np.linalg.lstsq(bla , returns)
            alphas[symbol] = result[0][1]
            

        selected_alphas = sorted(alphas.items(), key=lambda x: x[1], reverse=True)[:5]
        return [x[0] for x in selected_alphas]
        
    def SelectSymbols_betas(self, history):
        betas = dict()
        _bench = history[self._bench].pct_change().dropna()

        for symbol in self.symbols:
            returns = history[symbol].pct_change().dropna()
            bla = np.vstack([_bench, np.ones(len(returns))]).T
            result = np.linalg.lstsq(bla , returns)
            betas[symbol]= result[0][0]

        selected_betas = sorted(betas.items(), key=lambda x: x[1], reverse=False)[:5]
        return [x[0] for x in selected_betas]