"""
Tests for probability calibration (Sec 6 of revisions.txt).
"""
import pytest
import numpy as np
from scipy.stats import binomtest
from scipy.special import logit, expit
import os

# Try to import the modules; skip tests if not found (e.g., path issues)
try:
    from gru_sac_predictor.src import calibrate
except ImportError:
    calibrate = None

# --- Fixtures ---
@pytest.fixture(scope="module")
def calibration_data():
    """
    Generate sample raw probabilities and true outcomes.
    Simulates an overconfident model (T_implied < 1) where true probability drifts.
    """
    np.random.seed(42)
    n_samples = 2500
    # Simulate drifting true probability centered around 0.5
    drift = 0.05 * np.sin(np.linspace(0, 3 * np.pi, n_samples))
    true_prob = np.clip(0.5 + drift + np.random.randn(n_samples) * 0.05, 0.05, 0.95)
    # Simulate overconfidence (implied T ~ 0.7)
    raw_logits = logit(true_prob) / 0.7
    p_raw = expit(raw_logits)
    # Generate true outcomes
    y_true = (np.random.rand(n_samples) < true_prob).astype(int)
    return p_raw, y_true

# --- Tests ---
@pytest.mark.skipif(calibrate is None, reason="Module gru_sac_predictor.src.calibrate not found")
def test_optimise_temperature(calibration_data):
    """Check if optimise_temperature runs and returns a plausible value."""
    p_raw, y_true = calibration_data
    optimal_T = calibrate.optimise_temperature(p_raw, y_true)
    print(f"\nOptimised T: {optimal_T:.4f}")
    # Expect T > 0. A T near 0.7 would undo the simulated effect.
    assert optimal_T > 0.1 and optimal_T < 5.0, "Optimised temperature seems out of expected range."

@pytest.mark.skipif(calibrate is None, reason="Module gru_sac_predictor.src.calibrate not found")
def test_calibration_hit_rate_threshold(calibration_data):
    """
    Verify that the lower 95% CI of the hit-rate for non-zero calibrated
    signals is >= 0.55 (using the module's EDGE_THR).
    """
    p_raw, y_true = calibration_data
    optimal_T = calibrate.optimise_temperature(p_raw, y_true)
    p_cal = calibrate.calibrate(p_raw, optimal_T)
    action_signals = calibrate.action_signal(p_cal)

    # Filter for non-zero signals
    non_zero_idx = action_signals != 0
    if not np.any(non_zero_idx):
        pytest.fail("No non-zero action signals generated for hit-rate test.")

    signals_taken = action_signals[non_zero_idx]
    actual_direction = y_true[non_zero_idx]

    # Hit: signal matches actual direction (1 vs 1, -1 vs 0)
    hits = np.sum((signals_taken == 1) & (actual_direction == 1)) + \
           np.sum((signals_taken == -1) & (actual_direction == 0))
    total_trades = len(signals_taken)

    if total_trades < 30:
        pytest.skip(f"Insufficient non-zero signals ({total_trades}) for reliable CI.")

    # Calculate 95% lower CI using binomial test
    try:
        # Ensure hits is integer
        hits = int(hits)
        result = binomtest(hits, total_trades, p=0.5, alternative='greater')
        lower_ci = result.proportion_ci(confidence_level=0.95).low
    except Exception as e:
        pytest.fail(f"Binomial test failed: {e}")

    hit_rate = hits / total_trades
    required_threshold = calibrate.EDGE_THR # Use threshold from module

    print(f"\nCalibration Test: EDGE_THR={required_threshold:.3f}")
    print(f" Trades={total_trades}, Hits={hits}, Hit Rate={hit_rate:.4f}")
    print(f" 95% Lower CI: {lower_ci:.4f}")

    assert lower_ci >= required_threshold, \
        f"Hit rate lower CI ({lower_ci:.4f}) is below module threshold ({required_threshold:.3f})" 