"""Integration tests for the hierarchical classification pipeline.

These tests train a small synthetic model end-to-end and verify that the core
operations (fit, predict, evaluate, save/load round-trip) work correctly.
They are intentionally slow (a few seconds) because they exercise real sklearn
estimators and the full embedding pipeline.
"""
import os
import tempfile
import pytest
import numpy as np
import pandas as pd

from pwml.classifiers.features import InputFeature, OutputFeature
from pwml.classifiers.hierarchical import HierarchicalClassifierModel


# ---------------------------------------------------------------------------
# Synthetic dataset helpers
# ---------------------------------------------------------------------------

def _make_synthetic_data(n=120, seed=0):
    """Return a two-level classification DataFrame with numeric + category features.

    Level 1 (Animal): mammal / bird / reptile  (40 rows each)
    Level 2 (Species): depends on Animal
        mammal  -> cat / dog
        bird    -> eagle / parrot
        reptile -> snake / lizard
    """
    rng = np.random.default_rng(seed)

    animals   = ['mammal'] * 40 + ['bird'] * 40 + ['reptile'] * 40
    species_map = {
        'mammal':  ['cat', 'dog'],
        'bird':    ['eagle', 'parrot'],
        'reptile': ['snake', 'lizard'],
    }
    species = [rng.choice(species_map[a]) for a in animals]

    df = pd.DataFrame({
        'weight':  rng.uniform(0.5, 80.0, n),
        'length':  rng.uniform(0.1, 3.0, n),
        'habitat': rng.choice(['forest', 'desert', 'ocean'], n),
        'animal':  animals,
        'species': species,
    })
    return df


def _make_model(cache_path=None):
    input_features = [
        InputFeature('weight',  'numeric'),
        InputFeature('length',  'numeric'),
        InputFeature('habitat', 'category'),
    ]
    output_hierarchy = OutputFeature(
        'animal',
        child_feature=OutputFeature('species'))
    return HierarchicalClassifierModel(
        model_name='test_model',
        experiment_name='integration_test',
        input_features=input_features,
        output_feature_hierarchy=output_hierarchy,
        cache_path=cache_path,
    )


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

class TestEndToEnd:

    def test_fit_and_predict_single_row(self, tmp_path):
        """fit then predict a single dict input."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))
        model.load_from_dataframe(df)

        row = {'weight': 5.0, 'length': 0.5, 'habitat': 'forest'}
        result = model.predict(input=row)

        assert isinstance(result, list)
        assert len(result) >= 1
        for r in result:
            assert 'feature_name' in r
            assert 'value' in r

    def test_fit_and_predict_dataframe(self, tmp_path):
        """predict_dataframe returns a DataFrame with the expected columns."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))
        model.load_from_dataframe(df)

        preds = model.predict_dataframe(df)

        assert 'animal_predicted' in preds.columns
        assert 'species_predicted' in preds.columns
        assert len(preds) == len(df)

    def test_evaluate_returns_accuracy_per_level(self, tmp_path):
        """evaluate() returns a float accuracy for each output feature."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))
        model.load_from_dataframe(df)

        metrics, predictions = model.evaluate(df)

        assert 'animal' in metrics
        assert 'species' in metrics
        assert 0.0 <= metrics['animal'] <= 1.0
        assert 0.0 <= metrics['species'] <= 1.0

    def test_save_load_roundtrip(self, tmp_path):
        """A model saved and loaded must produce identical predictions."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))
        model.load_from_dataframe(df)

        model_path = str(tmp_path / 'model.joblib')
        model.save_model(model_path)

        loaded = HierarchicalClassifierModel.load_model(model_path)
        loaded_preds = loaded.predict_dataframe(df)
        original_preds = model.predict_dataframe(df)

        pd.testing.assert_frame_equal(original_preds, loaded_preds)

    def test_predict_unknown_category_does_not_raise(self, tmp_path):
        """An unseen category value at inference time must not raise."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))
        model.load_from_dataframe(df)

        row = {'weight': 5.0, 'length': 0.5, 'habitat': 'arctic'}  # unseen
        result = model.predict(input=row)
        assert isinstance(result, list)

    def test_predict_missing_input_feature_returns_empty(self, tmp_path):
        """If a required feature is absent, predict() returns an empty list."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))
        model.load_from_dataframe(df)

        result = model.predict(input={'weight': 5.0})  # missing length, habitat
        assert result == []

    def test_numeric_features_are_normalised(self, tmp_path):
        """NumericEmbedder must produce values in [0, 1] for in-range inputs."""
        from pwml.classifiers.embedders import NumericEmbedder
        from pwml.classifiers.features import InputFeature

        embedder = NumericEmbedder(cache_path=str(tmp_path))
        feature = InputFeature('weight', 'numeric')
        feature.range = (0.0, 100.0)

        result = embedder.embed_data(50.0, feature)
        assert result.shape == (1,)
        assert abs(result[0] - 0.5) < 1e-9

        result_min = embedder.embed_data(0.0, feature)
        assert abs(result_min[0] - 0.0) < 1e-9

        result_max = embedder.embed_data(100.0, feature)
        assert abs(result_max[0] - 1.0) < 1e-9

    def test_numeric_missing_range_falls_back_to_raw(self, tmp_path):
        """NumericEmbedder without range stores the raw float unchanged."""
        from pwml.classifiers.embedders import NumericEmbedder
        from pwml.classifiers.features import InputFeature

        embedder = NumericEmbedder(cache_path=str(tmp_path))
        feature = InputFeature('score', 'numeric')
        # No range set

        result = embedder.embed_data(42.0, feature)
        assert abs(result[0] - 42.0) < 1e-9

    def test_cross_validate_summary_structure(self, tmp_path):
        """cross_validate() returns a summary dict with mean/std per feature
        and per-fold metric dicts when run on sufficient data."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))

        summary, fold_results = model.cross_validate(df, n_splits=3)

        # Summary must contain mean and std for each output level
        assert 'animal' in summary
        assert 'species' in summary
        for key in ('mean', 'std'):
            assert key in summary['animal']
            assert key in summary['species']
        assert 0.0 <= summary['animal']['mean'] <= 1.0

        # One result dict per fold
        assert len(fold_results) == 3
        for fr in fold_results:
            assert 'animal' in fr
            assert 0.0 <= fr['animal'] <= 1.0

    def test_cross_validate_warns_on_sparse_classes(self, tmp_path):
        """cross_validate() emits UserWarning when a top-level class has
        fewer samples than n_splits, before any training begins."""
        import warnings

        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))

        # n_splits larger than any class count (40 per class) so warning fires
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter('always')
            try:
                model.cross_validate(df, n_splits=50)
            except Exception:
                pass  # training may fail with so few rows per fold

        user_warnings = [w for w in caught if issubclass(w.category, UserWarning)
                         and 'cross_validate' in str(w.message)]
        assert len(user_warnings) >= 1

    def test_cross_validate_deterministic(self, tmp_path):
        """Two cross_validate() calls with the same random_state produce identical fold metrics."""
        df = _make_synthetic_data()
        model1 = _make_model(cache_path=str(tmp_path / 'run1'))
        model2 = _make_model(cache_path=str(tmp_path / 'run2'))

        summary1, folds1 = model1.cross_validate(df, n_splits=3, random_state=42)
        summary2, folds2 = model2.cross_validate(df, n_splits=3, random_state=42)

        assert summary1 == summary2
        assert folds1 == folds2

    def test_save_load_cache_path_preserved(self, tmp_path):
        """cache_path survives a save_model/load_model round-trip."""
        df = _make_synthetic_data()
        cache_dir = str(tmp_path / 'cache')
        model = _make_model(cache_path=cache_dir)
        model.load_from_dataframe(df)

        model_path = str(tmp_path / 'model.joblib')
        model.save_model(model_path)

        loaded = HierarchicalClassifierModel.load_model(model_path)
        assert loaded.cache_path == model.cache_path

    def test_routing_uncertain_flag_propagated(self, tmp_path):
        """With min_routing_confidence=1.0, routing should always stop after
        the first level and set routing_uncertain=True on that result."""
        df = _make_synthetic_data()
        model = _make_model(cache_path=str(tmp_path))
        model.load_from_dataframe(df)

        row = {'weight': 5.0, 'length': 0.5, 'habitat': 'forest'}
        result = model.predict(input=row, min_routing_confidence=1.0)

        # With confidence threshold=1.0 (impossible to reach), routing must
        # stop after the first node, marking it as uncertain.
        last = result[-1]
        assert last.get('routing_uncertain') is True
