"""
On-device verification for the cross-compiled scikit-learn wheel.

Run after installing:
    pip install scikit_learn-1.7.1-cp312-cp312-android_24_x86_64.whl

Usage:
    python Test_ScikitLearn.py [--quick]

Exit code 0 = everything required PASSed.
Requires: numpy, scipy, joblib, threadpoolctl at runtime.

Generated by RIMI
"""
import sys

RESULTS = []


def test(name, fn):
    try:
        fn()
        RESULTS.append((name, "PASS", None))
    except NotImplementedError as exc:
        RESULTS.append((name, "SKIP", str(exc)))
    except Exception as exc:
        RESULTS.append((name, "FAIL", "%s: %s" % (type(exc).__name__, exc)))
        print("    ! %s -> %s: %s" % (name, type(exc).__name__, exc))


def section(title):
    print("=" * 60)
    print(title)
    print("=" * 60)


# ---------------------------------------------------------------------------
# 1. import / version
# ---------------------------------------------------------------------------
def import_sklearn():
    import sklearn
    print("    sklearn", sklearn.__version__)
    assert hasattr(sklearn, "__version__")
    assert hasattr(sklearn, "show_versions")

def check_c_extension():
    import sklearn
    # Check that at least one Cython extension loads (tree, metrics, etc.)
    from sklearn.tree import _tree
    assert hasattr(_tree, "Tree")


# ---------------------------------------------------------------------------
# 2. datasets
# ---------------------------------------------------------------------------
def load_iris():
    from sklearn.datasets import load_iris
    X, y = load_iris(return_X_y=True)
    assert X.shape == (150, 4)
    assert y.shape == (150,)

def load_digits():
    from sklearn.datasets import load_digits
    X, y = load_digits(return_X_y=True)
    assert X.shape[0] == 1797


# ---------------------------------------------------------------------------
# 3. preprocessing
# ---------------------------------------------------------------------------
def scaler_standard():
    from sklearn.preprocessing import StandardScaler
    import numpy as np
    X = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float64)
    scaler = StandardScaler()
    Xt = scaler.fit_transform(X)
    assert Xt.shape == X.shape
    assert abs(Xt.mean()) < 1e-6

def scaler_minmax():
    from sklearn.preprocessing import MinMaxScaler
    import numpy as np
    X = np.array([[1, 2], [3, 4]], dtype=np.float64)
    scaler = MinMaxScaler()
    Xt = scaler.fit_transform(X)
    assert Xt.min() >= 0 and Xt.max() <= 1


# ---------------------------------------------------------------------------
# 4. decomposition
# ---------------------------------------------------------------------------
def pca_test():
    from sklearn.decomposition import PCA
    import numpy as np
    rng = np.random.default_rng(42)
    X = rng.random((50, 10))
    pca = PCA(n_components=2)
    Xt = pca.fit_transform(X)
    assert Xt.shape == (50, 2)


# ---------------------------------------------------------------------------
# 5. cluster
# ---------------------------------------------------------------------------
def kmeans_test():
    from sklearn.cluster import KMeans
    import numpy as np
    rng = np.random.default_rng(42)
    X = rng.random((30, 2))
    km = KMeans(n_clusters=3, n_init=10, random_state=42)
    labels = km.fit_predict(X)
    assert labels.shape == (30,)
    assert len(set(labels)) == 3


# ---------------------------------------------------------------------------
# 6. linear model
# ---------------------------------------------------------------------------
def logistic_regression():
    from sklearn.linear_model import LogisticRegression
    from sklearn.datasets import load_iris
    X, y = load_iris(return_X_y=True)
    # binary iris (first 100 samples, 2 classes)
    Xb, yb = X[:100], y[:100]
    clf = LogisticRegression(max_iter=200)
    clf.fit(Xb, yb)
    pred = clf.predict(Xb[:5])
    assert pred.shape == (5,)

def linear_regression():
    from sklearn.linear_model import LinearRegression
    import numpy as np
    X = np.array([[1], [2], [3], [4]], dtype=np.float64)
    y = np.array([2, 4, 6, 8], dtype=np.float64)
    reg = LinearRegression()
    reg.fit(X, y)
    pred = reg.predict([[5]])
    assert abs(pred[0] - 10) < 1e-3


# ---------------------------------------------------------------------------
# 7. ensemble
# ---------------------------------------------------------------------------
def random_forest():
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.datasets import load_iris
    X, y = load_iris(return_X_y=True)
    clf = RandomForestClassifier(n_estimators=10, random_state=42)
    clf.fit(X[:100], y[:100])
    pred = clf.predict(X[:5])
    assert pred.shape == (5,)


# ---------------------------------------------------------------------------
# 8. metrics
# ---------------------------------------------------------------------------
def metrics_test():
    from sklearn.metrics import accuracy_score, mean_squared_error
    import numpy as np
    y_true = np.array([0, 1, 1, 0])
    y_pred = np.array([0, 1, 0, 0])
    acc = accuracy_score(y_true, y_pred)
    assert 0 <= acc <= 1
    mse = mean_squared_error([1, 2, 3], [1, 2, 3])
    assert mse == 0


# ---------------------------------------------------------------------------
# 9. model_selection
# ---------------------------------------------------------------------------
def train_test_split():
    from sklearn.model_selection import train_test_split
    import numpy as np
    X = np.random.rand(20, 4)
    y = np.random.randint(0, 2, 20)
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.25, random_state=42)
    assert Xtr.shape[0] == 15 and Xte.shape[0] == 5

def cross_val():
    from sklearn.model_selection import cross_val_score
    from sklearn.linear_model import LogisticRegression
    from sklearn.datasets import load_iris
    X, y = load_iris(return_X_y=True)
    clf = LogisticRegression(max_iter=200)
    scores = cross_val_score(clf, X[:100], y[:100], cv=3)
    assert len(scores) == 3


# ---------------------------------------------------------------------------
def main():
    quick = "--quick" in sys.argv
    section("1. import / version")
    test("import sklearn", import_sklearn)
    test("C extension _tree", check_c_extension)

    section("2. datasets")
    test("load_iris", load_iris)
    test("load_digits", load_digits)

    section("3. preprocessing")
    test("StandardScaler", scaler_standard)
    test("MinMaxScaler", scaler_minmax)

    section("4. decomposition")
    test("PCA", pca_test)

    section("5. cluster")
    test("KMeans", kmeans_test)

    section("6. linear_model")
    test("LogisticRegression", logistic_regression)
    test("LinearRegression", linear_regression)

    section("7. ensemble")
    test("RandomForest", random_forest)

    section("8. metrics")
    test("metrics", metrics_test)

    section("9. model_selection")
    test("train_test_split", train_test_split)
    test("cross_val_score", cross_val)

    print()
    print("=" * 60)
    print("SUMMARY")
    print("=" * 60)
    fails = 0
    skips = 0
    for name, status, why in RESULTS:
        mark = "  OK" if status == "PASS" else (" SKIP" if status == "SKIP" else "FAIL")
        print("%s %s" % (mark, name))
        if why:
            print("        -> %s" % why)
        if status == "FAIL":
            fails += 1
        elif status == "SKIP":
            skips += 1
    print()
    passed = len(RESULTS) - fails - skips
    print("passed=%d skipped=%d failed=%d" % (passed, skips, fails))
    if fails:
        print("RESULT: FAILED")
    elif skips and not quick:
        print("RESULT: PASSED (with informational skips)")
    else:
        print("RESULT: PASSED")
    sys.exit(1 if fails else 0)


if __name__ == "__main__":
    main()