#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, Set
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)


        # PARAMS:
        self.rebalance_period: RebalancePeriod = RebalancePeriod.BIWEEKLY

        # True -> 12 month rolling performance
        # False -> YTD performance
        self.rolling_perf_flag: bool = False

        # True -> Inversed volatility weighting
        # False -> Equally weighting
        self.inversed_volatily_weighting_flag: bool = False

        # pick max sector percentage of total traded stock count
        self.sector_max_percentage: float = 0.20

        # universe count by market cap
        self.fundamental_count: int = 500

        # alpha chart
        self.calculate_alpha: bool = False


        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.top_count: int = 20
        self.month_period: int = 11
        self.week_period: int = 51
        self.rolling_period: int = 12
        self.inverse_volatility_period: int = 3 * 21

        self.fundamental_sorting_key = lambda x: x.market_cap

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

        years_of_interest: List[int] = [2008, 2009, 2021]
        self.traded_stocks_by_year: Dict[int, Set] = {
            i: set() for i in years_of_interest
        }

        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 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(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[stock] = self.data[symbol].get_performance()

        # sort and divide
        if len(stock_performance) >= self.top_count:
            sorted_stocks: List[Fundamental] = sorted(stock_performance, key=stock_performance.get, reverse=True)
            for stock in sorted_stocks:
                if len(self.long) >= self.top_count:
                    break
                if self.sector_manager[stock.asset_classification.morningstar_sector_code] == 0:
                    continue
                self.long.append(stock.symbol)
                self.sector_manager[stock.asset_classification.morningstar_sector_code] -= 1

        if self.time.year in self.traded_stocks_by_year:
            for symbol in self.long:
                self.traded_stocks_by_year[self.time.year].add(symbol.value)

        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)))
        
        if self.calculate_alpha:
            if len(self.long) != 0:
                targets.append(PortfolioTarget(self.market, -1.))
            else:
                targets.append(PortfolioTarget(self.market, 0.))
        
        self.SetHoldings(targets, True)
        self.long.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.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()

    def on_end_of_algorithm(self) -> None:
        for year, stocks in self.traded_stocks_by_year.items():
            self.log(f'stocks traded in {year}: {stocks}')