"""
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

# --- Import VectorCalibrator (Task 4) --- #
try:
    from gru_sac_predictor.src.calibrator_vector import VectorCalibrator
except ImportError:
    VectorCalibrator = None
# --- End Import --- #

# --- Helper Function for ECE --- #
def _calculate_ece(probs: np.ndarray, y_true: np.ndarray, n_bins: int = 10) -> float:
    """
    Calculates the Expected Calibration Error (ECE).
    
    Args:
        probs (np.ndarray): Predicted probabilities for the positive class (N,) or all classes (N, K).
        y_true (np.ndarray): True labels (0 or 1 for binary, or class index for multi-class).
        n_bins (int): Number of bins to divide probabilities into.

    Returns:
        float: The calculated ECE score.
    """
    if len(probs.shape) == 1: # Binary case
        p_max = probs
        y_pred_class = (probs > 0.5).astype(int)
        y_true_class = y_true
    elif len(probs.shape) == 2: # Multi-class case
        p_max = np.max(probs, axis=1)
        y_pred_class = np.argmax(probs, axis=1)
        # If y_true is one-hot, convert to class index
        if len(y_true.shape) == 2 and y_true.shape[1] > 1:
            y_true_class = np.argmax(y_true, axis=1)
        else:
            y_true_class = y_true # Assume already class index
    else:
        raise ValueError("probs array must be 1D or 2D")

    ece = 0.0
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    
    for i in range(n_bins):
        in_bin = (p_max > bin_boundaries[i]) & (p_max <= bin_boundaries[i+1])
        prop_in_bin = np.mean(in_bin)
        
        if prop_in_bin > 0:
            accuracy_in_bin = np.mean(y_pred_class[in_bin] == y_true_class[in_bin])
            avg_confidence_in_bin = np.mean(p_max[in_bin])
            ece += np.abs(accuracy_in_bin - avg_confidence_in_bin) * prop_in_bin
            
    return ece
# --- End ECE Helper --- #

# --- 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})"

# --- Vector Scaling Test (Task 4.4) --- #
@pytest.mark.skipif(VectorCalibrator is None, reason="VectorCalibrator not found")
def test_vector_scaling_calibration():
    """Check if Vector Scaling reduces ECE on sample multi-class data."""
    np.random.seed(123)
    n_samples = 5000
    num_classes = 3

    # Simulate slightly miscalibrated logits (e.g., too peaky or too flat)
    # True distribution is uniform-ish
    true_labels = np.random.randint(0, num_classes, n_samples)
    y_onehot = tf.keras.utils.to_categorical(true_labels, num_classes=num_classes)
    
    # Generate logits - make class 1 slightly more likely, and make logits "peaky"
    logits_raw = np.random.randn(n_samples, num_classes) * 0.5 # Base noise
    logits_raw[:, 1] += 0.5 # Bias towards class 1
    # Add systematic miscalibration (e.g., scale up logits -> overconfidence)
    logits_miscalibrated = logits_raw * 1.8 

    # Instantiate calibrator
    vector_cal = VectorCalibrator()

    # Calculate ECE before calibration
    probs_uncal = vector_cal._softmax(logits_miscalibrated)
    ece_before = _calculate_ece(probs_uncal, true_labels)
    
    # Fit vector scaling
    vector_cal.fit(logits_miscalibrated, y_onehot)
    assert vector_cal.W is not None and vector_cal.b is not None, "Vector scaling fit failed"

    # Calibrate probabilities
    probs_cal = vector_cal.calibrate(logits_miscalibrated)

    # Calculate ECE after calibration
    ece_after = _calculate_ece(probs_cal, true_labels)

    print(f"\nVector Scaling Test: ECE Before = {ece_before:.4f}, ECE After = {ece_after:.4f}")

    # Assert that ECE improved (decreased)
    # Allow for slight numerical noise, but expect significant improvement
    assert ece_after < ece_before * 0.7, f"ECE did not improve significantly after Vector Scaling (Before: {ece_before:.4f}, After: {ece_after:.4f})"
    # Assert ECE is reasonably low after calibration
    assert ece_after < 0.05, f"ECE after Vector Scaling ({ece_after:.4f}) is higher than expected (< 0.05)" 