from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data import *
from datetime import timedelta
from collections import deque
from System.Collections.Generic import List
from QuantConnect.Data.UniverseSelection import *
from AlgorithmImports import *
from collections import deque
from joblib import Parallel, delayed
import pandas as pd

class QCUM(QCAlgorithm):
    def Initialize(self):
        
        # QCAlgorithm parameters

        ## Backtest start/end dates
        self.SetStartDate(2008, 1,1)
        self.SetEndDate(2023, 7, 15)
        
        ## Init Cash
        self.cap = 100000
        self.SetCash(self.cap)
        
        # Benchmark
        self.SetBenchmark("SPY")
        self.benchmark = ['SPY']
        
        # Bonds
        self.bonds = ["XLU", 'GLD', 'IEF', 'TLT']
        
        # Algorithm parameters 
        self.MarketCap = 50000000 # 50M


        self.TARGET_SECURITIES = 30
        self.TF_LOOKBACK = 100
        self.TF_CURRENT_LOOKBACK = 1

        
        self.CoarseSymbols = []
        self.FineSymbols = []
        self.SelectedSymbols = []
        self.stocks_to_hold =[]
        self.NeedReSelect = True

        
        self.spy_ten_days = False
        self.spy_five_days = False
        self.spy_ten_days_int = 10 #int(self.GetParameter("spy_ten_days_int"))
        self.spy_five_days_int = 5 #int(self.GetParameter("spy_five_days_int"))
        
        self.portfolioTargetCollect={}
        self.TrendUpAlpha = True

        # Universe
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)


        # Brokerage model and account type:
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        self.DefaultOrderProperties.TimeInForce = TimeInForce.Day

        # Bond 
        self.ETF_List = [
            "XLU",  # Utilities Select Sector SPDR Fund
            'GLD',
            'IEF',
            'TLT'
        ]
        self.ETFs = []
        
        for ticker in self.benchmark + self.bonds :
            symbol = self.AddEquity(ticker)
            symbol.SetDataNormalizationMode(DataNormalizationMode.Raw)
        

        ## RISK CONTROL
        self.Schedule.On(self.DateRules.EveryDay(),self.TimeRules.AfterMarketOpen("SPY",0), self.WarningCheck)
        ## REBALANCE 
        self.Schedule.On(self.DateRules.WeekStart(),self.TimeRules.AfterMarketOpen("SPY",122), self.Rebalance)


    def CoarseSelectionFunction(self, coarse):
        if not self.NeedReSelect:
            return Universe.Unchanged
        self.CoarseSymbols = [x.Symbol for x in coarse if x.HasFundamentalData ]        
        return self.CoarseSymbols

    def selection(self, x): 
        if x.MarketCap > self.MarketCap \
            and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS","ASE"] \
            and x.AssetClassification.ValueScore > 0 :
            return True
        else:
            return False 
    
    def M1(self, x):
        return  x.AssetClassification.ValueScore / x.Price

    def FineSelectionFunction(self, fine):
        if not self.NeedReSelect:
            return Universe.Unchanged

        m1_list = sorted([x for x in fine if self.selection(x)], key = lambda x: self.M1(x), reverse=True)[:self.TARGET_SECURITIES]
                                    
        self.SelectedSymbols = [x for x in m1_list]                            
        self.FineSymbols = [x.Symbol for x in m1_list]
        self.NeedReSelect = False

        return self.FineSymbols
    
    def WarningCheck(self):
        spy_close = self.History(
            self.benchmark, 
            15,
            Resolution.Daily).close.unstack(level=0)
        self.spy_ten_days =  (spy_close["SPY"][-1])/ spy_close["SPY"][-1 - self.spy_ten_days_int] < 0.90    
        self.spy_five_days =  (spy_close["SPY"][-1])/ spy_close["SPY"][-1 - self.spy_five_days_int] < 0.95
        if self.spy_five_days or self.spy_ten_days:
            self.Liquidate()
        
        
    def TrendUp(self):
        spy_fast = self.History(self.benchmark, self.TF_CURRENT_LOOKBACK, Resolution.Daily).close.mean()
        spy_slow = self.History(self.benchmark, self.TF_LOOKBACK, Resolution.Daily).close.mean()
        return spy_fast > spy_slow            

    def adjust(self, symbol: str , adjust_amount: float):
        #self.Log("adjust")
        price = self.Portfolio[symbol].Price
        if price == 0.0:
            #self.Log(f"price zero!!:symbol {symbol} self.Portfolio[symbol] {self.Portfolio[symbol]}, price {price}")
            return 0.0          
        try:
            qty = int (adjust_amount/price)
            direction = "Buy"
            if qty >= 0.0 :
                direction = "Buy"
            else:
                direction = "Sell"
            self.MarketOrder(symbol, qty)
            #self.Log(f"{direction} stock: {symbol} with qty {qty}, price: {price} amount: {qty * price}")
        except Exception as e:
                self.Log(f"{direction} exception {e}:symbol {symbol} self.Portfolio[symbol] {self.Portfolio[symbol]}, price {price}")
                return 0.0
        return abs(qty) * price

    def balance(self, stock : str , target_amount: float):
        #self.Log(f"balance {stock}")
        cur_amount = 0.0
        # check if the stock is in the position
        old_stocks = [x.Key for x in self.Portfolio if x.Value.Invested]
        if stock in old_stocks:
            cur_amount = self.Portfolio[stock].Quantity * self.Portfolio[stock].Price
        diff_amount = target_amount - cur_amount
        #self.Log(f"{stock} target_amount: {target_amount} cur_amount: {cur_amount} diff_amount: {diff_amount}")
        adjust_amount = self.adjust(stock, diff_amount)
        return adjust_amount
    
    def diff(self,li1, li2):
        li_dif = [i for i in li1 if i not in li2]
        return li_dif
    
    def diff_stocks(self):
        self.Log("diff_stocks")
        # get current stocks in positions
        current_stocks = [x.Key for x in self.Portfolio if x.Value.Invested]
        for stock in current_stocks:
            self.Log(f"Hold Equity: {stock}, Qty: {self.Portfolio[stock].Quantity}")
        # get new stocks read from file
        new_stocks = [x for x in self.portfolioTargetCollect.keys()] 
        for stock in new_stocks:
            self.Log(f"Target Equity: {stock} Rate: {self.portfolioTargetCollect[stock]}")   
        # return difference
        for stock in self.diff(current_stocks, new_stocks):
            self.Log(f"Sale Equity: {stock}") 
        return self.diff(current_stocks, new_stocks), self.diff(new_stocks, current_stocks)
    
    def Rebalance(self):
        self.Log("Rebalance")
        self.current_val = self.Portfolio.TotalPortfolioValue
        
        self.TrendUpAlpha = self.TrendUp()

        if self.TrendUpAlpha:
            self.stocks_to_hold = self.FineSymbols
        else: 
            current_hold = [x.Key for x in self.Portfolio if x.Value.Invested]
            self.stocks_to_hold = [x for x in self.FineSymbols if x in current_hold]
        # Stocks_to_hold    

            
        #calculate weights
        stock_weight = (1.0 / self.TARGET_SECURITIES)
        bond_weight = 0
        if len(self.bonds) >0:
            bond_weight  = max (1.0 - stock_weight* len(self.stocks_to_hold), 0) / len(self.bonds)
        
        self.portfolioTargetCollect = {}
        for x in self.stocks_to_hold:
            #self.Debug(f"stock : {x.Value} {stock_weight}")
            self.portfolioTargetCollect[x] = stock_weight
        if bond_weight > 0.0:
            for x in self.bonds:
                 #self.Debug(f"bond: {x} {bond_weight}")
                 self.portfolioTargetCollect[x] = bond_weight
        
        # # Print Hold Equities
        # self.PrintHoldEquities()
        
        # # Print TargetEquities()
        # self.PrintTargetEquities()
        
        # If holding is the same, skip adjust!!!

        # get equity
        equity = self.Portfolio.TotalPortfolioValue
        # iterate new stocks
        total_adjust_amount = 0.0
        sales, buy = self.diff_stocks()
        for symbol in sales:
            try:
                total_adjust_amount = total_adjust_amount + self.Portfolio[symbol].Quantity * self.Portfolio[symbol].Price
                self.Liquidate(symbol)
                self.Log(f"Sell stock: {symbol}")
            except Exception as e:
                self.Log(f"Sell {symbol} exception: {e}")
        
        for stock in [x for x in self.portfolioTargetCollect.keys()] :
            target_amount = self.portfolioTargetCollect[stock] * equity
            adjust_amount = self.balance(stock, target_amount)
            total_adjust_amount = total_adjust_amount + adjust_amount
        positions = [x.Key for x in self.Portfolio if x.Value.Invested]
        self.Log(f"After rebalance positions: {len(positions)}" )
        self.Log(f"Account total equity value : {self.Portfolio.TotalPortfolioValue}")
        self.Log(f"Adjust balance ratio:{total_adjust_amount/equity}")
        self.Plot("Adjust balance ratio", "Adjust balance ratio", total_adjust_amount/equity)
        #self.Plot("TrendUp", "inout_check", trend_up_val)
        
        self.NeedReSelect = True
        
    def PrintHoldEquities(self):
        self.Log("######## PrintHoldEquities: #########")
        current_hold = [x.Key for x in self.Portfolio if x.Value.Invested]
        for stock in current_hold:
            self.Log(f"Hold Equity: {stock}, Qty: {self.Portfolio[stock].Quantity}")
    
    def PrintTargetEquities(self):
        self.Log("######## PrintTargetEquities: #########")
        for stock in [x for x in self.portfolioTargetCollect.keys()]:
            self.Log(f"Target Equity: {stock} Rate: {self.portfolioTargetCollect[stock]}")