import pytest
from pydantic import ValidationError

from lmpm.training.dataset import TrainingDataset, TrainingRecord
from lmpm.training.targets import TargetSpec, split_targets, target_specs


def make_record(number: int = 1, *, target_overrides=None) -> TrainingRecord:
    targets = {
        "is_cut_through": True,
        "carbonized_edge_width": 0.2,
        "etching_depth": 30.0,
        "is_fire_smolder": False,
        "pattern_clarity_score": 9.0,
        "presentation_balance_score": 8.0,
    }
    if target_overrides:
        targets.update(target_overrides)

    return TrainingRecord(
        metadata={"experiment_id": f"LAS-2026-{number:04d}"},
        features={"actual_output_power": 8.5},
        targets=targets,
    )


def test_target_specs_describe_all_platform_targets():
    specs = target_specs()

    assert [(spec.name, spec.task) for spec in specs] == [
        ("is_cut_through", "classification"),
        ("carbonized_edge_width", "regression"),
        ("etching_depth", "regression"),
        ("is_fire_smolder", "classification"),
        ("pattern_clarity_score", "regression"),
        ("presentation_balance_score", "regression"),
    ]


def test_target_spec_rejects_unknown_task_type():
    with pytest.raises(ValidationError):
        TargetSpec(name="unsupported_target", task="unknown")


def test_split_targets_groups_by_task_type():
    dataset = TrainingDataset(
        records=(
            make_record(1),
            make_record(2, target_overrides={"is_cut_through": False}),
        )
    )

    split = split_targets(dataset)

    assert split.classification["is_cut_through"] == (True, False)
    assert split.classification["is_fire_smolder"] == (False, False)
    assert split.regression["carbonized_edge_width"] == pytest.approx((0.2, 0.2))
    assert split.regression["etching_depth"] == pytest.approx((30.0, 30.0))
    assert split.regression["pattern_clarity_score"] == pytest.approx((9.0, 9.0))
    assert split.regression["presentation_balance_score"] == pytest.approx((8.0, 8.0))


def test_empty_dataset_has_empty_target_groups():
    split = split_targets(TrainingDataset(records=()))

    assert split.classification == {"is_cut_through": (), "is_fire_smolder": ()}
    assert split.regression == {
        "carbonized_edge_width": (),
        "etching_depth": (),
        "pattern_clarity_score": (),
        "presentation_balance_score": (),
    }


@pytest.mark.parametrize("target_name", ["is_cut_through", "etching_depth"])
def test_missing_target_value_is_rejected(target_name):
    dataset = TrainingDataset(
        records=(make_record(target_overrides={target_name: None}),)
    )

    with pytest.raises(ValueError, match=target_name):
        split_targets(dataset)