import numpy as np
import pytest

from lmpm.training.dataset import (
    NUMERIC_FEATURE_FIELDS,
    TrainingDataset,
    TrainingRecord,
)
from lmpm.training.pipeline import prepare_training_data
from lmpm.training.preprocessing import FeaturePreprocessor
from lmpm.training.routing import default_model_routes, model_route_for


def make_record(
    number: int, *, category: str | None = "bidirectional"
) -> TrainingRecord:
    features = {
        field: float(number + index)
        for index, field in enumerate(NUMERIC_FEATURE_FIELDS)
    }
    features["apparent_density"] = None if number == 1 else 1.2
    features.update(
        {
            "filling_method": category,
            "processing_size": "20x20",
        }
    )
    return TrainingRecord(
        metadata={"experiment_id": f"LAS-2026-{number:04d}"},
        features=features,
        targets={
            "is_cut_through": number % 2 == 0,
            "carbonized_edge_width": 0.2,
            "etching_depth": 30.0,
            "is_fire_smolder": False,
            "pattern_clarity_score": 9.0,
            "presentation_balance_score": 8.0,
        },
    )


def test_preprocessor_imputes_scales_and_one_hot_encodes_features():
    dataset = TrainingDataset(records=(make_record(1, category=None), make_record(2)))

    prepared = FeaturePreprocessor().fit_transform(dataset)

    assert prepared.values.shape == (2, len(NUMERIC_FEATURE_FIELDS) + 2)
    assert prepared.experiment_ids == ("LAS-2026-0001", "LAS-2026-0002")
    assert "numeric__apparent_density" in prepared.feature_names
    assert "categorical__filling_method_bidirectional" in prepared.feature_names
    assert "categorical__processing_size_20x20" in prepared.feature_names
    assert np.isfinite(prepared.values).all()


def test_preprocessor_ignores_unseen_categories_after_fitting():
    preprocessor = FeaturePreprocessor().fit(
        TrainingDataset(records=(make_record(1), make_record(2)))
    )

    transformed = preprocessor.transform(
        TrainingDataset(records=(make_record(3, category="spiral"),))
    )

    category_index = transformed.feature_names.index(
        "categorical__filling_method_bidirectional"
    )
    assert transformed.values[0, category_index] == 0


def test_preprocessor_requires_fit_and_complete_feature_rows():
    dataset = TrainingDataset(records=(make_record(1), make_record(2)))

    with pytest.raises(RuntimeError, match="fitted"):
        FeaturePreprocessor().transform(dataset)

    incomplete = make_record(1)
    incomplete.features.pop("actual_output_power")
    with pytest.raises(ValueError, match="actual_output_power"):
        FeaturePreprocessor().fit(TrainingDataset(records=(incomplete,)))


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

    assert routes["etching_depth"].model_family == "gaussian_process_regressor"
    assert routes["is_cut_through"].model_family == "random_forest_classifier"
    assert model_route_for("is_fire_smolder").task == "classification"
    with pytest.raises(ValueError, match="unsupported"):
        model_route_for("unknown")


def test_prepared_training_data_connects_features_targets_and_routes():
    prepared = prepare_training_data(
        TrainingDataset(records=(make_record(1), make_record(2)))
    )

    assert prepared.features.values.shape[0] == 2
    assert prepared.targets.regression["etching_depth"] == pytest.approx((30.0, 30.0))
    assert [route.target_name for route in prepared.model_routes] == [
        "is_cut_through",
        "carbonized_edge_width",
        "etching_depth",
        "is_fire_smolder",
        "pattern_clarity_score",
        "presentation_balance_score",
    ]