"""Failing tests for HEPA backbone in train.py.

Run: cd ~/dev/AI/jepa-fx-risk && .venv/bin/python -m pytest tests/test_hepa.py -v
These tests define what the new backbone must satisfy BEFORE implementation.
"""
import math
import torch
import torch.nn as nn
import pytest

# ── Tests import the classes from train.py ────────────────────────────────────
# They will fail until train.py implements: CausalEncoder, HorizonPredictor, vicreg_loss


def _import():
    import importlib.util, sys
    spec = importlib.util.spec_from_file_location("train", "train.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


@pytest.fixture(scope="module")
def train_mod():
    return _import()


# 1. CausalEncoder exists and has correct output shape
def test_causal_encoder_shape(train_mod):
    enc = train_mod.CausalEncoder(n_channels=2, patch_len=10, d_model=32, n_heads=4, depth=1)
    x = torch.randn(4, 60, 2)
    tokens = enc(x)          # should return all tokens (B, N, D) for JEPA pretraining
    assert tokens.shape == (4, 6, 32), f"expected (4, 6, 32), got {tokens.shape}"


# 2. CausalEncoder is actually causal: earlier token outputs don't change when later inputs change
def test_causal_masking(train_mod):
    enc = train_mod.CausalEncoder(n_channels=2, patch_len=10, d_model=32, n_heads=4, depth=2)
    enc.eval()
    torch.manual_seed(0)
    x = torch.randn(1, 60, 2)
    x_perturbed = x.clone()
    # non-uniform noise (constant shift absorbed by per-patch LayerNorm; variance change is not)
    torch.manual_seed(99)
    x_perturbed[:, 30:, :] += torch.randn_like(x[:, 30:, :]) * 5.0

    with torch.no_grad():
        h1 = enc(x)
        h2 = enc(x_perturbed)

    # First 3 tokens must be identical (causal — don't see future patches)
    assert torch.allclose(h1[:, :3, :], h2[:, :3, :], atol=1e-5), \
        "causal masking broken: early tokens change when later input changes"
    # Last token should differ (it can see the perturbed patches)
    assert not torch.allclose(h1[:, -1, :], h2[:, -1, :], atol=1e-5), \
        "last token should differ when later input changes"


# 3. HorizonPredictor exists, takes (h, delta_t_float) → same shape as h
def test_horizon_predictor_shape(train_mod):
    pred = train_mod.HorizonPredictor(d_model=32)
    h = torch.randn(4, 32)
    dt = torch.tensor([1.0, 2.0, 3.0, 1.0])
    out = pred(h, dt)
    assert out.shape == (4, 32), f"expected (4, 32), got {out.shape}"


# 4. vicreg_loss is a scalar and backward doesn't error
def test_vicreg_loss_backward(train_mod):
    h_pred = torch.randn(8, 32, requires_grad=True)
    h_target = torch.randn(8, 32)
    loss = train_mod.vicreg_loss(h_pred, h_target, alpha=0.1)
    assert loss.shape == (), f"expected scalar, got {loss.shape}"
    loss.backward()
    assert h_pred.grad is not None


# 5. Full JEPA step: encode context, predict future, compute loss, backward
def test_jepa_step_end_to_end(train_mod):
    enc = train_mod.CausalEncoder(n_channels=2, patch_len=10, d_model=32, n_heads=4, depth=1)
    pred = train_mod.HorizonPredictor(d_model=32)
    opt = torch.optim.SGD(list(enc.parameters()) + list(pred.parameters()), lr=1e-3)

    x = torch.randn(4, 60, 2)
    tokens = enc(x)          # (4, 6, 32)
    c, dt = 2, 2             # context position 2, horizon 2
    h_ctx = tokens[:, c, :]
    h_tgt = tokens[:, c + dt, :].detach()
    h_hat = pred(h_ctx, torch.full((4,), float(dt)))
    loss = train_mod.vicreg_loss(h_hat, h_tgt, alpha=0.1)
    opt.zero_grad(); loss.backward(); opt.step()
    assert loss.item() < 100, "loss exploded"


# 6. build() returns year-based OOS split (2022-2023); hourly gives many more windows
def test_build_year_split(train_mod):
    (Xtr, ytr), (Xte, yte) = train_mod.build()
    assert Xtr.shape[1] == train_mod.WINDOW
    assert Xte.shape[1] == train_mod.WINDOW
    assert len(Xtr) > 0 and len(Xte) > 0
    # OOS: daily ≈ 600; hourly ≈ 17,000 (2 years × ~8,500 trading hours/year)
    assert len(Xte) > 400, f"OOS too small: {len(Xte)}"


# 7. hourly build gives > 10× more training windows than daily
def test_build_hourly_more_windows(train_mod):
    import os
    if not os.path.exists("data/processed/eurusd_hourly.parquet"):
        pytest.skip("eurusd_hourly.parquet not present — run data:prepare:hourly first")
    (Xtr, _), _ = train_mod.build()
    # Daily had ~877 train windows; hourly with 2008-2021 should have > 50,000
    assert len(Xtr) > 10_000, f"expected >10k hourly train windows, got {len(Xtr)}"