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

from setup import TFSABuyingPowerModel


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(2012, 1, 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.qqq = self.add_equity_symbol("QQQ")
        self.qld = self.add_equity_symbol("QLD")
        self.tqqq = self.add_equity_symbol("TQQQ")
        self.spy = self.add_equity_symbol("SPY")
        self.gld = self.add_equity_symbol("GLD")
        self.xlu = self.add_equity_symbol("UPW")
        self.vixy = self.add_equity_symbol("UVXY")

        self.qqqm = self.add_equity_symbol("QQQM")
        self.live_mode_symbols: dict[str, Symbol] = {
            self.qqq.value: self.qqqm,
        }

        self.strat_position_symbols: list[Symbol] = [
            self.tqqq,
            self.qld,
            self.qqq,
            self.vixy,
            self.xlu,
            self.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.spy, 0), self.add_cash)
            self.schedule.on(
                self.date_rules.month_start(1),
                self.time_rules.after_market_close(self.spy, 0),
                self.add_cash,
            )
            # self.schedule.on(
            #     self.date_rules.month_end(15),
            #     self.time_rules.after_market_close(self.spy, 0),
            #     self.add_cash,
            # )

        self.schedule.on(
            self.date_rules.every_day(),
            self.time_rules.before_market_close(self.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.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())
        s.set_buying_power_model(TFSABuyingPowerModel())
        return s.Symbol

    def add_cash(self):
        dcaCash = 600
        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
        # No idea why QC keeps 
        buying_power = self.portfolio.margin_remaining - 10
        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.BuyingPowerModel.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.BuyingPowerModel.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[self.spy].has_data:
            return

        qqq_rsi_10 = self.rsi_2(self.qqq, 10)
        # spy_rsi_10 = self.rsi_2(self.spy, 10)
        spy_price = self.securities[self.spy].price
        spy_sma_200 = self.sma_2(self.spy, 200)
        spy_sma_30 = self.sma_2(self.spy, 30)

        if qqq_rsi_10 >= 79:
            self.set_holdings_2(self.vixy, 0.4)
        elif qqq_rsi_10 < 31:
            self.set_holdings_2(self.tqqq, 1)

        elif spy_price > spy_sma_200:
            if spy_price > spy_sma_30:
                self.set_holdings_2(self.qld, 1)
            else:
                self.set_holdings_2(self.xlu, 1)
        else:
            if spy_price > spy_sma_30:
                self.set_holdings_2(self.qld, 1)
            else:
                self.set_holdings_2(self.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

# region imports
from AlgorithmImports import *
# endregion

# Your New Python File
from AlgorithmImports import *

class TFSABuyingPowerModel(BuyingPowerModel):
    """
    A custom buying power model that effectively ignores margin requirements.
    This is NOT for real trading and should only be used for highly experimental backtesting.
    """

    def __init__(self) -> None:
        # No need to initialize base BuyingPowerModel arguments—override methods instead.
        pass

    def get_leverage(self, security: Security) -> float:
        # Return effectively infinite leverage
        return float(1)

    def has_sufficient_buying_power_for_order(self, parameters: HasSufficientBuyingPowerForOrderParameters) -> HasSufficientBuyingPowerForOrderResult:
        # Always report that there is sufficient buying power
        return HasSufficientBuyingPowerForOrderResult(True)

    def get_initial_margin_requirement(self, parameters: InitialMarginParameters) -> InitialMargin:
        # No initial margin required
        return InitialMargin(0)

    def get_maintenance_margin(self, parameters: MaintenanceMarginParameters) -> MaintenanceMargin:
        # No maintenance margin required
        return MaintenanceMargin(0)

    # def get_maximum_order_quantity_for_delta_buying_power(self, parameters: GetMaximumOrderQuantityForDeltaBuyingPowerParameters) -> GetMaximumOrderQuantityResult:
    #     # No limit on how many shares or contracts can be ordered
    #     return GetMaximumOrderQuantityResult(float('inf'), "No limit on order quantity.")

    # def get_buying_power(self, parameters: BuyingPowerParameters) -> BuyingPower:
    #     # Effectively unlimited buying power
    #     return BuyingPower(float('inf'))