#region imports
from AlgorithmImports import *
#endregion

# Your New Python File
class SymbolData():
    def __init__(self, period: int, volatility_period: int, rolling_flag: bool) -> None:
        self._rolling_flag: bool = rolling_flag
        self._period: int = period
        self._daily_price: RollingWindow = RollingWindow[float](volatility_period)
        if self._rolling_flag:
            self._price: RollingWindow = RollingWindow[float](period)
        else:
            self._price: List[float] = []

    def update_daily_price(self, price: float) -> None:
        self._daily_price.add(price)
        
    def update_price(self, price: float) -> None:
        if self._rolling_flag:
            self._price.add(price)
        else:
            self._price.insert(0, price)

    def get_daily_returns(self) -> np.ndarray:
        daily_returns: np.ndarray = pd.Series(list(self._daily_price)[::-1]).pct_change().values[1:]
        return daily_returns

    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:
        if self._rolling_flag:
            return self._price.is_ready and self._daily_price.is_ready
        else:
            return len(self._price) > 1 and self._daily_price.is_ready

    def get_inverse_vol(self) -> float:
        daily_returns: np.ndarray = self.get_daily_returns()
        return 1 / (np.std(daily_returns) * np.sqrt(3*21))

# 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
from pandas.core.frame import DataFrame
from pandas.core.series import Series
# endregion

class RebalancePeriod(Enum):
    MONTHLY = 1
    BIWEEKLY = 2
    WEEKLY = 3

class Metatron005(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2000, 1, 1)
        self.set_cash(100_000)

        self.rebalance_period: RebalancePeriod = RebalancePeriod.MONTHLY

        self.universe_settings.leverage = 10

        # True -> 12 month rolling performance
        # False -> performance from expanding list, reset annually
        self.rolling_perf_flag: bool = False
        # True -> Use inversed volatility weighting
        # False -> Equally weighting
        self.inversed_volatily_weighting_flag: bool = True

        self.tickers_to_ignore: List[str] = ['EGAN', 'CR', 'PATH', 'ILIF', 'CRW']
        self.exchange_codes: List[str] = ['NYS', 'NAS', 'ASE']    

        self.top_count: int = 15
        self.month_period: int = 11
        self.week_period: int = 51
        self.rolling_period: int = 12
        self.inverse_volatility_period: int = 3 * 21

        self.fundamental_count: int = 500
        self.fundamental_sorting_key = lambda x: x.market_cap

        self.data: Dict[Symbol, data_tools.SymbolData] = {}
        self.long: List[Symbol] = []

        self.market: Symbol = self.AddEquity('SPY', Resolution.Daily).Symbol

        self.selection_flag: bool = False
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self.fundamental_selection_function)
        self.settings.minimum_order_margin_portfolio_percentage = 0.
        self.settings.daily_precise_end_time = False

        if self.rebalance_period == RebalancePeriod.MONTHLY:
            self.schedule.on(self.date_rules.month_start(self.market),
                            self.time_rules.after_market_open(self.market),
                            self.selection)
        else:
            self.schedule.on(self.date_rules.every_day(self.market),
                            self.time_rules.after_market_open(self.market),
                            self.selection)
            self.counter: int = 0
            self.current_week: int = -1
        self.curr_year: int = -1

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

    def fundamental_selection_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # update data
        for stock in fundamental:
            symbol: Symbol = stock.symbol
            
            if symbol in self.data:
                self.data[symbol].update_daily_price(stock.adjusted_price)
                if self.selection_flag:
                    self.data[symbol].update_price(stock.adjusted_price)

        if not self.selection_flag:
            return Universe.UNCHANGED

        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]]

        stock_performance: Dict[Symbol, float] = {}

        # price warmup
        for stock in selected:
            symbol: Symbol = stock.symbol

            if symbol not in self.data:
                self.data[symbol] = data_tools.SymbolData(self.rolling_period, self.inverse_volatility_period, self.rolling_perf_flag)

                if self.rolling_perf_flag:
                    history: DataFrame = self.history(symbol, start=self.time.date() - relativedelta(months=self.rolling_period), end=self.time.date())
                else:
                    if self.rebalance_period == RebalancePeriod.MONTHLY:
                        lookback: int = self.time.month - 1 if self.time.month != 1 else self.time.month + self.month_period
                        history: DataFrame = self.history(symbol, start=self.time.date() - relativedelta(months=lookback), end=self.time.date())
                    else:
                        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())

                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='MS')).first() if self.rebalance_period == RebalancePeriod.MONTHLY else data.groupby(pd.Grouper(freq='W')).first()
                for time, row in data_periodically.iterrows():
                    self.data[symbol].update_price(row.close)
                for time, row in data.iterrows():
                    self.data[symbol].update_daily_price(row.close)
                
            if self.data[symbol].is_ready():
                stock_performance[symbol] = self.data[symbol].get_performance()

        # sort and divide
        if len(stock_performance) >= self.top_count:
            sorted_stocks: List[Symbol] = sorted(stock_performance, key=stock_performance.get, reverse=True)
            self.long = sorted_stocks[:self.top_count]

        return self.long

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

        if self.inversed_volatily_weighting_flag:
            sum_ivol: float = sum([self.data[symbol].get_inverse_vol() for symbol in self.long])

        # trade execution
        targets: List[PortfolioTarget] = []
        for symbol in self.long:
            if slice.contains_key(symbol) and slice[symbol]:
                if slice[symbol].value != 0:
                    if self.inversed_volatily_weighting_flag:
                        if sum_ivol != 0:
                            targets.append(PortfolioTarget(symbol, self.data[symbol].get_inverse_vol() / sum_ivol))
                    else:
                        targets.append(PortfolioTarget(symbol, 1 / len(self.long)))
        
        targets.append(PortfolioTarget(self.market, -1.))
        
        self.SetHoldings(targets, True)
        self.long.clear()

    def selection(self) -> None:
        if self.rebalance_period == RebalancePeriod.MONTHLY:
                self.selection_flag = True
        else:
            if self.time.isocalendar().week == self.current_week:
                return
            self.current_week = self.time.isocalendar().week
            if self.rebalance_period == RebalancePeriod.WEEKLY:
                self.selection_flag = True
            if self.rebalance_period == RebalancePeriod.BIWEEKLY:
                self.counter += 1
                if self.counter == 2:
                    self.selection_flag = True
                    self.counter = 0
                    
        if not self.rolling_perf_flag:
            # 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()