# region imports
from AlgorithmImports import *
# endregion

# Your New Python File
class AssetArbitrageStrategy:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        self.long_trade_size = 0.05  # 5% of the portfolio for long trades
        self.short_trade_size = 0.03  # 3% of the portfolio for short trades
        self.long_stop_loss = 0.05  # 5% stop-loss for long trades
        self.short_stop_loss = 0.03  # 3% stop-loss for short trades
        self.max_portfolio_exposure = 0.80  # Maximum 80% portfolio exposure

    def Execute(self, indicators):
        contrarian_bands = indicators["contrarian_bands"]
        rsi = indicators["rsi"]
        trend = indicators["trend"]

        if not contrarian_bands.HasSignal() or not rsi.HasSignal():
            return

        price = self.algorithm.Securities[self.symbol].Price
        holdings = self.algorithm.Portfolio[self.symbol].Quantity
        average_price = self.algorithm.Portfolio[self.symbol].AveragePrice

        if price is None or price <= 0:
            self.algorithm.Debug(f"Skipping {self.symbol}: Invalid price {price}")
            return

        # Portfolio Exposure Check
        portfolio_exposure = sum([holding.HoldingsValue for holding in self.algorithm.Portfolio.Values]) / self.algorithm.Portfolio.TotalPortfolioValue
        if portfolio_exposure > self.max_portfolio_exposure:
            self.algorithm.Debug(f"Skipping trade for {self.symbol}: Portfolio exposure exceeds limit ({portfolio_exposure:.2%})")
            return

        # Long Entry
        if holdings == 0 and price < contrarian_bands.bbands.LowerBand.Current.Value and rsi.rsi.Current.Value < 30 and trend.IsUptrend():
            self.algorithm.SetHoldings(self.symbol, self.long_trade_size)

        # Short Entry
        elif holdings == 0 and price > contrarian_bands.bbands.UpperBand.Current.Value and rsi.rsi.Current.Value > 70 and trend.IsDowntrend():
            self.algorithm.SetHoldings(self.symbol, -self.short_trade_size)

        # Stop-Loss for Long Positions
        if holdings > 0 and price < average_price * (1 - self.long_stop_loss):
            self.algorithm.Debug(f"Stop-loss triggered for long {self.symbol} at price {price}")
            self.algorithm.Liquidate(self.symbol)

        # Stop-Loss for Short Positions
        if holdings < 0 and price > average_price * (1 + self.short_stop_loss):
            self.algorithm.Debug(f"Stop-loss triggered for short {self.symbol} at price {price}")
            self.algorithm.Liquidate(self.symbol)

        # Long Exit
        if holdings > 0 and price >= contrarian_bands.bbands.MiddleBand.Current.Value:
            self.algorithm.Liquidate(self.symbol)

        # Short Exit
        if holdings < 0 and price <= contrarian_bands.bbands.MiddleBand.Current.Value:
            self.algorithm.Liquidate(self.symbol)

# region imports
from AlgorithmImports import *
# endregion

# Your New Python File
class CustomBollingerBands:
    def __init__(self, algorithm, symbol, period=20, deviations=2):
        self.algorithm = algorithm
        self.symbol = symbol
        self.bbands = algorithm.BB(symbol, period, deviations, MovingAverageType.Simple, Resolution.Daily)

    def Update(self, data):
        if self.symbol in data and data[self.symbol] is not None:
            self.bbands.Update(data[self.symbol].EndTime, data[self.symbol].Close)

    def HasSignal(self):
        return self.bbands.IsReady


class RSIIndicator:
    def __init__(self, algorithm, symbol, period=14):
        self.algorithm = algorithm
        self.symbol = symbol
        self.rsi = algorithm.RSI(symbol, period, MovingAverageType.Simple, Resolution.Daily)

    def Update(self, data):
        if self.symbol in data and data[self.symbol] is not None:
            self.rsi.Update(data[self.symbol].EndTime, data[self.symbol].Close)

    def HasSignal(self):
        return self.rsi.IsReady


class HistoricalVolatility:
    def __init__(self, algorithm, symbol, period=20):
        self.algorithm = algorithm
        self.symbol = symbol
        self.returns = RollingWindow[float](period)

    def Update(self, data):
        if self.symbol in data and data[self.symbol] is not None:
            price = data[self.symbol].Close
            if self.returns.Count > 0 and self.returns[0] != 0:
                self.returns.Add(np.log(price / self.returns[0]))
            else:
                self.returns.Add(0)

    def GetVolatility(self):
        if self.returns.IsReady:
            return np.std(list(self.returns))
        return None


class TrendFilter:
    def __init__(self, algorithm, symbol, sma_period=50):
        self.algorithm = algorithm
        self.symbol = symbol
        self.sma = algorithm.SMA(symbol, sma_period, Resolution.Daily)

    def Update(self, data):
        if self.symbol in data and data[self.symbol] is not None:
            self.sma.Update(data[self.symbol].EndTime, data[self.symbol].Close)

    def IsUptrend(self):
        if self.sma.IsReady:
            price = self.algorithm.Securities[self.symbol].Price
            return price > self.sma.Current.Value
        return False

    def IsDowntrend(self):
        if self.sma.IsReady:
            price = self.algorithm.Securities[self.symbol].Price
            return price < self.sma.Current.Value
        return False

from AlgorithmImports import *
from Indicators import CustomBollingerBands, RSIIndicator, HistoricalVolatility, TrendFilter
from AssetStrategy import AssetArbitrageStrategy

class VolatilityArbitrage(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 10, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100000)

        # Add cryptocurrencies
        self.crypto_symbols = ["BTCUSD", "ETHUSD", "SOLUSD"]
        for crypto in self.crypto_symbols:
            self.AddCrypto(crypto, Resolution.Daily)

        # Universe selection for 100 stocks
        self.AddUniverse(self.CoarseSelectionFunction)

        # Extended warm-up period for better historical volatility data
        self.SetWarmUp(timedelta(days=60))

        self.symbol_data = {}

    def CoarseSelectionFunction(self, coarse):
        # Select liquid equities with price > $20 and sort by dollar volume in descending order
        filtered = [x for x in coarse if x.Price > 20 and x.DollarVolume > 1e7]
        sorted_by_volume = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        selected = [x.Symbol for x in sorted_by_volume[:100]]  # Top 100 stocks
        return selected + self.crypto_symbols  # Add cryptos to the universe

    def OnSecuritiesChanged(self, changes):
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            if symbol not in self.symbol_data:
                self.symbol_data[symbol] = {
                    "indicators": {
                        "contrarian_bands": CustomBollingerBands(self, symbol),
                        "rsi": RSIIndicator(self, symbol),
                        "hist_vol": HistoricalVolatility(self, symbol),
                        "trend": TrendFilter(self, symbol)
                    },
                    "strategy": AssetArbitrageStrategy(self, symbol)
                }

        for removed in changes.RemovedSecurities:
            symbol = removed.Symbol
            if symbol in self.symbol_data:
                del self.symbol_data[symbol]

    def OnData(self, data):
        if self.IsWarmingUp:
            return

        for symbol, data_set in self.symbol_data.items():
            if symbol in data:
                indicators = data_set["indicators"]
                indicators["contrarian_bands"].Update(data)
                indicators["rsi"].Update(data)
                indicators["hist_vol"].Update(data)
                indicators["trend"].Update(data)
                data_set["strategy"].Execute(indicators)