# region imports
from AlgorithmImports import *
import statsmodels.api as sm
# endregion

def multiple_linear_regression(x:np.ndarray, y:np.ndarray):
    x = sm.add_constant(x, has_constant='add')
    result = sm.OLS(endog=y, exog=x).fit()
    return result

class SymbolData():
    def __init__(self, bond_yield_symbol: Symbol, period: int) -> None:
        self._bond_yield_symbol: Symbol = bond_yield_symbol
        self._daily_price: RollingWindow = RollingWindow[float](period)
        self._bond_yield_values: RollingWindow = RollingWindow[float](period)
    
    def update_data(self, price: float, bond_yield_value: float) -> None:
        self._daily_price.add(price)
        self._bond_yield_values.add(bond_yield_value)

    def get_returns(self) -> List[float]:
        log_returns: np.ndarray = pd.Series(list(self._daily_price)[::-1]).pct_change().dropna()
        return log_returns

    def get_bond_diff(self) -> List[float]:
        return pd.Series(list(self._bond_yield_values)[::-1]).diff().dropna()

    def get_daily_price(self) -> List[float]:
        return list(self._daily_price)[::-1]

    def is_ready(self) -> bool:
        return self._daily_price.is_ready and self._bond_yield_values.is_ready

class LastDateHandler():
    _last_update_date: Dict[Symbol, datetime.date] = {}

    @staticmethod
    def get_last_update_date() -> Dict[Symbol, datetime.date]:
       return LastDateHandler._last_update_date

# Quantpedia data.
# NOTE: IMPORTANT: Data order must be ascending (datewise)
class QuantpediaFutures(PythonData):
    def GetSource(self, config: SubscriptionDataConfig, date: datetime, isLiveMode: bool) -> SubscriptionDataSource:
        return SubscriptionDataSource("data.quantpedia.com/backtesting_data/futures/{0}.csv".format(config.Symbol.Value), SubscriptionTransportMedium.RemoteFile, FileFormat.Csv)

    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, isLiveMode: bool) -> BaseData:
        data = QuantpediaFutures()
        data.Symbol = config.Symbol
        
        if not line[0].isdigit(): return None
        split = line.split(';')
        
        data.Time = datetime.strptime(split[0], "%d.%m.%Y") + timedelta(days=1)
        data['back_adjusted'] = float(split[1])
        data['spliced'] = float(split[2])
        data.Value = float(split[1])

        if config.Symbol not in LastDateHandler._last_update_date:
            LastDateHandler._last_update_date[config.Symbol] = datetime(1,1,1).date()
        if data.Time.date() > LastDateHandler._last_update_date[config.Symbol]:
            LastDateHandler._last_update_date[config.Symbol] = data.Time.date()

        return data

# Quantpedia bond yield data.
# NOTE: IMPORTANT: Data order must be ascending (datewise)
class QuantpediaBondYield(PythonData):
    def GetSource(self, config: SubscriptionDataConfig, date: datetime, isLiveMode: bool) -> SubscriptionDataSource:
        return SubscriptionDataSource("data.quantpedia.com/backtesting_data/bond_yield/{0}.csv".format(config.Symbol.Value), SubscriptionTransportMedium.RemoteFile, FileFormat.Csv)

    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, isLiveMode: bool) -> BaseData:
        data = QuantpediaBondYield()
        data.Symbol = config.Symbol
        
        if not line[0].isdigit(): return None
        split = line.split(',')
        
        data.Time = datetime.strptime(split[0], "%Y-%m-%d") + timedelta(days=1)
        data['yield'] = float(split[1])
        data.Value = float(split[1])

        if config.Symbol not in LastDateHandler._last_update_date:
            LastDateHandler._last_update_date[config.Symbol] = datetime(1,1,1).date()
        if data.Time.date() > LastDateHandler._last_update_date[config.Symbol]:
            LastDateHandler._last_update_date[config.Symbol] = data.Time.date()

        return data

class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters: OrderFeeParameters) -> OrderFee:
        fee: float = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))

# https://quantpedia.com/strategies/stock-bond-correlations-and-the-expected-country-stock-returns/
# 
# The investment universe for this strategy includes equity and bond markets from a broad set of countries. The selection of individual instruments is based on the 
# availability of data for stock returns and 10-year bond yields. (The research paper mentions 30 countries, including those from Europe, Asia, the Americas, 
# Oceania, and Africa.)
# (Financial data is sourced from Bloomberg. The MSCI index is used to proxy stock returns,
# and 10-year bond yields compute the stock-bond relationship. Use MSCI net total returns data. The survey data is sourced from the International Monetary Fund (IMF). 
# Further macroeconomic data is obtained from the World Bank.)
# Rationale: The primary tool used in the research paper is the SB beta, which measures the relationship between stock returns and changes in bond yields. The 
# methodology involves calculating the SB beta using a regression of stock returns on the negative change in bond yields.
# Calculation: Do regression (1) where the dependent variable is the log return of a country i’s equity index denominated in local currency, and the SB beta is the 
# slope of the regression; independent variable first-difference of the ten-year Treasury bond yield of country i also in local currency. (Similarly, the SB 
# correlation is calculated as the negative correlation between stock returns and first-order differences in bond yields.)
# Ranking: The strategy ranks countries based on their SB beta values, focusing on countries exhibiting a positive SB correlation: The portfolios are formed after 
# sorting countries based on their lagged SB betas, estimated using daily over a rolling 12-month or 52-week window.
# Execution: Based on the sorting, perform:
# Portfolio 5 consists of countries with the highest SB beta: The buy rule involves investing in countries within the top quintile of SB beta values, indicating a 
# positive correlation.
# Portfolio 1 consists of countries with the lowest SB beta: The sell rule involves shorting countries within the bottom quintile of SB beta values, indicating a 
# negative correlation.
# Positions are equally weighted to ensure diversification and mitigate country-specific risks. Rebalancing occurs monthly to adjust for changes in SB beta values.
# 
# QC Implementation:
#   - QC ETFs are used as trading universe. 

# region imports
from AlgorithmImports import *
import data_tools
from pandas.core.frame import DataFrame
import statsmodels.api as sm
# endregion

class TradedStrategy(Enum):
    QP_FUTURES = 1
    ETF = 2

class StockBondCorrelationsAndTheExpectedCountryStockReturns(QCAlgorithm):

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

        self._period: int = 252
        self._quantile: int = 5
        self._data: Dict[Symbol, SymbolData] = {}
        leverage: int = 5

        self._traded_strategy: TradedStrategy = TradedStrategy.QP_FUTURES

        if self._traded_strategy == TradedStrategy.QP_FUTURES:
            tickers_bonds: List[str] = {               
                "CME_ES1"     : "US10YT", # E-mini S&P 500 Futures, Continuous Contract #1
                "EUREX_FDAX1" : "DE10YT", # DAX Futures, Continuous Contract #1
                "EUREX_FSMI1" : "CH10YT", # SMI Futures, Continuous Contract #1
                # "LIFFE_FCE1"  : "FR10YT", # CAC40 Index Futures, Continuous Contract #1 # data ended 2022
                "LIFFE_Z1"    : "GB10YT", # FTSE 100 Index Futures, Continuous Contract #1       
                "SGX_NK1"     : "JP10YT"  # SGX Nikkei 225 Index Futures, Continuous Contract #1      
            }

        else:
            tickers_bonds: List[str] = {
                "SPY" : "US10YT",  # SPDR S&P 500 ETF
                "EWA" : "AU10YT",  # iShares MSCI Australia Index ETF
                "EWO" : "AS10YT",  # iShares MSCI Austria Investable Mkt Index ETF
                "EWK" : "BE10YT",  # iShares MSCI Belgium Investable Market Index ETF
                "EWZ" : "BR10YT",  # iShares MSCI Brazil Index ETF
                "EWC" : "CA10YT",  # iShares MSCI Canada Index ETF
                "FXI" : "CN10YT",  # iShares China Large-Cap ETF
                "EWQ" : "FR10YT",  # iShares MSCI France Index ETF
                "EWG" : "DE10YT",  # iShares MSCI Germany ETF 
                "EWH" : "HK10YT",  # iShares MSCI Hong Kong Index ETF
                "EWI" : "IT10YT",  # iShares MSCI Italy Index ETF
                "EWJ" : "JP10YT",  # iShares MSCI Japan Index ETF
                "EWM" : "MY10YT",  # iShares MSCI Malaysia Index ETF
                "EWW" : "MX10YT",  # iShares MSCI Mexico Inv. Mt. Idx
                "EWN" : "NL10YT",  # iShares MSCI Netherlands Index ETF
                "EWS" : "SG10YT",  # iShares MSCI Singapore Index ETF
                "EZA" : "ZA10YT",  # iShares MSCI South Africa Index ETF
                "EWY" : "KR10YT",  # iShares MSCI South Korea ETF
                "EWP" : "ES10YT",  # iShares MSCI Spain Index ETF
                "EWL" : "CH10YT",  # iShares MSCI Switzerland Index ETF
                "EWU" : "GB10YT",  # iShares MSCI United Kingdom Index ETF
                "INDA": "IN10YT",  # iShares MSCI India ETF
                "TUR" : "TR10YT",  # iShares MSCI Turkey ETF
            }

        for ticker, bond_yield in tickers_bonds.items():
            data: Security = self.add_data(data_tools.QuantpediaFutures, ticker, Resolution.DAILY) \
                if self._traded_strategy == TradedStrategy.QP_FUTURES \
                else self.add_equity(ticker, Resolution.DAILY)
            if self._traded_strategy == TradedStrategy.QP_FUTURES:
                data.set_fee_model(data_tools.CustomFeeModel())
            data.set_leverage(leverage)

            self._data[data.symbol] = data_tools.SymbolData(self.add_data(data_tools.QuantpediaBondYield, bond_yield, Resolution.DAILY).symbol, self._period)

        self.set_warm_up(timedelta(days=self._period), Resolution.DAILY)

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

        self.selection_flag: bool = False
        self.schedule.on(self.date_rules.month_start(next(iter(self._data))),
            self.time_rules.at(0, 0),
            self.selection)

    def on_data(self, slice: Slice) -> None:
        for symbol, symbol_data in self._data.items():
            custom_data: List[Symbol] = [symbol, symbol_data._bond_yield_symbol] \
                if self._traded_strategy == TradedStrategy.QP_FUTURES \
                else [symbol_data._bond_yield_symbol]
            if any(self.securities[x].get_last_data() and self.time.date() > data_tools.LastDateHandler.get_last_update_date()[x] for x in custom_data):
                self.liquidate()
                self.log('Data stopped coming.')
                return

            if self.securities[symbol_data._bond_yield_symbol].get_last_data() and slice.contains_key(symbol) and slice[symbol]:
                symbol_data.update_data(slice[symbol].price, self.securities[symbol_data._bond_yield_symbol].get_last_data().price)

        if self.is_warming_up:
            return

        if not self.selection_flag:
            return
        self.selection_flag = False

        returns_dict: Dict[Symbol, List[float]] = {
            symbol: data.get_returns() 
            for symbol, data in self._data.items() 
            if data.is_ready()
        }
        funds_returns: List[List[float]] = list(zip(*[[i for i in x] for x in returns_dict.values()]))
        bond_diff: List[List[float]] = list(zip(
            *[[i for i in x] 
            for x in [data.get_bond_diff() 
            for data in self._data.values() 
            if data.is_ready()]]
        ))

        if len(returns_dict) < self._quantile:
            self.log('Not enough data for further calculation.')
            return

        # get beta
        x: np.ndarray = np.array(bond_diff)
        y: np.ndarray = np.array(funds_returns)
        model = data_tools.multiple_linear_regression(x, y)
        betas: np.ndarray = model.params[1]
            
        # store betas
        beta_values: Dict[Symbol, float] = {symbol : betas[i] for i, symbol in enumerate(list(returns_dict))}

        long: List[Symbol] = []
        short: List[Symbol] = []

        # sort and divide
        sorted_betas: List[Symbol] = sorted(beta_values, key=beta_values.get, reverse=True)
        quantile: int = int(len(sorted_betas) / self._quantile)
        long = sorted_betas[:quantile]
        short = sorted_betas[-quantile:]

        # trade execution
        targets: List[PortfolioTarget] = []
        for i, portfolio in enumerate([long, short]):
            for symbol in portfolio:
                if slice.contains_key(symbol) and slice[symbol]:
                    targets.append(PortfolioTarget(symbol, ((-1) ** i) / len(portfolio)))
        
        self.set_holdings(targets, True)

    def selection(self) -> None:
        self.selection_flag = True