#region imports
from AlgorithmImports import *
import bz2
import pickle
import base64
import numpy as np
from typing import List, Dict, OrderedDict
#endregion

def initialize_QP_custom_data(algo: QCAlgorithm, period: int, leverage: int, tickers_to_ignore: List[str]) -> Dict:
    QP_data: Dict[Symbol, SymbolData] = {}
    ticker_file_str: str = algo.download('data.quantpedia.com/backtesting_data/equity/india_stocks/nse_500_tickers.csv')
    ticker_lines: List[str] = ticker_file_str.split('\r\n')
    tickers = [ ticker_line.split(',')[0] for ticker_line in ticker_lines[1:] ]

    for t in tickers:
        if t in tickers_to_ignore:
            continue
        # price data subscription
        data: Security = algo.add_data(IndiaStocks, t, Resolution.DAILY)
        data.set_fee_model(CustomFeeModel())
        data.set_leverage(leverage)
        stock_symbol: Symbol = data.symbol

        QP_data[stock_symbol] = SymbolData(period)

    return QP_data

class SymbolData():
    def __init__(
        self, 
        period: int
    ) -> None:
        self._period: int = period
        self._prices: RollingWindow = RollingWindow[float](period)

    def update_price(self, price: float) -> None:
        self._prices.add(price)
    
    def is_ready(self) -> bool:
        return self._prices.is_ready

    def get_momentum(self) -> float:
        return self._prices[1] / self._prices[self._prices.count - 1] - 1

    def get_volatility(self) -> float:
        returns: np.ndarray = pd.Series(list(self._prices)[::-1]).pct_change().values[1:]
        return np.std(returns) * np.sqrt(12)

# India stock price/volume data
# NOTE: IMPORTANT: Data order must be ascending (datewise)
class IndiaStocks(PythonData):
    def GetSource(self, config:SubscriptionDataConfig, date:datetime, isLiveMode:bool) -> SubscriptionDataSource:
        return SubscriptionDataSource("data.quantpedia.com/backtesting_data/equity/india_stocks/nse_500_prices.dat", SubscriptionTransportMedium.REMOTE_FILE, FileFormat.UNFOLDING_COLLECTION)
    
    _last_update_date:Dict[Symbol, datetime.date] = {}

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

    def Reader(self, config:SubscriptionDataConfig, line:str, date:datetime, isLiveMode:bool) -> BaseDataCollection:
        objects:list[IndiaStocks] = []

        base64_bytes = line.encode('ascii')
        data_to_decompress = base64.b64decode(base64_bytes)
        decompressed_data = bz2.decompress(data_to_decompress)
        data:list[dict] = pickle.loads(decompressed_data)

        for index, sample in enumerate(data):
            custom_data:IndiaStocks = IndiaStocks()
            custom_data.symbol = config.symbol
            
            curr_date: datetime = datetime.strptime(sample['date'], '%Y-%m-%d') + timedelta(days=1)
            custom_data.Time = curr_date

            if config.symbol.value in sample['stocks']:
                price: float = float(sample['stocks'][config.symbol.value]['c']) if sample['stocks'][config.symbol.value]['c'] != 'null' else 0.
                custom_data['price'] = price
                custom_data.Value = price

                # store last date of the symbol
                if config.symbol not in IndiaStocks._last_update_date:
                    IndiaStocks._last_update_date[config.symbol] = datetime(1,1,1).date()
                if custom_data.Time.date() > IndiaStocks._last_update_date[config.symbol]:
                    IndiaStocks._last_update_date[config.symbol] = custom_data.Time.date()
            else:
                custom_data['price'] = 0
                custom_data.Value = 0

            objects.append(custom_data)

        return BaseDataCollection(objects[-1].EndTime, config.symbol, objects)

# Custom fee model.
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))

# https://quantpedia.com/strategies/adjusted-momentum-strategies-in-indian-stocks/
# 
# This strategy’s investment universe focuses on a broad selection of Indian equities (NSE or BSE exchanges), such as those in the Nifty 200 index.
# (Data is from LSEG Datastream.)
# Calculation: WML_Vol_Adj scales the returns used in MOM by the volatility of each stock’s return. Specifically, the returns over the 12-month ranking period are 
# divided by the standard deviation of the returns during the same period: Volatility-adjusted momentum is calculated by taking each stock’s 12-month return, 
# excluding the most recent month, and adjusting it by the standard deviation of returns over the same period, annualized by multiplying by the square root of 12.
# Ranking: Stocks are ranked according to WML_Vol_Adj, volatility-adjusted momentum scores, with the top decile identified as “winners” and the bottom decile as “losers.”
# Execution: The buy rule involves going long on stocks in the top decile of momentum scores, while the sell rule involves shorting stocks in the bottom decile.
# Rebalancing & Weighting: Use (market) value-weighted decile portfolios. Rebalance monthly.
# 
# Implementation changes:
#   - QP NSE 100 custom data are used as trading universe (data available until 2023).
#   - Due to the absence of market cap data, an equally weighted decile portfolios are used.

# region imports
from AlgorithmImports import *
from typing import List, Dict
import data_tools
# endregion

class AdjustedMomentumStrategiesInIndianStocks(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2010, 1, 1)
        self.set_cash(10_000_000) # INR

        excluded_tickers: List[str] = ['TATAMTRDVR', 'LODHA']
        self._excluded_tickers: List[str] = ['RELIANCE', 'TCS', 'HDFCBANK', 'INFY']

        leverage: int = 3
        period: int = 12
        self._quantile: int = 10

        self._data: Dict[Symbol, data_tools.SymbolData] = data_tools.initialize_QP_custom_data(self, period, leverage, excluded_tickers)

        self.settings.minimum_order_margin_portfolio_percentage = 0.
        self.settings.daily_precise_end_time = False
        self._current_month: int = -1

    def on_data(self, slice: Slice) -> None:
        price_last_update_date: Dict[Symbol, datetime.date] = data_tools.IndiaStocks.get_last_update_date()
        
        # check if custom data is still comming in
        if all(
            [
                self.securities[x].get_last_data() for x in list(self._data.keys())]) and any([self.time.date() >= price_last_update_date[x] \
                for x in price_last_update_date
            ]
        ):
            self.log('QP India custom data stopped coming.')
            self.liquidate()
            return

        # monthly rebalance
        if self.time.month == self._current_month:
            return
        self._current_month = self.time.month

        # store monthly price data
        for symbol, symbol_data in self._data.items():
            if symbol.value in self._excluded_tickers:
                continue
            if slice.contains_key(symbol) and slice[symbol]:
                price: float = slice[symbol].value
                if price != 0:
                    self._data[symbol].update_price(price)

        WML_Vol_Adj: Dict[Symbol, float] = {
            symbol: symbol_data.get_momentum() / symbol_data.get_volatility() 
            for symbol, symbol_data in self._data.items() 
            if symbol_data.is_ready() and
            symbol_data.get_momentum() != 0
        }

        # sort and divide into quantiles
        if len(WML_Vol_Adj) < self._quantile:
            self.log('Not enough stocks for further selection.')
            return

        sorted_WML_Vol_Adj: List[Symbol] = sorted(WML_Vol_Adj, key=WML_Vol_Adj.get, reverse=True)
        quantile: int = len(sorted_WML_Vol_Adj) // self._quantile
        long: List[Symbol] = sorted_WML_Vol_Adj[:quantile]
        short: List[Symbol] = sorted_WML_Vol_Adj[-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)