#region imports
from AlgorithmImports import *
#endregion

class SymbolData():
    def __init__(self) -> None:
        self._price: List[float] = []
        
    def update_price(self, price: float) -> None:
        self._price.insert(0, price)

    def get_performance(self) -> float:
        return self._price[0] / list(self._price)[-1] - 1

    def reset(self) -> None:
        self._price = self._price[:1]

    def is_ready(self) -> bool:
        return len(self._price) > 1

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

# region imports
from AlgorithmImports import *
import data_tools
from dateutil.relativedelta import relativedelta
from typing import List, Dict, Set
from pandas.core.frame import DataFrame
from pandas.core.series import Series
# endregion

class MetatronYTDHigh(QCAlgorithm):

    _notional_value: int = 60_000
    _trade_exec_minute_offset: int = 15

    # pick max sector percentage of total traded stock count
    _sector_max_percentage: float = 0.20

    _top_count: int = 10
    _week_period: int = 51

    def initialize(self) -> None:
        self.set_start_date(2022, 1, 1)
        self.set_cash(self._notional_value)

        # universe count by market cap
        self._fundamental_count: int = 500
        self._fundamental_sorting_key = lambda x: x.market_cap

        self.universe_settings.leverage = 5
        self._tickers_to_ignore: List[str] = ['EGAN', 'CR', 'PATH', 'ILIF', 'CRW']
        self._exchange_codes: List[str] = ['NYS', 'NAS', 'ASE']    

        self._data: Dict[Symbol, data_tools.SymbolData] = {}
        self._sector_manager: Dict[int, int] = {}
        self._last_selection: List[Symbol] = []

        market: Symbol = self.add_equity('SPY', Resolution.MINUTE).Symbol

        self._selection_flag: bool = True # Selection and trading is executed on start of backtest.
        self._trade_flag: bool = False
        self.universe_settings.resolution = Resolution.MINUTE
        self.add_universe(self.fundamental_selection_function)
        self.settings.minimum_order_margin_portfolio_percentage = 0.

        self.schedule.on(
            self.date_rules.every_day(market),
            self.time_rules.before_market_close(market),
            self._selection
        )

        self.schedule.on(
            self.date_rules.every_day(market),
            self.time_rules.before_market_close(market, self._trade_exec_minute_offset),
            self._rebalance
        )

        self._counter: int = 0
        self._current_week: int = -1
        self._curr_year: int = -1

    def on_securities_changed(self, changes: SecurityChanges) -> None:
        if not self.live_mode:
            for security in changes.added_securities:
                security.set_fee_model(data_tools.CustomFeeModel())
        else:
            pass

    def fundamental_selection_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        if not self._selection_flag:
            return Universe.UNCHANGED

        # update data
        for stock in fundamental:
            symbol: Symbol = stock.symbol
            
            if symbol in self._data:
                self._data[symbol].update_price(stock.adjusted_price)

        selected: List[Fundamental] = [
            x for x in fundamental 
            if x.SecurityReference.ExchangeId in self._exchange_codes
            and x.Market == 'usa'
            and x.symbol.value not in self._tickers_to_ignore
        ]
        if len(selected) > self._fundamental_count:
            selected = [
                x for x in sorted(selected, key=self._fundamental_sorting_key, reverse=True)[:self._fundamental_count]
            ]

        self.log(f'New selection of {len(selected)} assets')

        # price warmup
        for stock in selected:
            symbol: Symbol = stock.symbol
            
            if stock.asset_classification.morningstar_sector_code not in self._sector_manager:
                self._sector_manager[stock.asset_classification.morningstar_sector_code] = self._top_count * self._sector_max_percentage

            if symbol not in self._data:
                self._data[symbol] = data_tools.SymbolData()

                lookback: int = self.time.isocalendar().week - 1 if self.time.isocalendar().week != 1 else self.time.isocalendar().week + self._week_period
                # history: DataFrame = self.history(symbol, start=self.time.date() - relativedelta(weeks=lookback), end=self.time.date())
                history: DataFrame = self.history(symbol, self.time.date() - relativedelta(weeks=lookback), self.time.date(), Resolution.DAILY)

                if history.empty:
                    self.log(f"Not enough data for {symbol} yet.")
                    continue
                data: Series = history.loc[symbol]
                data_periodically: Series = data.groupby(pd.Grouper(freq='W')).first()
                for time, row in data_periodically.iterrows():
                    self._data[symbol].update_price(row.close)

        self._last_selection = [x.symbol for x in selected if self._data[x.symbol].is_ready()]

        return self._last_selection

    def on_data(self, slice: Slice) -> None:
        # order execution
        if not self._trade_flag:
            return
        self._trade_flag = False

        stock_performance: Dict[Symbol, float] = {
            symbol: self._data[symbol].get_performance() for symbol in self._last_selection
        }

        long: List[Symbol] = []

        # sort and divide
        if len(stock_performance) >= self._top_count:
            sorted_stocks: List[Symbol] = sorted(stock_performance, key=stock_performance.get, reverse=True)
            for symbol in sorted_stocks:
                sector_code: str = self.securities[symbol].fundamentals.asset_classification.morningstar_sector_code
                if sector_code == 0:
                    self.log(f'Sector code missing for ticker: {symbol.value}.')
                    continue
                if len(long) >= self._top_count:
                    break
                if self._sector_manager[sector_code] == 0:
                    continue
                long.append(symbol)
                self._sector_manager[sector_code] -= 1
        else:
            self.log(f'Not enough stocks for further selection: {len(stock_performance)}')

        invested: List[Symbol] = [x.key for x in self.portfolio if x.value.invested]
        for symbol in invested:
            if symbol not in long:
                self.liquidate(symbol)

        # trade execution
        self.log(f'Rebalancing portfolio with {len(long)} assets')

        for symbol in long:
            if self.securities[symbol].is_tradable:
                if self.securities[symbol].price != 0:
                    q: int = (self._notional_value / len(long)) // self.securities[symbol].price
                    self.market_order(
                        symbol, 
                        q - self.portfolio[symbol].quantity,
                    )
                else:
                    self.log(f'{symbol} price is {self.securities[symbol].price}')
            else:
                self.log(f'{symbol} not present in current slice. Asset is tradable: {self.securities[symbol].is_tradable}; Last time available: {self.securities[symbol].get_last_data().time}')
        self._last_selection.clear()

        for sector, _ in self._sector_manager.items():
            self._sector_manager[sector] = self._top_count * self._sector_max_percentage

    def _selection(self) -> None:
        if self.time.isocalendar().week == self._current_week:
            return
        self._current_week = self.time.isocalendar().week
        self._counter += 1
        if self._counter == 2:
            self._selection_flag = True
            self._counter = 0

        # reset price data
        if self.time.year != self._curr_year:
            self._curr_year = self.time.year
            for symbol, symbol_data in self._data.items():
                symbol_data.reset()

    def _rebalance(self) -> None:
        if self._selection_flag:
            self._trade_flag = True
            self._selection_flag = False