tests/test_e2e.py

"""
Tests for the End-to-End Differentiable Portfolio Optimization pipeline.
Covers the differentiable layer, neural network, training loop, caching,
and integration with the production CVXPY solver.
"""

import os
import sys
import shutil
import tempfile
import numpy as np
import pandas as pd
import pytest

# ─────────────────────────────────────────────────────────────────────
# Skip entire module if torch / cvxpylayers are not installed
# ─────────────────────────────────────────────────────────────────────
torch = pytest.importorskip("torch")
pytest.importorskip("cvxpylayers")

from differentiable_optimizer import DifferentiablePortfolioLayer
from e2e_forecast_model import (
    PortfolioForecastNetwork,
    E2EPortfolioTrainer,
    save_model,
    load_model,
)


# ─────────────────────────────────────────────────────────────────────
# Fixtures
# ─────────────────────────────────────────────────────────────────────

N_ASSETS   = 5
N_FEATURES = 7
BATCH      = 4
TICKERS    = ["SPY", "QQQ", "TLT", "GLD", "AAPL"]


@pytest.fixture
def mu_batch():
    """Random expected-return vector, shape (BATCH, N_ASSETS)."""
    return torch.randn(BATCH, N_ASSETS, dtype=torch.float32) * 0.01


@pytest.fixture
def L_batch():
    """Random lower-Cholesky factor, shape (BATCH, N, N)."""
    # Create a valid PD matrix, then Cholesky
    A = torch.randn(BATCH, N_ASSETS, N_ASSETS) * 0.01
    Sigma = torch.bmm(A, A.transpose(1, 2)) + 0.01 * torch.eye(N_ASSETS)
    return torch.linalg.cholesky(Sigma)


@pytest.fixture
def synthetic_returns_df():
    """Synthetic daily returns DataFrame aligned to TICKERS."""
    np.random.seed(42)
    T = 600
    dates = pd.bdate_range("2020-01-01", periods=T)
    data = np.random.randn(T, N_ASSETS) * 0.01
    return pd.DataFrame(data, index=dates, columns=TICKERS)


@pytest.fixture
def synthetic_features_dict(synthetic_returns_df):
    """Minimal features_dict mimicking ``data.build_ml_features()``."""
    features = {}
    for t in TICKERS:
        n = len(synthetic_returns_df)
        df = pd.DataFrame({
            "ret": synthetic_returns_df[t],
            "target": synthetic_returns_df[t].shift(-21).rolling(21).sum(),
            "mom_1m": synthetic_returns_df[t].rolling(21).sum().shift(1),
            "mom_3m": synthetic_returns_df[t].rolling(63).sum().shift(1),
            "mom_6m": synthetic_returns_df[t].rolling(126).sum().shift(1),
            "rev_5d": synthetic_returns_df[t].rolling(5).sum().shift(1),
            "vol_21d": synthetic_returns_df[t].rolling(21).std().shift(1),
            "beta_63d": np.random.randn(n) * 0.5 + 1.0,
            "smb_21d": np.random.randn(n) * 0.01,
        }, index=synthetic_returns_df.index)
        features[t] = df.dropna()
    return features


# ─────────────────────────────────────────────────────────────────────
# 1. DifferentiablePortfolioLayer
# ─────────────────────────────────────────────────────────────────────

class TestDifferentiableLayer:

    def test_forward_produces_valid_weights(self, mu_batch, L_batch):
        layer = DifferentiablePortfolioLayer(N_ASSETS, risk_factor=3.0)
        w = layer(mu_batch, L_batch)

        assert w.shape == (BATCH, N_ASSETS)
        # Weights sum to 1
        sums = w.sum(dim=-1)
        np.testing.assert_allclose(sums.detach().numpy(), 1.0, atol=1e-3)
        # Non-negative (long-only)
        assert (w.detach().numpy() >= -1e-4).all()

    def test_gradient_flows_through_mu(self, L_batch):
        layer = DifferentiablePortfolioLayer(N_ASSETS, risk_factor=3.0)
        mu = torch.randn(BATCH, N_ASSETS, requires_grad=True)
        w = layer(mu, L_batch)

        # Dummy loss: negative portfolio return
        loss = -(w * mu).sum()
        loss.backward()

        assert mu.grad is not None
        assert mu.grad.shape == mu.shape
        assert torch.isfinite(mu.grad).all()

    def test_short_allowed(self, mu_batch, L_batch):
        layer = DifferentiablePortfolioLayer(
            N_ASSETS, risk_factor=3.0, allow_short=True
        )
        w = layer(mu_batch, L_batch)
        sums = w.sum(dim=-1)
        np.testing.assert_allclose(sums.detach().numpy(), 1.0, atol=1e-3)
        # Short positions are now permitted — no non-negativity check


# ─────────────────────────────────────────────────────────────────────
# 2. PortfolioForecastNetwork
# ─────────────────────────────────────────────────────────────────────

class TestForecastNetwork:

    def test_forward_shapes(self):
        net = PortfolioForecastNetwork(N_ASSETS, N_FEATURES, hidden_dim=32)
        x = torch.randn(BATCH, N_ASSETS, N_FEATURES)
        mu, vol_scale = net(x)

        assert mu.shape == (BATCH, N_ASSETS)
        assert vol_scale.shape == (BATCH, N_ASSETS)
        # vol_scale must be positive (Softplus output)
        assert (vol_scale.detach().numpy() > 0).all()


# ─────────────────────────────────────────────────────────────────────
# 3. E2EPortfolioTrainer
# ─────────────────────────────────────────────────────────────────────

class TestTrainer:

    def test_training_runs_without_error(
        self, synthetic_features_dict, synthetic_returns_df
    ):
        trainer = E2EPortfolioTrainer(
            n_assets=N_ASSETS,
            n_features=N_FEATURES,
            risk_factor=3.0,
            loss_type="sharpe",
            hidden_dim=16,
            lr=1e-3,
        )
        history = trainer.train(
            synthetic_features_dict,
            synthetic_returns_df,
            n_epochs=3,       # Tiny run for speed
            batch_size=16,
            silent=True,
        )
        assert "train_loss" in history
        assert len(history["train_loss"]) == 3

    def test_predict_returns_valid_series(
        self, synthetic_features_dict, synthetic_returns_df
    ):
        trainer = E2EPortfolioTrainer(
            n_assets=N_ASSETS,
            n_features=N_FEATURES,
            risk_factor=3.0,
            loss_type="sharpe",
            hidden_dim=16,
        )
        # Quick train so the model has learned params
        trainer.train(
            synthetic_features_dict, synthetic_returns_df,
            n_epochs=2, batch_size=16, silent=True,
        )

        cov = synthetic_returns_df.cov()
        mu, w = trainer.predict(synthetic_features_dict, cov)

        assert isinstance(mu, pd.Series)
        assert isinstance(w, pd.Series)
        assert set(mu.index) == set(TICKERS)
        assert set(w.index) == set(TICKERS)
        np.testing.assert_allclose(w.sum(), 1.0, atol=1e-3)


# ─────────────────────────────────────────────────────────────────────
# 4. Cache save / load
# ─────────────────────────────────────────────────────────────────────

class TestCache:

    def test_save_and_load_roundtrip(self):
        trainer = E2EPortfolioTrainer(
            n_assets=N_ASSETS, n_features=N_FEATURES,
            risk_factor=3.0, hidden_dim=16,
        )

        tmp_dir = tempfile.mkdtemp()
        try:
            save_model(trainer, tmp_dir, TICKERS, 500, 0.04)
            assert load_model(trainer, tmp_dir, TICKERS, 500, 0.04)
            assert not load_model(trainer, tmp_dir, TICKERS, 999, 0.04)
        finally:
            shutil.rmtree(tmp_dir)


# ─────────────────────────────────────────────────────────────────────
# 5. Integration with solver.py warm-start
# ─────────────────────────────────────────────────────────────────────

class TestSolverIntegration:

    def test_warm_start_injection(self, synthetic_returns_df):
        """Verify that E2E weights survive the full build_and_optimize path."""
        from solver import build_and_optimize
        from core_types import PortfolioState

        returns_df = synthetic_returns_df
        bench_rets = returns_df["SPY"]
        state = PortfolioState.empty(list(returns_df.columns))
        state.total_capital = 100_000.0

        # Fake E2E warm-start
        e2e_ws = np.array([0.3, 0.2, 0.2, 0.2, 0.1])

        cfg = {
            "risk_free_rate": 0.04,
            "sector_limit": 0.5,
            "single_asset_min": 0.0,
            "single_asset_max": 0.5,
            "gross_leverage_cap": 1.0,
            "sector_map": {t: "Other" for t in TICKERS},
            "transaction_cost": 0.001,
            "bond_metadata": {},
            "default_adv_proxy": 50_000_000.0,
            "cvar_enabled": False,
            "garch_enabled": False,
            "dynamic_risk": False,
            "hmm_regime": False,
            "tax_enabled": False,
            "_e2e_warm_start": e2e_ws,
        }

        result = build_and_optimize(
            returns_df, bench_rets, 5, 3.0,
            state=state, cfg=cfg, model=1,
            silent=True,
        )

        assert result is not None
        # Weights should be valid
        w = result.weights
        assert abs(w.sum() - 1.0) < 0.05


if __name__ == "__main__":
    pytest.main([__file__, "-v"])