import numpy as np
import pytest

from hunt import models as M
from hunt.rng import Mulberry32


def test_auc_matches_sklearn_with_ties():
    from sklearn.metrics import roc_auc_score
    rng = np.random.default_rng(0)
    for _ in range(5):
        y = rng.integers(0, 2, 200)
        s = np.round(rng.normal(size=200) + 0.8 * y, 1)      # rounding makes ties
        assert abs(M.auc(s, y) - roc_auc_score(y, s)) < 1e-12
    assert M.auc([1, 2, 3], [0, 0, 0]) == 0.5


def test_logreg_separates_easy_data_and_is_deterministic():
    rng = np.random.default_rng(1)
    X = rng.normal(size=(300, 3))
    y = (X[:, 0] + 0.5 * X[:, 1] > 0).astype(int)
    m1 = M.logreg_fit(X, y)
    m2 = M.logreg_fit(X, y)
    assert np.array_equal(m1["w"], m2["w"])
    assert M.accuracy(M.logreg_predict(m1, X), y) > 0.95
    assert abs(m1["w"][0]) > abs(m1["w"][2])
    # constant columns get zero weight and do not break anything
    Xc = np.column_stack([X, np.ones(300)])
    mc = M.logreg_fit(Xc, y)
    assert mc["w"][3] == 0.0
    assert mc["lr"] <= 10


def test_logreg_step_is_stable_on_correlated_features():
    rng = np.random.default_rng(2)
    base = rng.normal(size=(120, 1))
    X = base + 0.01 * rng.normal(size=(120, 40))          # 40 nearly identical columns
    y = (base[:, 0] > 0).astype(int)
    m = M.logreg_fit(X, y)
    assert np.all(np.isfinite(m["w"]))
    assert M.accuracy(M.logreg_predict(m, X), y) > 0.9


def test_knn1_ties_and_distances():
    Xtr = np.array([[0.0, 0.0], [1.0, 1.0], [0.0, 0.0]])
    ytr = np.array([1, 0, 0])
    pred, dist = M.knn1(Xtr, ytr, np.array([[0.0, 0.0], [1.0, 1.0]]))
    assert pred.tolist() == [1, 0]                           # exact tie -> lowest index wins
    assert dist[0] == 0.0 and dist[1] == 0.0


def test_stump_finds_threshold_and_reports_per_feature():
    X = np.array([[1, 9], [2, 8], [3, 7], [4, 1], [5, 2], [6, 3]], dtype=float)
    y = np.array([0, 0, 0, 1, 1, 1])
    s = M.stump_fit(X, y)
    assert s["acc"] == 1.0 and s["feature"] == 0 and s["threshold"] == 3.5 and s["above"] == 1
    assert s["per_feature"] == [1.0, 1.0]
    assert M.stump_predict(s, X).tolist() == y.tolist()
    # direction: predicts 0 above the threshold when that is what fits
    s2 = M.stump_fit(X[:, [1]], y)
    assert s2["above"] == 0 and s2["acc"] == 1.0


def test_metrics():
    pred = np.array([0, 0, 0, 0, 1])
    y = np.array([0, 0, 1, 1, 1])
    assert M.accuracy(pred, y) == pytest.approx(0.6)
    assert M.confusion(pred, y) == {"tp": 1, "fp": 0, "fn": 2, "tn": 2}
    assert M.balanced_accuracy(pred, y) == pytest.approx(0.5 * (1 / 3 + 1))
    assert M.recall(pred, y) == pytest.approx(1 / 3)


def test_mulberry32_known_values():
    # first draws of mulberry32(7), as produced by the usual JavaScript one-liner
    r = Mulberry32(7)
    xs = [r.random() for _ in range(3)]
    assert all(0 <= x < 1 for x in xs)
    assert len(set(xs)) == 3
    assert Mulberry32(7).random() == xs[0]
    items = Mulberry32(1).shuffle(list(range(10)))
    assert sorted(items) == list(range(10)) and items != list(range(10))