"""
Tests for hedged_rf package.
"""

import numpy as np
import pytest
from sklearn.ensemble import RandomForestRegressor

from hedged_rf import (
    fit_hrf,
    extract_tree_predictions,
    estimate_mu_ewma,
    estimate_mu_complete,
    estimate_sigma_ewma,
    estimate_sigma_complete,
    create_cvc_target,
    solve_hrf_optimization,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture(scope="module")
def residual_matrix():
    """Small deterministic residual matrix for unit tests."""
    rng = np.random.default_rng(42)
    T, N = 120, 50
    return rng.standard_normal((T, N))


@pytest.fixture(scope="module")
def fitted_rf():
    """Small fitted RandomForestRegressor for integration tests."""
    rng = np.random.default_rng(0)
    X = rng.standard_normal((200, 10))
    y = X[:, 0] * 2 + rng.standard_normal(200) * 0.5
    rf = RandomForestRegressor(n_estimators=50, random_state=42, max_depth=4)
    rf.fit(X, y)
    return rf, X, y


# ---------------------------------------------------------------------------
# Unit tests — EWMA estimation
# ---------------------------------------------------------------------------

class TestEWMAMean:
    def test_output_shape(self, residual_matrix):
        T, N = residual_matrix.shape
        mu = estimate_mu_ewma(residual_matrix)
        assert mu.shape == (N,)

    def test_weights_sum_to_one(self):
        """EWMA weights should be normalized."""
        T = 100
        t = np.arange(1, T + 1)
        lam = 0.15
        w = lam * (1 - lam) ** (T - t)
        w /= w.sum()
        assert abs(w.sum() - 1.0) < 1e-10

    def test_recent_bias(self, residual_matrix):
        """With high lambda, EWMA should be closer to recent rows."""
        R = residual_matrix
        mu_high = estimate_mu_ewma(R, lambda_param=0.8)
        mu_low = estimate_mu_ewma(R, lambda_param=0.01)
        # High lambda → closer to last row
        diff_high = np.abs(mu_high - R[-1]).mean()
        diff_low = np.abs(mu_low - R[-1]).mean()
        assert diff_high < diff_low


class TestEWMASigma:
    def test_output_shape(self, residual_matrix):
        T, N = residual_matrix.shape
        S = estimate_sigma_ewma(residual_matrix)
        assert S.shape == (N, N)

    def test_symmetry(self, residual_matrix):
        S = estimate_sigma_ewma(residual_matrix)
        np.testing.assert_allclose(S, S.T, atol=1e-12)

    def test_positive_diagonal(self, residual_matrix):
        S = estimate_sigma_ewma(residual_matrix)
        assert np.all(np.diag(S) > 0)


# ---------------------------------------------------------------------------
# Unit tests — CVC target
# ---------------------------------------------------------------------------

class TestCVCTarget:
    def test_shape(self, residual_matrix):
        S = estimate_sigma_ewma(residual_matrix)
        F = create_cvc_target(S)
        assert F.shape == S.shape

    def test_constant_diagonal(self, residual_matrix):
        S = estimate_sigma_ewma(residual_matrix)
        F = create_cvc_target(S)
        diag = np.diag(F)
        assert np.allclose(diag, diag[0])

    def test_constant_off_diagonal(self, residual_matrix):
        S = estimate_sigma_ewma(residual_matrix)
        F = create_cvc_target(S)
        N = F.shape[0]
        mask = ~np.eye(N, dtype=bool)
        off_diag = F[mask]
        assert np.allclose(off_diag, off_diag[0])


# ---------------------------------------------------------------------------
# Unit tests — Shrinkage
# ---------------------------------------------------------------------------

class TestShrinkage:
    def test_mu_complete_shape(self, residual_matrix):
        mu = estimate_mu_complete(residual_matrix, use_fast=False)
        assert mu.shape == (residual_matrix.shape[1],)

    def test_sigma_complete_shape(self, residual_matrix):
        S = estimate_sigma_complete(residual_matrix, use_fast=False)
        assert S.shape == (residual_matrix.shape[1], residual_matrix.shape[1])

    def test_sigma_complete_symmetry(self, residual_matrix):
        S = estimate_sigma_complete(residual_matrix, use_fast=False)
        np.testing.assert_allclose(S, S.T, atol=1e-10)


# ---------------------------------------------------------------------------
# Unit tests — Optimization
# ---------------------------------------------------------------------------

class TestOptimization:
    def test_weights_sum_to_one(self, residual_matrix):
        mu = estimate_mu_complete(residual_matrix, use_fast=False)
        Sigma = estimate_sigma_complete(residual_matrix, use_fast=False)
        w = solve_hrf_optimization(mu, Sigma, kappa=2.0)
        assert abs(w.sum() - 1.0) < 1e-4

    def test_l1_constraint_respected(self, residual_matrix):
        mu = estimate_mu_complete(residual_matrix, use_fast=False)
        Sigma = estimate_sigma_complete(residual_matrix, use_fast=False)
        kappa = 2.0
        w = solve_hrf_optimization(mu, Sigma, kappa=kappa)
        assert np.sum(np.abs(w)) <= kappa + 1e-4

    def test_non_negative_with_kappa_one(self, residual_matrix):
        mu = estimate_mu_complete(residual_matrix, use_fast=False)
        Sigma = estimate_sigma_complete(residual_matrix, use_fast=False)
        w = solve_hrf_optimization(mu, Sigma, kappa=1.0)
        assert np.all(w >= -1e-5)


# ---------------------------------------------------------------------------
# Integration tests — fit_hrf
# ---------------------------------------------------------------------------

class TestFitHRF:
    def test_returns_array(self, residual_matrix):
        w = fit_hrf(residual_matrix, use_fast=False)
        assert isinstance(w, np.ndarray)

    def test_weights_shape(self, residual_matrix):
        w = fit_hrf(residual_matrix, use_fast=False)
        N = residual_matrix.shape[1]
        assert w.shape == (N,)

    def test_weights_sum_to_one(self, residual_matrix):
        w = fit_hrf(residual_matrix, use_fast=False)
        assert abs(w.sum() - 1.0) < 1e-4

    def test_verbose_runs(self, residual_matrix, capsys):
        fit_hrf(residual_matrix, verbose=True, use_fast=False)
        captured = capsys.readouterr()
        assert "HRF" in captured.out


class TestExtractTreePredictions:
    def test_shape(self, fitted_rf):
        rf, X, y = fitted_rf
        R = extract_tree_predictions(rf, X, y)
        assert R.shape == (len(y), len(rf.estimators_))

    def test_residual_correctness(self, fitted_rf):
        rf, X, y = fitted_rf
        R = extract_tree_predictions(rf, X, y)
        tree0_pred = rf.estimators_[0].predict(X)
        np.testing.assert_allclose(R[:, 0], y - tree0_pred)


# ---------------------------------------------------------------------------
# Edge-case tests
# ---------------------------------------------------------------------------

class TestEdgeCases:
    def test_single_tree(self):
        rng = np.random.default_rng(1)
        R = rng.standard_normal((100, 1))
        w = fit_hrf(R, use_fast=False)
        assert abs(w[0] - 1.0) < 1e-4

    def test_lambda_near_zero(self):
        rng = np.random.default_rng(2)
        R = rng.standard_normal((100, 20))
        w = fit_hrf(R, lambda_param=0.01, use_fast=False)
        assert abs(w.sum() - 1.0) < 1e-4

    def test_lambda_near_one(self):
        rng = np.random.default_rng(3)
        R = rng.standard_normal((100, 20))
        w = fit_hrf(R, lambda_param=0.99, use_fast=False)
        assert abs(w.sum() - 1.0) < 1e-4
