#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.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"))

european_countries_iso3 = [
    "ALB",  # Albania
    "AND",  # Andorra
    "AUT",  # Austria
    "BLR",  # Belarus
    "BEL",  # Belgium
    "BIH",  # Bosnia and Herzegovina
    "BGR",  # Bulgaria
    "HRV",  # Croatia
    "CYP",  # Cyprus
    "CZE",  # Czech Republic
    "DNK",  # Denmark
    "EST",  # Estonia
    "FIN",  # Finland
    "FRA",  # France
    "DEU",  # Germany
    "GRC",  # Greece
    "HUN",  # Hungary
    "ISL",  # Iceland
    "IRL",  # Ireland
    "ITA",  # Italy
    "XKX",  # Kosovo (not officially ISO recognized but commonly used)
    "LVA",  # Latvia
    "LIE",  # Liechtenstein
    "LTU",  # Lithuania
    "LUX",  # Luxembourg
    "MLT",  # Malta
    "MDA",  # Moldova
    "MCO",  # Monaco
    "MNE",  # Montenegro
    "NLD",  # Netherlands
    "MKD",  # North Macedonia
    "NOR",  # Norway
    "POL",  # Poland
    "PRT",  # Portugal
    "ROU",  # Romania
    "RUS",  # Russia
    "SMR",  # San Marino
    "SRB",  # Serbia
    "SVK",  # Slovakia
    "SVN",  # Slovenia
    "ESP",  # Spain
    "SWE",  # Sweden
    "CHE",  # Switzerland
    "UKR",  # Ukraine
    "GBR",  # United Kingdom
    "VAT",  # Vatican City (Holy See)
]

eu_countries_iso3 = [
    "AUT",  # Austria
    "BEL",  # Belgium
    "BGR",  # Bulgaria
    "HRV",  # Croatia
    "CYP",  # Cyprus
    "CZE",  # Czech Republic
    "DNK",  # Denmark
    "EST",  # Estonia
    "FIN",  # Finland
    "FRA",  # France
    "DEU",  # Germany
    "GRC",  # Greece
    "HUN",  # Hungary
    "IRL",  # Ireland
    "ITA",  # Italy
    "LVA",  # Latvia
    "LTU",  # Lithuania
    "LUX",  # Luxembourg
    "MLT",  # Malta
    "NLD",  # Netherlands
    "POL",  # Poland
    "PRT",  # Portugal
    "ROU",  # Romania
    "SVK",  # Slovakia
    "SVN",  # Slovenia
    "ESP",  # Spain
    "SWE",  # Sweden
]

# region imports
from AlgorithmImports import *
from data_tools import SymbolData, CustomFeeModel, eu_countries_iso3
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

    _minute_resolution: bool = True

    _volume_period: int = 90 # Days
    _min_volume: int = 1_000_000
    _market_cap: int = 1_000_000_000

    def initialize(self) -> None:
        self.set_start_date(2000, 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, SymbolData] = {}
        self._sector_manager: Dict[int, int] = {}
        self._last_selection: List[Symbol] = []

        resolution = Resolution.MINUTE if self._minute_resolution else Resolution.DAILY

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

        self._selection_flag: bool = False
        self._trade_flag: bool = False
        self.universe_settings.resolution = resolution
        self.add_universe(self.fundamental_selection_function)
        self.settings.minimum_order_margin_portfolio_percentage = 0.
        self.settings.daily_precise_end_time = False

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

        if self._minute_resolution:
            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:
        for security in changes.added_securities:
            security.set_fee_model(CustomFeeModel())

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

        # 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.security_reference.exchange_id in self._exchange_codes
            and x.company_reference.country_id in eu_countries_iso3
            and x.market_cap > self._market_cap
        ]

        filtered_by_volume: list[Fundamental] = []
        for security in selected:
            history: DataFrame = self.history(security.symbol, start=self.time.date() - timedelta(days=self._volume_period), end=self.time.date())
            symbol_volume = history.loc[security.symbol]["volume"]
            total_volume = symbol_volume.sum()
            if total_volume > self._min_volume:
                filtered_by_volume.append(security)

        if len(filtered_by_volume) > self._fundamental_count:
            filtered_by_volume = [x for x in sorted(filtered_by_volume, key=self._fundamental_sorting_key, reverse=True)[:self._fundamental_count]]

        # price warmup
        for stock in filtered_by_volume:
            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] = 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())

                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 filtered_by_volume if self._data[x.symbol].is_ready()]

        if not self._minute_resolution:
            self._trade_flag = True

        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 if self._data[symbol].is_ready()}

        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

        # trade execution
        targets: List[PortfolioTarget] = []
        for symbol in long:
            if slice.contains_key(symbol) and slice[symbol]:
                if slice[symbol].value != 0:
                    targets.append(PortfolioTarget(symbol, 1 / len(long)))
        
        self.set_holdings(targets, True)
        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