# main.py
from AlgorithmImports import *
from datetime import timedelta, datetime
from time import sleep
from QuantConnect.Indicators import *
from symbols import SYMBOLS


class OptimizedBuySellStrategy(QCAlgorithm):
    def initialize(self):
        self.webhook_url = None  # set to "https://wwww.mysite.com/webhook" in order to work
        self.webhook_headers = {}  # set to { 'Authorization': 'Basic MY_PASSWORD' } if you have basic auth in place
        self.bar_counts = {}

        # Get parameters with defaults
        start_date_str = self.get_parameter("start_date") or "2024-01-01"
        end_date_str = self.get_parameter("end_date") or "2024-12-31"
        cash_str = self.get_parameter("cash") or 100_000 * len(SYMBOLS)

        # Parse parameters
        try:
            start_date = datetime.strptime(start_date_str, "%Y-%m-%d")
            end_date = datetime.strptime(end_date_str, "%Y-%m-%d")
            cash = float(cash_str)
        except ValueError as e:
            raise ValueError(f"Invalid date or cash format: {e}")

        # Set values
        self.set_start_date(start_date.year, start_date.month, start_date.day)
        self.set_end_date(end_date.year, end_date.month, end_date.day)
        self.set_cash(cash)

        self.timeframe_map = {
            "15m": timedelta(minutes=15),
            "30m": timedelta(minutes=30),
            "45m": timedelta(minutes=45),
            "1H": timedelta(hours=1),
            "2H": timedelta(hours=2),
            "3H": timedelta(hours=3),
            "4H": timedelta(hours=4),
            "1D": timedelta(days=1),
            "1W": timedelta(weeks=1)
        }

        # set brokerage model for backtesting
        self.set_brokerage_model(BrokerageName.ALPACA, AccountType.MARGIN)  # ALPACA has zero fees
        self.universe_settings.leverage = 2

        self.symbol_data = {}

        # [OUTDATED] Get single symbol_info from parameters
        # symbol_info_str = self.get_parameter("symbol_info")
        # if symbol_info_str:
        #     import json
        #     symbol_info = json.loads(symbol_info_str)
        # else:
        #     raise ValueError("No symbol_info parameter provided")

        for symbol_info in SYMBOLS:
            # Process single symbol
            symbol = self.add_equity(symbol_info["ticker"], Resolution.MINUTE).symbol
            timeframe = symbol_info["timeframe"]

            if timeframe not in self.timeframe_map:
                raise ValueError(f"Invalid timeframe '{timeframe}' for {symbol_info['ticker']}")

            self.symbol_data[symbol] = {
                "ticker": symbol_info["ticker"],
                "timeframe": self.timeframe_map[timeframe],
                "risk_index": symbol_info.get("risk_index", 2),
                "last_trade_time": None,
                "entry_price": None,
                "stop_loss": None,
                "take_profit": None,
                "is_liquidating": False,
                "consolidator": self.consolidate(symbol, self.timeframe_map[timeframe], lambda bar: self.bar_handler(bar)),
                "rsi": RelativeStrengthIndex(14),
                "macd": MovingAverageConvergenceDivergence(12, 26, 9),
                "atr": AverageTrueRange(14),
                "support": Minimum(50),
                "resistance": Maximum(50),
                "history": RollingWindow[TradeBar](50),
                "bullish_run_window": RollingWindow[TradeBar](4)
            }

        # Warm up indicators
        self.set_warm_up(50)

    def bar_handler(self, bar: TradeBar):
        symbol = bar.symbol
        data = self.symbol_data[symbol]

        # Update indicators
        data["rsi"].update(bar.end_time, bar.close)
        data["macd"].update(bar.end_time, bar.close)
        data["atr"].update(bar)
        data["support"].update(bar.end_time, bar.low)
        data["resistance"].update(bar.end_time, bar.high)
        data["history"].add(bar)
        data["bullish_run_window"].add(bar)

    def on_data(self, data):
        for symbol in self.symbol_data:
            if not self.symbol_data[symbol]["rsi"].is_ready:
                continue

            data = self.symbol_data[symbol]
            holding = self.portfolio[symbol]

            # Configurable parameters
            risk_index = data["risk_index"]
            rsi_overbought = 70 if risk_index == 2 else 65 if risk_index == 1 else 75
            rsi_oversold = 30 if risk_index == 2 else 35 if risk_index == 1 else 25
            atr_multiplier = 2.0 if risk_index == 2 else 1.5 if risk_index == 1 else 2.5
            cooldown_bars = 10

            # Get current bar and previous bars
            current_bar = data["history"][0]
            prev_bar = data["history"][1] if data["history"].count > 1 else None

            if prev_bar is None or data["bullish_run_window"].count < 4:
                continue

            # Cooldown management
            cooldown_over = (data["last_trade_time"] is None or
                             (self.time - data["last_trade_time"]) > (data["timeframe"] * cooldown_bars))

            # Bullish run detection
            current_highs = [data["bullish_run_window"][i].high for i in range(3)]
            prev_high = data["bullish_run_window"][3].high
            current_lows = [data["bullish_run_window"][i].low for i in range(3)]
            prev_low = data["bullish_run_window"][3].low
            bullish_run = max(current_highs) > prev_high and min(current_lows) > prev_low

            # Candle patterns
            bullish_engulfing = (current_bar.close > prev_bar.open and
                                 current_bar.close > current_bar.open)
            hammer = (current_bar.close > current_bar.open and
                      (prev_bar.low - current_bar.low) > 2 * (current_bar.close - current_bar.open))

            # Buy signal
            buy_signal = cooldown_over and not holding.invested and not data["is_liquidating"] and (
                    (data["rsi"].current.value < rsi_oversold and
                     data["macd"].fast.current.value > data["macd"].signal.current.value) or
                    bullish_engulfing or
                    hammer or
                    bullish_run
            )

            # Sell signal
            sell_signal = holding.invested and not data["is_liquidating"] and (
                    (data["rsi"].current.value > rsi_overbought and
                     data["macd"].fast.current.value < data["macd"].signal.current.value) or
                    current_bar.close >= data["resistance"].current.value
            )

            # Execute trades
            if buy_signal:
                weight = self.universe_settings.leverage / len(self.symbol_data) / 2
                self.set_holdings(symbol, weight)
                data["stop_loss"] = current_bar.low - (data["atr"].current.value * atr_multiplier)
                data["take_profit"] = data["resistance"].current.value
                data["last_trade_time"] = self.time
                data["entry_price"] = current_bar.close
                data["is_liquidating"] = False
                self.log(f"Buy {symbol} at {current_bar.close}")
                self.send_notification(self.buy_text(data))

            if holding.invested and not data["is_liquidating"]:
                if bullish_run and data["stop_loss"] is not None:
                    data["stop_loss"] = max(data["stop_loss"],
                                            current_bar.close - (data["atr"].current.value * atr_multiplier))

                stop_loss_triggered = data["stop_loss"] is not None and current_bar.close <= data["stop_loss"]
                take_profit_triggered = data["take_profit"] is not None and current_bar.close >= data["take_profit"]

                if stop_loss_triggered or take_profit_triggered or sell_signal:
                    self.liquidate(symbol)
                    data["is_liquidating"] = True
                    price = current_bar.close
                    profit = price - data["entry_price"]
                    profit_percent = (profit / data["entry_price"]) * 100
                    self.log(f"Sell {symbol} at {price}, Profit: {profit_percent:.2f}%")
                    if stop_loss_triggered:
                        reason = "stop_loss"
                    elif take_profit_triggered:
                        reason = "take_profit"
                    else:
                        reason = "sell_signal"
                    self.send_notification(self.sell_text(data, reason, price, profit_percent))
                    data["stop_loss"] = None
                    data["take_profit"] = None
                    data["entry_price"] = None
                    data["last_trade_time"] = self.time

            if not holding.invested and data["is_liquidating"]:
                data["is_liquidating"] = False

    def buy_text(self, data, reason="buy_signal"):
        # format: Smart BUY (ticker, buy_reason, purchase price, take profit price, stop loss price, risk_index, time)
        return (f'Smart BUY ({data["ticker"]}, {reason}, {round(data["entry_price"], 2)}, {round(data["take_profit"], 2)}, '
                f'{round(data["stop_loss"], 2)}, {data["risk_index"]}, {self.time})')

    def sell_text(self, data, reason, price, profit_percent):
        # format: Smart SELL (ticker, reason, sell price, profit %, risk_index, time)
        return (f'Smart SELL ({data["ticker"]}, {reason}, {round(price, 2)}, {round(profit_percent, 2)}%, '
                f'{data["risk_index"]}, {self.time})')

    def send_notification(self, text):
        if not self.live_mode:
            self.log(f"Webhook message: {text}")
            return

        if not self.webhook_url:
            return

        # try sending notification 3 times in case of exception
        for _ in range(3):
            try:
                self.notify.web(address=self.webhook_url, data=text, headers=self.webhook_headers)
                break
            except Exception as e:
                self.debug(f"Exception while trying to send web hook update: {type(e)} {e}")
                sleep(1)

# This file defines the symbols and their respective candlestick intervals.
SYMBOLS = [
    {"ticker": "AAPL", "timeframe": "3H", "risk_index": 2},  # 2 is default risk_index
    {"ticker": "MSFT", "timeframe": "1D"},
    {"ticker": "TSLA", "timeframe": "1W"},
    {"ticker": "F", "timeframe": "1D"},
    {"ticker": "GE", "timeframe": "4H"},
    {"ticker": "NVDA", "timeframe": "4H"},
    {"ticker": "BAC", "timeframe": "1D"},
    {"ticker": "C", "timeframe": "2H"},
    {"ticker": "JPM", "timeframe": "1D"},
    {"ticker": "AMZN", "timeframe": "1D"},
    {"ticker": "GOOGL", "timeframe": "3H"},
    {"ticker": "GOOG", "timeframe": "2H"},
    {"ticker": "META", "timeframe": "2H"},
    {"ticker": "NFLX", "timeframe": "4H"},
    {"ticker": "INTC", "timeframe": "1D"},
    {"ticker": "QCOM", "timeframe": "4H"},
    {"ticker": "CSCO", "timeframe": "4H"},
    {"ticker": "VZ", "timeframe": "2H"},
    {"ticker": "T", "timeframe": "1D"},
    {"ticker": "TMUS", "timeframe": "1D"},
    {"ticker": "CMCSA", "timeframe": "1H"},
    {"ticker": "CHTR", "timeframe": "3H"},
    {"ticker": "MCD", "timeframe": "3H"},
    {"ticker": "SBUX", "timeframe": "1D"}

    # Add additional symbols as needed...
]