#region imports
from AlgorithmImports import *
#endregion


# Your New Python File
class SymbolData():
    def __init__(self) -> None:
        self._price: List[float] = []

    def update_price(self, price: float) -> None:
        self._price.append(price)

    def get_performance(self) -> float:
        return self._price[-1] / self._price[0] - 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
# endregion

class Metatron005(QCAlgorithm):

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

        # True -> monthly rebalance
        # False -> rebalance every 2 weeks
        self.month_rebalance_flag: bool = False

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

        self.top_count: int = 25
        self.month_period: int = 11
        self.week_period: int = 51

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

        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.
        if self.month_rebalance_flag:
            self.schedule.on(self.date_rules.month_start(market),
                            self.time_rules.after_market_open(market),
                            self.selection)
        else:
            self.schedule.on(self.date_rules.every_day(market),
                            self.time_rules.after_market_open(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]:
        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.company_reference.short_name is not None
        #     and 'crane' in x.company_reference.short_name.lower()
        # ]
        # if len(selected) != 0: 
        #     foo = 3

        selected: List[Fundamental] = [
            x for x in fundamental 
            if x.SecurityReference.ExchangeId in self.exchange_codes
            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()
                if self.month_rebalance_flag:
                    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: Series = data.groupby(pd.Grouper(freq='MS')).first() if self.month_rebalance_flag else data.groupby(pd.Grouper(freq='W')).first()
                for time, row in data.iterrows():
                    self.data[symbol].update_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

        # trade execution
        targets: List[PortfolioTarget] = []
        for symbol in self.long:
            if slice.contains_key(symbol) and slice[symbol]:
                targets.append(PortfolioTarget(symbol, 1 / len(self.long)))
        
        self.SetHoldings(targets, True)
        self.long.clear()

    def selection(self) -> None:
        if self.month_rebalance_flag:
                self.selection_flag = True
        else:
            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()