"""Smoke tests for ai_toolkit.ml."""

import numpy as np
import pandas as pd
import pytest
from sklearn.linear_model import LogisticRegression

from ai_toolkit.ml import (
    build_classification_pipeline,
    build_results_table,
    encode_labels,
    error_analysis,
    full_classification_report,
    scale_features,
    train_and_evaluate,
)


def test_scale_features():
    X_tr = np.random.randn(50, 3)
    X_te = np.random.randn(10, 3)
    scaler, Z_tr, Z_te = scale_features(X_tr, X_te)
    assert Z_tr.shape == X_tr.shape
    assert Z_te.shape == X_te.shape


def test_encode_labels():
    enc, y = encode_labels(["a", "b", "a"])
    assert len(y) == 3
    assert set(y) <= {0, 1}


def test_build_classification_pipeline():
    pipe = build_classification_pipeline(classifier=LogisticRegression(max_iter=200))
    assert pipe is not None
    assert hasattr(pipe, "fit")


def test_train_and_evaluate():
    X = np.random.randn(80, 5)
    y = np.array([0, 1] * 40)
    X_tr, X_te = X[:60], X[60:]
    y_tr, y_te = y[:60], y[60:]
    pipe = LogisticRegression(max_iter=200)
    metrics = train_and_evaluate(pipe, X_tr, X_te, y_tr, y_te, task="classification")
    assert "accuracy" in metrics
    assert "f1_weighted" in metrics


def test_build_results_table():
    tbl = build_results_table([{"accuracy": 0.8, "f1": 0.79}, {"accuracy": 0.85, "f1": 0.84}])
    assert isinstance(tbl, pd.DataFrame)
    assert len(tbl) == 2


def test_full_classification_report():
    model = LogisticRegression(max_iter=200)
    X = np.random.randn(40, 4)
    y = np.array([0, 1] * 20)
    model.fit(X, y)
    report = full_classification_report(model, X, y, output_dict=False)
    assert isinstance(report, str)
    assert "precision" in report.lower() or "accuracy" in report.lower()


def test_error_analysis():
    model = LogisticRegression(max_iter=200)
    X = np.random.randn(30, 4)
    y = np.array([0, 1] * 15)
    model.fit(X, y)
    df = error_analysis(model, X, y, top_n=5)
    assert isinstance(df, pd.DataFrame)
    assert "true_label" in df.columns
    assert "predicted_label" in df.columns
