# tests/test_wave1_domain_invariant.py
"""Tests for DomainInvariantSelector (wave1 29.1A)."""
from __future__ import annotations

import sys
from pathlib import Path
import pytest
import numpy as np

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))


def _make_features(n=30, d=12, seed=0):
    rng = np.random.RandomState(seed)
    return rng.randn(n, d).astype(np.float32)


def _make_chains(n=30, seed=0):
    """Minimal chain dicts for testing."""
    rng = np.random.RandomState(seed)
    chains = []
    for i in range(n):
        chains.append({
            "question": f"Test question {i}",
            "final_answer": "Test answer",
            "steps": [{"thought": "thinking", "action_type": "Search",
                        "action_arg": "query", "observation": "Some observation text."}],
            "correct": bool(rng.rand() > 0.5),
            "finished": True,
        })
    return chains


def test_import():
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector
    assert DomainInvariantSelector is not None


def test_instantiates_with_defaults():
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector
    dis = DomainInvariantSelector()
    assert dis.alpha == 1.0
    assert not dis.is_fitted


def test_fit_returns_self():
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector
    dis = DomainInvariantSelector()
    X_src = _make_features(30)
    y_src = np.array([0] * 15 + [1] * 15)
    X_tgt = _make_features(20, seed=1)
    result = dis.fit(X_src, y_src, X_tgt)
    assert result is dis
    assert dis.is_fitted


def test_transform_output_shape():
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector
    dis = DomainInvariantSelector()
    X_src = _make_features(30)
    y_src = np.array([0] * 15 + [1] * 15)
    X_tgt = _make_features(20, seed=1)
    dis.fit(X_src, y_src, X_tgt)
    X_new = _make_features(10, seed=2)
    X_out = dis.transform(X_new)
    assert X_out.shape == X_new.shape


def test_feature_weights_shape():
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector
    dis = DomainInvariantSelector()
    X_src = _make_features(30)
    y_src = np.array([0] * 15 + [1] * 15)
    X_tgt = _make_features(20, seed=1)
    dis.fit(X_src, y_src, X_tgt)
    assert dis.feature_weights_.shape == (12,)
    assert np.all(dis.feature_weights_ > 0), "all weights must be positive"


def test_predict_returns_float_in_range():
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector
    dis = DomainInvariantSelector()
    X_src = _make_features(30)
    y_src = np.array([0] * 15 + [1] * 15)
    X_tgt = _make_features(20, seed=1)
    dis.fit(X_src, y_src, X_tgt)
    X_test = _make_features(5, seed=3)
    preds = dis.predict_proba(X_test)
    assert preds.shape == (5,)
    assert np.all(preds >= 0) and np.all(preds <= 1)


def test_high_alpha_downweights_shifted_features():
    """With high alpha, features with large distribution shift get lower weight."""
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector
    rng = np.random.RandomState(42)
    # Feature 0 has large shift; feature 1 has zero shift
    X_src = rng.randn(30, 2).astype(np.float32)
    X_tgt = X_src.copy()
    X_tgt[:, 0] += 5.0  # large shift in feature 0
    y_src = np.array([0] * 15 + [1] * 15)
    dis = DomainInvariantSelector(alpha=10.0)
    dis.fit(X_src, y_src, X_tgt)
    assert dis.feature_weights_[0] < dis.feature_weights_[1], (
        "Shifted feature should have lower weight"
    )
