import numpy as np
import pytest
from sklearn.ensemble import RandomForestClassifier
from sklearn.gaussian_process import GaussianProcessRegressor

from lmpm.training.dataset import (
    CATEGORICAL_FEATURE_FIELDS,
    NUMERIC_FEATURE_FIELDS,
    TrainingDataset,
    TrainingRecord,
)
from lmpm.training.models import (
    build_estimator,
    load_training_bundle,
    save_training_bundle,
    train_models,
)
from lmpm.training.routing import default_model_routes


def make_record(number: int, material_name: str) -> TrainingRecord:
    features = {
        field: float(number + index + 1)
        for index, field in enumerate(NUMERIC_FEATURE_FIELDS)
    }
    features.update(
        {
            CATEGORICAL_FEATURE_FIELDS[0]: "bidirectional",
            CATEGORICAL_FEATURE_FIELDS[1]: "20x20",
        }
    )
    return TrainingRecord(
        metadata={
            "experiment_id": f"LAS-2026-{number:04d}",
            "material_name": material_name,
        },
        features=features,
        targets={
            "is_cut_through": number % 2 == 0,
            "carbonized_edge_width": number / 10,
            "etching_depth": number * 10.0,
            "is_fire_smolder": False,
            "pattern_clarity_score": float(number),
            "presentation_balance_score": float(10 - number),
        },
    )


def training_dataset() -> TrainingDataset:
    return TrainingDataset(
        records=(
            make_record(1, "Acrylic"),
            make_record(2, "Acrylic"),
            make_record(3, "Basswood"),
            make_record(4, "Basswood"),
        )
    )


def test_model_factory_builds_the_routed_estimator_types():
    routes = {route.target_name: route for route in default_model_routes()}

    assert isinstance(
        build_estimator(routes["etching_depth"]), GaussianProcessRegressor
    )
    assert isinstance(build_estimator(routes["is_cut_through"]), RandomForestClassifier)


def test_training_runs_leakage_safe_lomo_and_persists_bundle(tmp_path):
    bundle = train_models(training_dataset())

    evaluation = bundle.evaluations["etching_depth"]
    assert [fold.held_out_material for fold in evaluation.folds] == [
        "Acrylic",
        "Basswood",
    ]
    assert evaluation.mean_metrics["mae"] is not None
    assert bundle.models["etching_depth"].route.model_family == (
        "gaussian_process_regressor"
    )

    features = bundle.models["etching_depth"].preprocessor.transform(training_dataset())
    predictions = bundle.models["etching_depth"].estimator.predict(features.values)
    assert isinstance(predictions, np.ndarray)
    assert predictions.shape == (4,)

    artifact_path = save_training_bundle(bundle, tmp_path / "models.joblib")
    restored = load_training_bundle(artifact_path)
    assert restored.report() == bundle.report()


def test_training_requires_two_distinct_materials():
    dataset = TrainingDataset(
        records=(make_record(1, "Acrylic"), make_record(2, "Acrylic"))
    )

    with pytest.raises(ValueError, match="at least two materials"):
        train_models(dataset)