# region imports
from AlgorithmImports import *
from pandas.core.frame import DataFrame
# endregion

class TacticalAllocation(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2004, 1, 1) # VOO Inception Date Sep 07, 2010
        self.init_cash:float = 10000.
        self.SetCash(self.init_cash)

        # one drive file hash - not used at the moment
        ticker_hash:Dict[str, str] = {
            'MDY' : '1KCeGIxxtIOSPx0ygnmouwv6f_XZyCLVvlCQsUvQkVjU',
            'EFA' : '1QJsulsdxY4zjxXXCzXpTzrqrIk3rFaq6NCwe_Qo6xkg',
            'TLT' : '185hK7chN1e25_phYL47vjuSusnBy4cweUxZhZ3BIDlo',
            'QQQ' : '1iqBsgueshFKwYxxeKf5kSmXen-OWlynjyiPFPsoP_dY',
            'FEMKX' : '1uLMoCAGxzcvEufsy50xeeKL2vdXksKdBb8MfbVfFgUw',
        }

        self.traded_symbols:List[Symbol] = []
       # self.AddRiskManagement(MaximumDrawdownPercentPortfolio(-0.10))
        
        self.ma_period:int = 8
        self.momentum_period:int = 4
        self.SetWarmup(max([self.ma_period, self.momentum_period]), Resolution.Daily)
        self.traded_count:int = 1 # number of assets to hold

        # subscribe data
        for ticker, hash_ in ticker_hash.items():
            data:Security = self.AddData(CustomTicker, ticker, Resolution.Daily)
            data.SetFeeModel(CustomFeeModel())

            self.traded_symbols.append(data.Symbol)
        
        self.recent_month:int = self.Time.month
        
        # benchmark
        self.print_benchmark:bool = False
        self.benchmark:Symbol = self.AddEquity("VOO", Resolution.Daily).Symbol
        self.benchmark_values:List[float] = []

        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.

    def OnEndOfDay(self) -> None:
        if self.print_benchmark:
            mkt_price_df:DataFrame = self.History(self.benchmark, 2, Resolution.Daily)
            if not mkt_price_df.empty:
                mkt_price:float = mkt_price_df['close'].unstack(level= 0).iloc[-1]
                if len(self.benchmark_values) == 2:
                    self.benchmark_values[-1] = mkt_price
                    mkt_perf:float = self.init_cash * self.benchmark_values[-1] / self.benchmark_values[0] 
                    self.Plot('Strategy Equity', self.benchmark, mkt_perf)
                else:
                    self.benchmark_values.append(mkt_price)

    def OnData(self, data: Slice) -> None:
        if self.IsWarmingUp: return

        # rebalance once a month
        if self.Time.month == self.recent_month:
            return
        self.recent_month = self.Time.month

        last_update_date:Dict[str, datetime.date] = CustomTicker.get_last_update_date()

        # custom data stopped comming in
        if all(x.Value in last_update_date and self.Time.date() < last_update_date[x.Value] for x in self.traded_symbols):
            # momentum sort
            traded_symbols:List[Symbol] = []

            m_period:int = max([self.ma_period, self.momentum_period])
            hist_period:int = m_period * 31
            price_df:DataFrame = self.History(self.traded_symbols, hist_period, Resolution.Daily).unstack(level=0)['adj close']
            price_df = price_df.groupby(pd.Grouper(freq='M')).last()[-m_period:]
            if len(price_df) == m_period:
                momentum_df:DataFrame = price_df.pct_change(periods=self.momentum_period)
                ma_df:DataFrame = price_df.rolling(self.ma_period).mean()
                
                top_by_momentum:str = momentum_df.iloc[-1].sort_values(ascending=0).index[0]
                if price_df.iloc[-1][top_by_momentum] >= ma_df.iloc[-1][top_by_momentum]:
                    traded_symbols.append(self.Symbol(top_by_momentum))
        
            # liquidate
            invested:List[Symbol] = [x.Key for x in self.Portfolio if x.Value.Invested]
            for symbol in invested:
                if symbol not in traded_symbols:
                    self.Liquidate(symbol)

            # rebalance
            for symbol in traded_symbols:
                if not self.Portfolio[symbol].Invested:
                    self.SetHoldings(symbol, 1 / len(traded_symbols))
        else:
            self.Liquidate()
            return

# Custom Ticker data
# NOTE Data order must be ascending (datewise)
class CustomTicker(PythonData):
    _last_update_date:Dict[str, datetime.date] = {}

    @staticmethod
    def get_last_update_date() -> Dict[str, datetime.date]:
       return CustomTicker._last_update_date

    def GetSource(self, config: SubscriptionDataConfig, date: datetime, isLiveMode: bool) -> SubscriptionDataSource:
        return SubscriptionDataSource(f"data.quantpedia.com/backtesting_data/equity/{config.Symbol.Value}.csv", SubscriptionTransportMedium.RemoteFile, FileFormat.Csv)

    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, isLiveMode: bool) -> BaseData:
        if not (line.strip() and line[0].isdigit()): return None

        data = CustomTicker()
        data.Symbol = config.Symbol
        
        split:List[str] = line.split(',')
        
        # Date,Open,High,Low,Close,Adj Close,Volume
        data.Time = datetime.strptime(split[0], "%Y-%m-%d") + timedelta(days=1)
        data['Open'] = float(split[1])
        data['High'] = float(split[2])
        data['Low'] = float(split[3])
        data['Close'] = float(split[4])
        data['Adj Close'] = float(split[5])
        data['Volume'] = float(split[6])
        data.Value = float(split[5])

        # store last update date
        if config.Symbol.Value not in CustomTicker._last_update_date:
            CustomTicker._last_update_date[config.Symbol.Value] = datetime(1,1,1).date()

        if data.Time.date() > CustomTicker._last_update_date[config.Symbol.Value]:
            CustomTicker._last_update_date[config.Symbol.Value] = data.Time.date()

        return data

# custom fee model
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))