from AlgorithmImports import *
import numpy as np
import pandas as pd
from sklearn.covariance import LedoitWolf
from datetime import timedelta

class FactorMinVariancePortfolio(QCAlgorithm):
    """
    Factor-based long-only minimum variance portfolio.
    Trades 4 ETFs using factor exposures from SPY and BND.
    Uses Ledoit-Wolf shrinkage for factor covariance estimation.
    Rebalances roughly every 100 days using a 1-year lookback of daily data.
    """

    def initialize(self):
        self.set_start_date(2018, 1, 1)
        self.set_cash(100000)
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0
        self.universe_settings.resolution = Resolution.DAILY

        # -----------------------------
        # Add traded assets
        # -----------------------------
        asset_tickers = ["TLT", "GLD", "XLE", "QQQ"]
        self.asset_symbols = [
            self.add_equity(t, Resolution.DAILY).symbol
            for t in asset_tickers
        ]

        # -----------------------------
        # Add factor proxies (not traded)
        # -----------------------------
        factor_tickers = ["SPY", "BND"]
        self.factor_symbols = [
            self.add_equity(t, Resolution.DAILY).symbol
            for t in factor_tickers
        ]

        # -----------------------------
        # Parameters
        # -----------------------------
        self.lookback = 252  # Use 1 year of daily data
        self.rebalance_interval = timedelta(days=100)  # Rebalance roughly every 4 months
        self.last_rebalance = datetime.min
        self.epsilon = 1e-6  # Small ridge for numerical stability
        self.factor_exposure_limit = 0.5  # Not enforced yet

    def on_data(self, data: Slice):
        # -----------------------------
        # Rebalance check
        # -----------------------------
        if self.time - self.last_rebalance < self.rebalance_interval:
            return

        # -----------------------------
        # Fetch historical data
        # -----------------------------
        symbols = self.asset_symbols + self.factor_symbols
        history = self.history(symbols, self.lookback, Resolution.DAILY)
        if history.empty:
            return

        prices = history.close.unstack(level=0)
        returns = prices.pct_change().dropna()
        if len(returns) < 100:
            return

        # -----------------------------
        # Separate asset and factor returns
        # -----------------------------
        asset_returns = returns[self.asset_symbols].to_numpy()
        factor_returns = returns[self.factor_symbols].to_numpy()
        T, N = asset_returns.shape
        K = factor_returns.shape[1]

        # -----------------------------
        # Estimate factor betas (B) via OLS and residuals
        # -----------------------------
        X = np.column_stack([np.ones(T), factor_returns])
        B = np.zeros((N, K))
        residuals = np.zeros((N, T))
        for i in range(N):
            y = asset_returns[:, i]
            beta = np.linalg.lstsq(X, y, rcond=None)[0]
            B[i, :] = beta[1:]
            residuals[i, :] = y - X @ beta

        # -----------------------------
        # Factor covariance (Ledoit-Wolf shrinkage)
        # -----------------------------
        lw_f = LedoitWolf().fit(factor_returns)
        Sigma_f = lw_f.covariance_

        # -----------------------------
        # Idiosyncratic covariance
        # -----------------------------
        idio_var = np.var(residuals, axis=1)
        idio_var = np.maximum(idio_var, 1e-6)
        Sigma_eps = np.diag(idio_var)

        # -----------------------------
        # Total covariance matrix
        # -----------------------------
        Sigma = B @ Sigma_f @ B.T + Sigma_eps
        Sigma += self.epsilon * np.eye(N)

        # -----------------------------
        # Long-only minimum variance weights
        # -----------------------------
        w_closed = np.linalg.pinv(Sigma) @ np.ones(N)
        w_closed = np.maximum(w_closed, 0)
        w_closed /= w_closed.sum()
        self.Debug(f"Closed-form min-var weights (long-only): {w_closed}")

        # -----------------------------
        # Execute trades
        # -----------------------------
        for i, symbol in enumerate(self.asset_symbols):
            self.set_holdings(symbol, float(w_closed[i]))

        # -----------------------------
        # Post-trade holdings weights (on same date)
        # -----------------------------
        holdings_array = np.array([
            self.portfolio[symbol].quantity * self.portfolio[symbol].price 
            for symbol in self.asset_symbols
        ])
        holdings_weights = holdings_array / holdings_array.sum()
        self.Debug(f"Post-trade holdings weights: {holdings_weights}")

        # -----------------------------
        # Update last rebalance AFTER trading
        # -----------------------------
        self.last_rebalance = self.time

        # -----------------------------
        # Diagnostics
        # -----------------------------
        asset_vols = np.sqrt(np.diag(Sigma))
        self.Debug(f"Asset volatilities: {asset_vols}")
        self.Debug(f"B matrix (factor exposures):\n{B}")