from datetime import timedelta
import math
from AlgorithmImports import *
import time
import numpy as np


class ActualIbkrFeeModel(FeeModel):
    def get_order_fee(self, parameters: OrderFeeParameters) -> OrderFee:
        security = parameters.security
        order = parameters.order
        # Optional check if it's equity
        if security.type != SecurityType.EQUITY:
            return OrderFee(CashAmount(0, "USD"))

        quantity = abs(order.quantity)
        price = security.price

        return OrderFee(CashAmount(get_fee(quantity, price), "USD"))


def get_fee(quantity, price, buffer=0):
    # Commission: $0.006 per share
    commission = 0.006 * quantity
    # Minimum commission: $0.80
    min_fee = 0.80
    # Maximum commission: 0.4% of trade value
    max_fee = 0.004 * quantity * price

    # Final commission per order
    final_fee = max(min_fee, min(commission, max_fee))

    return final_fee + buffer


class RSIRebalanceStrategy(QCAlgorithm):

    def Initialize(self):
        self.set_start_date(2015, 6, 1)
        # self.set_end_date(2022, 1, 1)
        # https://www.interactivebrokers.ca/en/accounts/tradingPermissions.php?ib_entity=ca
        # 5% ish for market orders
        self.ibkr_market_order_buffer = 0.02
        self.ibkr_fee_buffer = 25

        self.set_brokerage_model(
            BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.CASH
        )
        # self.fee_model = ActualIbkrFeeModel()
        self.SetSecurityInitializer(lambda s: s.SetLeverage(1))

        self.symbols = [
            "QQQ",
            "QLD",
            "TQQQ",
            "SPY",
            "GLD",
            "UPW",
            "UVXY",
            "LABU",
            "XLP",
            "QQQM",
        ]

        self.equities = {}

        for name in self.symbols:
            self.equities[name] = self.add_equity_symbol(name)

        self.live_mode_symbols: dict[str, Symbol] = {
            self.equities["QQQ"].value: self.equities["QQQM"],
        }

        self.strat_position_symbols: list[Symbol] = [
            self.equities["TQQQ"],
            self.equities["QLD"],
            self.equities["QQQ"],
            self.equities["UVXY"],
            self.equities["UPW"],
            self.equities["GLD"],
        ]

        if not self.live_mode:
            self.set_warmup(timedelta(days=250))
            self.set_cash(8000)

            # self.schedule.on(self.date_rules.week_end(), self.time_rules.after_market_close(self.equities["SPY"], 0), self.add_cash)
            self.schedule.on(
                self.date_rules.month_start(1),
                self.time_rules.after_market_close(self.equities["SPY"], 0),
                self.add_cash,
            )
            self.schedule.on(
                self.date_rules.month_start(15),
                self.time_rules.after_market_close(self.equities["SPY"], 0),
                self.add_cash,
            )
            # self.schedule.on(
            #     self.date_rules.month_end(15),
            #     self.time_rules.after_market_close(self.equities["SPY"], 0),
            #     self.add_cash,
            # )

        self.schedule.on(
            self.date_rules.every_day(),
            self.time_rules.before_market_close(self.equities["SPY"], 6),
            self.reset_liquidated,
        )

        for offset in np.arange(4, 0, -1):
            self.schedule.on(
                self.date_rules.every_day(),
                self.time_rules.before_market_close(self.equities["SPY"], offset),
                self.rebalance,
            )

        self.cashInvested = self.portfolio.cash_book["USD"].amount
        self.liquidated = False
        self.stupid_order_log = ""

    # def on_data(self, slice: Slice) -> None:
    #     # Obtain the mapped TradeBar of the symbol if any
    #     # if not self.portfolio[self.spy].invested:
    #     #     self.set_holdings(self.spy, 1, True)
    #     if self.live_mode and self.securities[self.qqq].has_data:
    #         self.rebalance()

    def reset_liquidated(self):
        self.liquidated = False

    def add_equity_symbol(self, symbol: str) -> Symbol:
        s = self.add_equity(
            symbol, Resolution.MINUTE, data_normalization_mode=DataNormalizationMode.RAW
        )
        s.set_settlement_model(ImmediateSettlementModel())
        # s.set_fee_model(self.fee_model)
        s.set_fill_model(ImmediateFillModel())
        return s.Symbol

    def add_cash(self):
        dcaCash = 230
        self.portfolio.cash_book["USD"].add_amount(dcaCash)
        self.cashInvested += dcaCash

    def check_and_get_live_symbol(self, symbol: Symbol):
        if (
            self.portfolio.total_portfolio_value < 20000
            and self.live_mode
            and symbol.value in self.live_mode_symbols
        ):
            return self.live_mode_symbols[symbol.value]
        return symbol

    def rsi_2(self, symbol, period):
        warmup = int(round_up(11 * math.sqrt(period) + 5.5 * period, 0))
        extension = min(warmup, 250)
        r_w = RollingWindow[float](extension)
        history = self.history(symbol.value, extension - 1, Resolution.DAILY)
        for historical_bar in history:
            r_w.add(historical_bar.close)
        while r_w.count < extension:
            current_price = self.securities[symbol.value].price
            r_w.add(current_price)
        if r_w.is_ready:
            average_gain = 0
            average_loss = 0
            gain = 0
            loss = 0
            for i in range(extension - 1, extension - period - 1, -1):
                gain += max(r_w[i - 1] - r_w[i], 0)
                loss += abs(min(r_w[i - 1] - r_w[i], 0))
            average_gain = gain / period
            average_loss = loss / period
            for i in range(extension - period - 1, 0, -1):
                average_gain = (
                    average_gain * (period - 1) + max(r_w[i - 1] - r_w[i], 0)
                ) / period
                average_loss = (
                    average_loss * (period - 1) + abs(min(r_w[i - 1] - r_w[i], 0))
                ) / period
            if average_loss == 0:
                return 100
            else:
                rsi = 100 - (100 / (1 + average_gain / average_loss))
                return rsi
        else:
            return None

    def sma_2(self, symbol: Symbol, period):
        r_w = RollingWindow[float](period)
        history = self.history(symbol.value, period - 1, Resolution.DAILY)
        for historical_bar in history:
            r_w.add(historical_bar.close)
        while r_w.count < period:
            current_price = self.securities[symbol.value].price
            r_w.add(current_price)
        if r_w.is_ready:
            sma = sum(r_w) / period
            return sma
        else:
            return 0

    def get_current_strat_positions(self):
        """
        Returns:
            list[Symbol]: A list of stock symbols that this strategy buys and sells, not symbols for indicators.
        """
        current_positions = []
        for strat_symbol in self.strat_position_symbols:
            strat_symbol = self.check_and_get_live_symbol(strat_symbol)

            if self.portfolio[strat_symbol].invested:
                current_positions.append(strat_symbol)

        return current_positions

    def custom_log(self, message: str, only_live_mode: bool = True):
        if only_live_mode and self.live_mode:
            self.log(message)

    def set_holdings_2(self, symbol: Symbol, portion=1):
        """IBKR Brokage Model somehow doesn't wait till liquidation finishes in set_holdings(symbol, 1, True)
        So we liquidate explicitly first and set_holdings after
        """
        symbol = self.check_and_get_live_symbol(symbol)

        # liquidate any other symbols when switching
        if (
            not self.liquidated
            and not self.portfolio[symbol].invested
            and self.portfolio.invested
        ):
            # for curr_pos in current_positions:
            #     self.liquidate(curr_pos, tag="Liquidated")
            self.liquidate()
            self.liquidated = True
            return

        if (
            len(self.transactions.get_open_orders()) > 0
            or self.portfolio.unsettled_cash > 0
            or (
                self.liquidated
                and self.portfolio.total_portfolio_value != self.portfolio.cash
            )
            or not self.securities[symbol].has_data
        ):
            self.custom_log(f"Skip order submissions")
            return

        # Usually portion < 1 is for risky asset like UVXY, no need to keep it's portion
        # as we don't hold it for long
        if portion < 1:
            if self.liquidated and not self.portfolio.invested:
                self.set_holdings(symbol, portion)
                return
            elif self.portfolio[symbol].invested:
                self.set_holdings(symbol, portion)
                return

        # Calculate for limit order
        # using 99% of buying power to avoid "Insufficient buying power to complete orders" error
        # no idea why QC's engine has weird initial margin stuff
        # TFSA doesn't have any initial margin requirements, so we can use 100% - 10 dollars
        buying_power = (
            self.portfolio.margin_remaining - 10
            if self.live_mode
            else self.portfolio.margin_remaining * 0.99
        )
        symbol_price = self.securities[symbol].ask_price

        # 5 cents buffer for limit order, IBKR will find lowest ask
        limit_price = round_up(symbol_price + 0.03)

        shares_num: int = math.floor(buying_power / limit_price)

        security = self.securities[symbol]
        initial_margin_params = InitialMarginParameters(security, shares_num)
        initial_margin_required = (
            security.buying_power_model.get_initial_margin_requirement(
                initial_margin_params
            )
        )
        # Extract the numeric value from the InitialMargin object
        required_margin_value = initial_margin_required.value

        while shares_num >= 2 and required_margin_value > buying_power - get_fee(
            shares_num, limit_price, 10
        ):
            shares_num = shares_num - 1
            initial_margin_params = InitialMarginParameters(security, shares_num)
            initial_margin_required = (
                security.buying_power_model.get_initial_margin_requirement(
                    initial_margin_params
                )
            )
            required_margin_value = initial_margin_required.value

        # order_amount = shares_num * limit_price
        # satisfy_ibkr_market_order = (1 - (order_amount / cash)) > self.ibkr_market_order_buffer

        if shares_num >= 2 and required_margin_value < buying_power - get_fee(
            shares_num, limit_price, 10
        ):
            self.stupid_order_log = f"Placing order for {symbol.value}, {shares_num} shares, {symbol_price} price, {required_margin_value} amount while buying power is {buying_power} and fee is {get_fee(shares_num, limit_price, 10)}"

            self.custom_log(
                f"Placing order for {symbol.value}, {shares_num} shares, {symbol_price} price, {required_margin_value} amount"
            )
            self.limit_order(symbol, shares_num, limit_price)
            return

    def rebalance(self):
        if self.time < self.start_date:
            return

        if not self.securities["SPY"].has_data:
            return

        rsi = self.rsi_2
        equities = self.equities

        qqq_rsi_10 = rsi(equities["QQQ"], 10)
        # spy_rsi_10 = rsi(equities["SPY"], 10)
        spy_price = self.securities[equities["SPY"]].price
        spy_sma_200 = self.sma_2(equities["SPY"], 200)
        spy_sma_30 = self.sma_2(equities["SPY"], 30)

        if qqq_rsi_10 > 79 or rsi(equities["XLP"], 10) > 75:
            self.set_holdings_2(equities["UVXY"], 0.4)
            
        elif qqq_rsi_10 < 31:
            self.set_holdings_2(equities["TQQQ"], 1)
        elif rsi(equities["LABU"], 10) < 25:
            self.set_holdings_2(equities["LABU"], 1)
            
        elif spy_price > spy_sma_200:
            if spy_price > spy_sma_30:
                self.set_holdings_2(equities["QLD"], 1)
            else:
                self.set_holdings_2(equities["UPW"], 1)
        else:
            if spy_price > spy_sma_30:
                self.set_holdings_2(equities["QLD"], 1)
            else:
                self.set_holdings_2(equities["GLD"], 1)

    def on_end_of_algorithm(self) -> None:
        self.debug(f"Cash invested: {self.cashInvested}")

    def on_order_event(self, order_event: OrderEvent) -> None:
        # order = self.transactions.get_order_by_id(order_event.order_id)
        if order_event.status == OrderStatus.INVALID:
            # amount = order.quantity * order.limit_price
            self.debug(self.stupid_order_log)


def round_up(n, decimals=2):
    multiplier = 10**decimals
    return math.ceil(n * multiplier) / multiplier