"""Unit tests for ucml.prediction helpers on synthetic arrays (no data access)."""
import numpy as np

from ucml import prediction as P


def test_make_targets_rules():
    day = 86400.0
    entry = np.zeros(5)
    first_cut = np.array([10 * day, 40 * day, np.nan, 5 * day, np.nan])
    duration = np.array([50.0, 50.0, 20.0, 12.0, 45.0])
    exit_conf = np.array([True, False, True, True, False])
    y = P.make_targets(entry, first_cut, duration, exit_conf, horizon=30)
    assert y["y_markdown30"].tolist() == [1, 0, 0, 1, 0]        # cut at 40 d is beyond the horizon
    assert y["y_exit30"].tolist() == [0, 0, 1, 1, 0]            # exit at 12 d with cut at 5 d: both events
    cr = y["y_markdown30cr"]
    assert np.isnan(cr[2]) and cr[3] == 1 and cr[0] == 1 and cr[1] == 0   # exit-before-cut removed only when no cut


def test_calibration_table_and_ece():
    y = np.array([0, 0, 1, 1, 1, 0, 1, 1])
    p = np.array([0.05, 0.15, 0.15, 0.95, 0.85, 0.85, 0.95, 0.95])
    t = P.calibration_table(y, p, n_bins=10)
    assert len(t) == 10 and t["n"].sum() == 8
    b1 = t.iloc[1]; assert b1["n"] == 2 and abs(b1["frac_pos"] - 0.5) < 1e-12 and abs(b1["mean_pred"] - 0.15) < 1e-12
    e = P.expected_calibration_error(y, p, n_bins=10)
    # bins: [0.05]->|0-0.05|, [0.15,0.15]->|0.5-0.15|, [0.85,0.85]->|0.5-0.85|, [0.95x3]->|1-0.95|
    assert abs(e - (1 * 0.05 + 2 * 0.35 + 2 * 0.35 + 3 * 0.05) / 8) < 1e-12
    perfect = P.expected_calibration_error(np.array([0, 1, 0, 1]), np.array([0.5, 0.5, 0.5, 0.5]))
    assert abs(perfect) < 1e-12


def test_bootstrap_metrics_keys_and_ci_order():
    rng = np.random.default_rng(0)
    p = rng.uniform(size=400); y = (rng.uniform(size=400) < p).astype(float)
    m = P.bootstrap_metrics(y, p, n_boot=50, seed=1)
    for k in ("log_loss", "brier", "auc"):
        lo, hi = m[f"{k}_ci"]; assert lo <= m[k] <= hi or abs(m[k] - lo) < 0.05 or abs(hi - m[k]) < 0.05
    assert m["n_test"] == 400 and 0 < m["base_rate"] < 1 and m["auc"] > 0.6
    g = np.repeat(["A", "B"], 200)
    mg = P.bootstrap_metrics(y, p, groups=g, n_boot=20, seed=1)
    assert abs(mg["log_loss"] - np.mean([P._metrics(y[g == k], p[g == k])["log_loss"] for k in ("A", "B")])) < 1e-12


def test_target_encode_no_self_leak_and_prior_for_unseen():
    g_tr = np.array(["d1"] * 10 + ["d2"] * 10)
    y_tr = np.array([1.0] * 10 + [0.0] * 10)
    enc_tr, enc_te = P.target_encode(g_tr, y_tr, np.array(["d1", "d2", "d9"]), m=2.0, k=5, seed=0)
    prior = 0.5
    assert abs(enc_te[2] - prior) < 1e-12                       # unseen dealer -> training prior
    assert abs(enc_te[0] - (10 + 2 * prior) / 12) < 1e-12 and abs(enc_te[1] - (0 + 2 * prior) / 12) < 1e-12
    # training rows: own label excluded, so d1 encodings are strictly below the full-training value
    assert np.all(enc_tr[:10] < enc_te[0]) and np.all(enc_tr[10:] > enc_te[1])
    # NaN labels (competing-risk target) are ignored, not propagated
    y_nan = y_tr.copy(); y_nan[0] = np.nan
    a, b = P.target_encode(g_tr, y_nan, np.array(["d1"]), m=2.0)
    assert np.isfinite(a).all() and np.isfinite(b).all()