# region imports
from datetime import timedelta
from AlgorithmImports import *
import pandas as pd
from io import StringIO
import pytz

# endregion


class HyperActiveBlueRat(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 9, 29)
        self.SetCash(100000)

        # parameters
        self.CPT = self.GetParameter("CPT", 1000.0)
        self.SLP = self.GetParameter("SLP", 25.0)
        self.EBE = self.GetParameter("EBE", 120.0)
        self.THEP = self.GetParameter("THEP", 33.0)
        self.PMSP = self.GetParameter("PMSP", 50.0)
        self.OTL = self.GetParameter("OTL", 5)
        self.VP = self.GetParameter("VP", 15)

        csv = self.Download(
            "https://www.dropbox.com/scl/fi/nmhib4wgzybsi7klinm73/UOA-Signals_2023_10_27.csv?rlkey=ktunyhjpgqbho6a2fhnsot3xf&dl=1"
        )
        self.df = pd.read_csv(StringIO(csv))

        self.symbols = []
        self.options_symbols = []
        self.signals = {}
        self.entry_order = {}
        self.take_profit = {}
        self.stop_loss = {}
        self.look_for_trigger = []
        self.look_for_validation = []
        self.SMAs = {}

        self.AddEquity("SPY", dataNormalizationMode=DataNormalizationMode.Raw)
        for i in range(len(self.df)):
            if self.df.iloc[i, 2] not in self.symbols:
                try:
                    symbol = self.AddEquity(
                        self.df.iloc[i, 2],
                        Resolution.Minute,
                        dataNormalizationMode=DataNormalizationMode.Raw,
                    ).Symbol
                    self.SMAs[symbol.Value] = self.SMA(
                        symbol, self.VP, Resolution.Minute
                    )
                    options_symbol = self.AddSecurity(
                        SecurityType.Option, symbol, Resolution.Minute
                    ).Symbol
                    self.options_symbols.append(options_symbol)
                    self.symbols.append(self.df.iloc[i, 2])
                except:
                    self.Debug(f"Error adding {self.df.iloc[i, 2]}")

        self.fill_signals()
        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.BeforeMarketClose("SPY", 1),
            self.market_close,
        )
        # TODO: Add consolidators
        # TODO: Add Warmup

    def OnData(self, data: Slice):
        self.check_for_expiry()
        self.check_for_validation()
        # check for exits
        for order_ticket in list(self.entry_order.values()):
            if (
                order_ticket.Status == OrderStatus.Submitted
                and order_ticket.Time.astimezone(pytz.timezone("US/Eastern")).replace(
                    tzinfo=None
                )
                + timedelta(minutes=self.OTL)
                <= self.Time
            ):
                order_ticket.Cancel("OTL")

        for symbol in self.symbols:
            for signal in self.signals.get(self.Time.strftime("%m/%d/%Y %H:%M"), []):
                if signal["symbol"] == symbol:
                    contracts = self.OptionChainProvider.GetOptionContractList(
                        symbol, data.Time
                    )
                    if contracts != None:
                        calls = [
                            x for x in contracts if x.ID.OptionRight == OptionRight.Call
                        ]
                        puts = [
                            x for x in contracts if x.ID.OptionRight == OptionRight.Put
                        ]
                        expiry = signal["expiration"]

                        if signal["Alert"] in ("Bullish Call Buying", "Bullside Idea"):
                            contract = None
                            current_price = self.Securities[symbol].Price
                            for call in calls:
                                if call.ID.Date.strftime("%d-%b") != expiry:
                                    continue
                                if contract is None:
                                    contract = call
                                elif (
                                    call.ID.StrikePrice > contract.ID.StrikePrice
                                    and call.ID.StrikePrice < current_price
                                ):
                                    contract = call

                        if signal["Alert"] in ("Bearish Put Buying", "Bearside Idea"):
                            contract = None
                            current_price = self.Securities[symbol].Price
                            for put in puts:
                                if put.ID.Date.strftime("%d-%b") != expiry:
                                    continue
                                if contract is None:
                                    contract = put
                                elif (
                                    put.ID.StrikePrice < contract.ID.StrikePrice
                                    and put.ID.StrikePrice > current_price
                                ):
                                    contract = put

                        if contract is not None:
                            self.AddOptionContract(contract, Resolution.Minute)
                            self.look_for_trigger.append((signal, contract))
        self.check_for_trigger()

    def OnOrderEvent(self, orderEvent: OrderEvent):
        if orderEvent.Status != OrderStatus.Filled:
            return

        if self.entry_order[orderEvent.Symbol].OrderId == orderEvent.OrderId:
            self.take_profit[orderEvent.Symbol] = self.LimitOrder(
                orderEvent.Symbol,
                -orderEvent.Quantity,
                round(orderEvent.FillPrice + orderEvent.FillPrice * self.THEP / 100, 2),
                "TP",
            )
            self.stop_loss[orderEvent.Symbol] = self.StopMarketOrder(
                orderEvent.Symbol,
                -orderEvent.Quantity,
                round(orderEvent.FillPrice * (1 - self.SLP / 100), 2),
                "SL",
            )
        elif self.take_profit[orderEvent.Symbol].OrderId == orderEvent.OrderId:
            self.stop_loss[orderEvent.Symbol].Cancel()

        elif self.stop_loss[orderEvent.Symbol].OrderId == orderEvent.OrderId:
            self.take_profit[orderEvent.Symbol].Cancel()

    def make_order(self, symbol: OptionSymbol):
        underlying_price = self.Securities[symbol.Value.split(" ")[0]].Price
        if symbol.ID.OptionRight == OptionRight.Put and underlying_price < self.PMSP:
            self.Log(f"Ignoring {symbol.Value} order as the price is < ${self.PMSP}")
            return
        if self.Securities.ContainsKey(symbol) is False:
            self.Log(f"couldn't find security {symbol} ignoring signal")
            return
        if self.Portfolio[symbol.Value].Invested:
            self.Log(f"symbol already invested {symbol.Value} ignoring signal")
            return
        if (
            symbol in self.entry_order
            and self.entry_order[symbol].Status == OrderStatus.Submitted
        ):
            self.Log(f"symbol already in entry order {symbol.Value} ignoring signal")
            return
        bid = self.Securities[symbol.Value].BidPrice
        ask = self.Securities[symbol.Value].AskPrice
        if bid == 0:
            self.Log(f"couldn't find bid/ask for {symbol} ignoring signal")
            return
        lmt_price = round(((bid + ask) / 2 + ask) / 2, 2)
        qty = self.CPT // (lmt_price * 100)
        if qty == 0 or lmt_price == 0:
            return
        self.entry_order[symbol] = self.LimitOrder(symbol, qty, lmt_price)

    def market_close(self):
        self.Liquidate(tag="End of day")
        self.look_for_validation.clear()
        self.look_for_trigger.clear()

    def check_for_expiry(self):
        for position in self.Portfolio.items():
            symbol = position[0]
            qty = position[1].AbsoluteQuantity
            if qty > 0:
                expiry = symbol.ID.Date
                expiry.replace(hour=16)
                if self.Time + timedelta(minutes=self.EBE) >= expiry:
                    self.Liquidate(symbol, "EBE hit")

    def check_for_trigger(self):
        for item in list(self.look_for_trigger):
            if (
                item[1].ID.OptionRight == OptionRight.Call
                and self.Securities[item[0]["symbol"]].High >= item[0]["Trigger price"]
            ) or (
                item[1].ID.OptionRight == OptionRight.Put
                and self.Securities[item[0]["symbol"]].Low <= item[0]["Trigger price"]
            ):
                self.look_for_validation.append((item, self.Time))
                self.look_for_trigger.remove(item)

    def check_for_validation(self):
        for item in list(self.look_for_validation):
            if item[1] + timedelta(minutes=self.VP) <= self.Time and (
                (
                    item[0][1].ID.OptionRight == OptionRight.Call
                    and self.SMAs[item[0][0]["symbol"]].Current.Value
                    >= item[0][0]["Trigger price"]
                )
                or (
                    item[0][1].ID.OptionRight == OptionRight.Put
                    and self.SMAs[item[0][0]["symbol"]].Current.Value
                    <= item[0][0]["Trigger price"]
                )
            ):
                self.make_order(item[0][1])
                self.look_for_validation.remove(item)

    def fill_signals(self):
        for i in range(len(self.df)):
            date = self.df.iloc[i, 1]
            if date not in self.signals:
                self.signals[date] = []
            self.signals[date].append(
                {
                    "symbol": self.df.iloc[i, 2],
                    "Alert": self.df.iloc[i, 3],
                    "Stock price": float(self.df.iloc[i, 5]),
                    "Trigger price": float(self.df.iloc[i, 6]),
                    "expiration": self.df.iloc[i, 8],
                }
            )