#region imports
from AlgorithmImports import *
from typing import List, Dict
from pandas.core.frame import DataFrame
from pandas.core.series import Series
from io import StringIO
#endregion

def get_ranked_data(algo: QCAlgorithm, diff_threshold: float) -> DataFrame:
    # source: https://polymarket.com/event/presidential-election-winner-2024?tid=1741681510458
    load: str = algo.download(f'data.quantpedia.com/backtesting_data/economic/us_elections_2024.csv')
    df: DataFrame = pd.read_csv(StringIO(load), delimiter=';')
    # df['date'] = pd.to_datetime(df['date']).dt.date
    df['date'] = pd.to_datetime(df['date']) + pd.Timedelta(days=1)
    df['date'] = df['date'].dt.date
    df.set_index('date', inplace=True)

    def assign_ranks(row) -> Series:
        if abs(row['Donald Trump'] - row['Kamala Harris']) > diff_threshold:
            if row['Donald Trump'] > row['Kamala Harris']:
                return pd.Series({'Trump_rank': 1, 'Harris_rank': 0})
            else:
                return pd.Series({'Trump_rank': 0, 'Harris_rank': 1})
        else:
            return pd.Series({'Trump_rank': 0, 'Harris_rank': 0})

    ranks: DataFrame = df.apply(assign_ranks, axis=1)
    return pd.concat([df, ranks], axis=1)

# https://quantpedia.com/strategies/election-driven-arbitrage-strategy/
# 
# The investment universe for this strategy includes stocks, currencies, and digital assets that align with the policy agendas of the 2024 U.S. Presidential candidates, 
# Donald Trump and Kamala Harris. The initial universe consists of over 1,000 assets, which are filtered down to approximately 300 based on thematic baskets tailored to 
# each candidate’s policies. For Trump, the focus is on sectors such as energy, defense, technology, and onshoring. For Harris, the emphasis is on renewable energy, 
# healthcare, and fiscal stimulus beneficiaries. The final selection involves approximately 40 assets for each candidate’s long-only portfolio, further narrowed down to 
# the top 15 best-performing assets.
# The strategy utilizes prediction market data to inform trading decisions, specifically focusing on daily changes in election outcome probabilities. Critical events are 
# defined by daily changes in prediction market odds exceeding a specified threshold (e.g., 4%). For positive events, where odds increase for a candidate, assets that 
# exceed a dynamically calculated positive threshold are selected for long positions. For negative events, where odds decrease, assets that remain above a dynamically 
# calculated negative floor are selected. The methodology involves calculating dynamic thresholds based on Year-to-Date (YTD) performance, allowing for asset-specific 
# sensitivity to market conditions.
# Portfolios are constructed with a focus on long-only positions. Each portfolio consists of the top 15 assets selected based on their alignment with candidate policies 
# and historical performance. Rebalancing occurs periodically to reflect changes in prediction market odds and policy developments.

# region imports
from AlgorithmImports import *
import data_tools
from typing import List
from pandas.core.frame import DataFrame
# endregion

class ElectionDrivenArbitrageStrategy(QCAlgorithm):

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

        self._top_count: int = 5
        diff_threshold: float = .04
        trade_start_date: datetime = datetime(2024, 11, 1)
        trade_end_date: datetime = datetime(2025, 1, 20)

        self._election_ranks: DataFrame = data_tools.get_ranked_data(self, diff_threshold)

        # source: https://etfdb.com/etfs/industry/
        etf_tickers: List[str] = [
            'QQQ', 'XLF', 'VNQ', 'XLV', 'SMH', 'XLE', 'XLI', 'XLY', 'XLP', 'XLU', 
            'XLC', 'GDX', 'KWEB', 'PAVE', 'AMLP', 'IBB', 'ITA', 'IGV', 'XLB', 'CIBR', 
            'GUNR', 'VOX', 'KBWB', 'IHI', 'PHO', 'KRE', 'ITB', 'SKYY', 'FAS', 'URA', 
            'TSLL', 'MLPX', 'ICLN', 'COPX', 'XME', 'SIL', 'GRID', 'EIPI', 'KIE', 'PPH', 
            'BIZD', 'IYC', 'IAI', 'OIH', 'IYK', 'CONL', 'XRT', 'NLR', 'LIT', 'DFGR', 
        ]

        self._universe: List[Symbol] = [
            self.add_equity(etf_ticker, Resolution.DAILY).symbol for etf_ticker in etf_tickers
        ]

        self.settings.daily_precise_end_time = False
        self.settings.minimum_order_margin_portfolio_percentage = 0.

        self._trade_flag: bool = False
        self.schedule.on(
            self.date_rules.on(trade_start_date),
            self.time_rules.before_market_close(self._universe[0]), 
            self._trade
        )

        self.schedule.on(
            self.date_rules.on(trade_end_date),
            self.time_rules.at(0, 0), 
            lambda: self.liquidate()
        )

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

        # Get ETFs performance history.
        history: DataFrame = self.history(self._universe, start=self._election_ranks.index[0], end=self._election_ranks.index[-1]).unstack(level=0)
        if len(history.close.columns) != len(self._universe):
            self.log('Not enough data for further calculation.')
            return
            
        returns: DataFrame = history.close.pct_change().dropna()
        df_combined: DataFrame = pd.merge(self._election_ranks, returns, left_index=True, right_index=True, how='inner')

        trump_portfolio: List[Symbol] = list(df_combined[df_combined['Trump_rank'] == 1].iloc[:, 4:].sum().sort_values(ascending=False)[:self._top_count].index)
        harris_portfolio: List[Symbol] = list(df_combined[df_combined['Harris_rank'] == 1].iloc[:, 4:].sum().sort_values(ascending=False)[:self._top_count].index)

        # Trade execution.
        for i, portfolio in enumerate([trump_portfolio, harris_portfolio]):
            for symbol in portfolio:
                if slice.contains_key(symbol) and slice[symbol]:
                    self.market_order(symbol, self.calculate_order_quantity(symbol, 1 / len(portfolio)))

    def _trade(self) -> None:
        self._trade_flag = True