#region === Imports ===
from AlgorithmImports import *
import pywt
import numpy as np
from sklearn.svm import SVR
from sklearn.model_selection import GridSearchCV
#endregion


class FX_Daily_Wavelet_SVR_DePrado(QCAlgorithm):
    """
    DAILY FX — Lopez de Prado style:
    - Daily bars only
    - Wavelet denoise + SVR next-day forecast
    - Meta-label: skip on high daily vol
    - Walk-forward SVR re-training
    - Simple daily signal: trade daily open
    - Risk: fixed drawdown + trailing TP
    """

    def initialize(self):
        self.set_start_date(2020, 1, 1)
        self.set_end_date(2025, 1, 1)
        self.set_cash(100000)

        self._period = 50  # daily window, ~6 months of daily bars
        self._leverage = 100

        self.set_brokerage_model(BrokerageName.OANDA_BROKERAGE, AccountType.MARGIN)
        self.set_benchmark(SecurityType.FOREX, "EURUSD") 
        self.universe_settings.leverage = self._leverage
        self.universe_settings.resolution = Resolution.DAILY

        pairs = ["AUDUSD", "EURAUD", "EURCAD", "EURJPY", "GBPNZD", "AUDCAD", "NZDCHF", "AUDCHF", "CHFJPY", "AUDJPY", "NZDJPY"] 
        symbols = [Symbol.create(pair, SecurityType.FOREX, Market.OANDA) for pair in pairs]
        self.set_universe_selection(ManualUniverseSelectionModel(symbols))

        self.set_alpha(DailyWaveletAlphaModelDePrado(self._period))
        self.set_portfolio_construction(
            LeveragedWeightingPortfolioConstructionModel(Resolution.DAILY, self._leverage)
        )

        self.add_risk_management(MaximumDrawdownPercentPerSecurity(0.02))  # ~2% daily max SL
        self.add_risk_management(TrailingStopRiskManagementModel(0.02))     # ~2% trailing TP


class DailyWaveletAlphaModelDePrado(AlphaModel):
    def __init__(self, period):
        self._period = period
        self._wavelet = SVMWavelet()
        self._symbol_data = {}
        self._last_day = -1
        self._retrain_freq = 5  # re-train every 5 days

    def update(self, algorithm, data):
        insights = []
        day = algorithm.time.day
        if self._last_day == day: return []
        self._last_day = day

        for symbol, sd in self._symbol_data.items():
            prices = sd.prices()
            if len(prices) < self._period: continue

            forecast = self._wavelet.forecast(prices)
            ret = (forecast / prices[-1]) - 1

            daily_vol = np.std(prices[-20:]) / prices[-1]
            if daily_vol < 0.002 or daily_vol > 0.01:  
                algorithm.debug(f"Meta-label: skip {symbol} — daily vol {daily_vol:.4f}")
                continue

            if sd.retrain_counter % self._retrain_freq == 0:
                sd.train_wavelet = True
            sd.retrain_counter += 1

            if ret > 0.005:
                insights.append(
                    Insight.price(symbol, timedelta(days=1), InsightDirection.UP, weight=min(abs(ret), 0.05))
                )
            elif ret < -0.005:
                insights.append(
                    Insight.price(symbol, timedelta(days=1), InsightDirection.DOWN, weight=min(abs(ret), 0.05))
                )

            algorithm.insights.cancel([symbol])

        return insights

    def on_securities_changed(self, algorithm, changes):
        for security in changes.added_securities:
            self._symbol_data[security.symbol] = SymbolDataDePrado(
                algorithm, security.symbol, self._period)
        for security in changes.removed_securities:
            data = self._symbol_data.pop(security.symbol, None)
            if data:
                data.dispose()


class SymbolDataDePrado:
    def __init__(self, algorithm, symbol, period):
        self._algorithm = algorithm
        self._symbol = symbol
        self._close = RollingWindow[float](period)
        self.retrain_counter = 0

        self._consolidator = QuoteBarConsolidator(timedelta(days=1))
        self._consolidator.data_consolidated += self._on_consolidated
        algorithm.subscription_manager.add_consolidator(symbol, self._consolidator)

        history = algorithm.history[QuoteBar](symbol, period, Resolution.DAILY)
        for bar in history:
            self._consolidator.update(bar)

    def _on_consolidated(self, _, bar):
        self._close.add(bar.close)

    def prices(self):
        return np.array(list(self._close))[::-1]

    def dispose(self):
        self._close.reset()
        self._algorithm.subscription_manager.remove_consolidator(self._symbol, self._consolidator)


class LeveragedWeightingPortfolioConstructionModel(EqualWeightingPortfolioConstructionModel):
    def __init__(self, rebalance=Resolution.DAILY, leverage=100):
        super().__init__(rebalance)
        self.leverage = leverage

    def should_create_target_for_insight(self, insight):
        return insight.weight is not None

    def determine_target_percent(self, active_insights):
        result = {}
        for insight in active_insights:
            direction = insight.direction if self.respect_portfolio_bias(insight) else InsightDirection.FLAT
            result[insight] = direction * insight.weight * self.leverage
        return result


class SVMWavelet:
    def forecast(self, data):
        w = pywt.Wavelet('sym10')
        coeffs = pywt.wavedec(data, w)
        threshold = 0.75

        for i in range(len(coeffs)):
            if i > 0:
                coeffs[i] = pywt.threshold(coeffs[i], threshold * max(coeffs[i]))
            forecast = self._svm_forecast(coeffs[i])
            coeffs[i] = np.roll(coeffs[i], -1)
            coeffs[i][-1] = forecast

        return pywt.waverec(coeffs, w)[-1]

    def _svm_forecast(self, data, sample_size=20):
        X, y = self._partition_array(data, size=sample_size)
        grid = {'C': [.05, .1, .5, 1, 5, 10], 'epsilon': [0.001, 0.005, 0.01, 0.05, 0.1]}
        model = GridSearchCV(SVR(), grid, scoring='neg_mean_squared_error').fit(X, y).best_estimator_
        return model.predict(data[np.newaxis, -sample_size:])[0]

    def _partition_array(self, arr, size=None, splits=None):
        arrs, values = [], []
        if not (bool(size is None) ^ bool(splits is None)):
            raise ValueError('Size XOR Splits must be set')
        if size:
            arrs = [arr[i:i + size] for i in range(len(arr) - size)]
            values = [arr[i] for i in range(size, len(arr))]
        elif splits:
            size = len(arr) // splits
            arrs = [arr[i:i + size] for i in range(size - 1, len(arr) - 1, size)]
            values = [arr[i] for i in range(2 * size - 1, len(arr), size)]
        return np.array(arrs), np.array(values)