"""Tests for the confidence feature pipeline.

Covers feature extraction, calibrator fit/predict, conservative defaults,
artifact save/load, and ECE/Brier computation.
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np
import pytest

from services.intelligence_pipeline_v3.confidence.artifacts import (
    list_versions,
    load_artifact,
    load_metadata,
    save_artifact,
)
from services.intelligence_pipeline_v3.confidence.calibrator import (
    ConfidenceCalibrator,
    _compute_brier,
    _compute_ece,
    compare_methods,
)
from services.intelligence_pipeline_v3.confidence.defaults import (
    get_default_confidence,
    is_underrepresented,
)
from services.intelligence_pipeline_v3.confidence.features import (
    AgreementStageResult,
    ConfidenceFeatureExtractor,
    EvidenceStageResult,
    ExtractionStageResult,
    ResolutionStageResult,
    SentimentStageResult,
)
from services.intelligence_pipeline_v3.confidence.models import (
    ConfidenceFeatures,
    ConfidenceResult,
)

# --- Fixtures ---


def _make_extraction_result(
    entity_scores: list[float] | None = None,
    relation_scores: list[float] | None = None,
    total_facts: int = 10,
    valid_numeric_facts: int = 8,
    populated_fields: int = 7,
    expected_fields: int = 10,
) -> ExtractionStageResult:
    return ExtractionStageResult(
        entity_scores=[0.9, 0.85, 0.7] if entity_scores is None else entity_scores,
        relation_scores=[0.8, 0.75] if relation_scores is None else relation_scores,
        total_facts=total_facts,
        valid_numeric_facts=valid_numeric_facts,
        populated_fields=populated_fields,
        expected_fields=expected_fields,
    )


def _make_resolution_result(
    margins: list[float] | None = None,
) -> ResolutionStageResult:
    return ResolutionStageResult(
        ambiguity_margins=[0.9, 0.6] if margins is None else margins,
    )


def _make_evidence_result(
    total: int = 10,
    supported: int = 8,
) -> EvidenceStageResult:
    return EvidenceStageResult(
        total_claims=total,
        supported_claims=supported,
    )


def _make_sentiment_result(
    probs: list[float] | None = None,
) -> SentimentStageResult:
    return SentimentStageResult(
        max_class_probabilities=[0.85, 0.9] if probs is None else probs,
        calibration_version="v1.0",
    )


def _make_agreement_result() -> AgreementStageResult:
    return AgreementStageResult(
        agreement_ratio=0.8,
        novelty_certainty=0.7,
        hard_case_score=0.2,
    )


def _make_features(
    entity_span_score: float = 0.85,
    document_type: str = "news",
) -> ConfidenceFeatures:
    return ConfidenceFeatures(
        entity_span_score=entity_span_score,
        alias_resolution_margin=0.75,
        numeric_parser_validity=0.8,
        evidence_coverage=0.8,
        relation_score=0.775,
        sentiment_calibration_confidence=0.875,
        cross_stage_agreement=0.8,
        duplicate_novelty_certainty=0.7,
        document_completeness=0.7,
        document_type=document_type,
        known_hard_case_patterns=0.2,
    )


def _generate_training_data(
    n_samples: int = 100,
    seed: int = 42,
) -> tuple[list[ConfidenceFeatures], list[bool]]:
    """Generate synthetic training data for calibrator tests."""
    rng = np.random.default_rng(seed)
    features = []
    labels = []
    doc_types = ["news", "filing", "transcript", "press_release", "macro_event"]

    for _ in range(n_samples):
        # Generate features with some correlation to the label
        base_quality = rng.uniform(0.3, 0.95)
        noise = rng.normal(0, 0.1)

        f = ConfidenceFeatures(
            entity_span_score=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)),
            alias_resolution_margin=float(np.clip(base_quality + rng.normal(0, 0.15), 0, 1)),
            numeric_parser_validity=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)),
            evidence_coverage=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)),
            relation_score=float(np.clip(base_quality + rng.normal(0, 0.15), 0, 1)),
            sentiment_calibration_confidence=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)),
            cross_stage_agreement=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)),
            duplicate_novelty_certainty=float(np.clip(base_quality + rng.normal(0, 0.15), 0, 1)),
            document_completeness=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)),
            document_type=rng.choice(doc_types),
            known_hard_case_patterns=float(np.clip(rng.uniform(0, 0.5), 0, 1)),
        )
        features.append(f)

        # Label correlates with base quality
        label = bool(rng.random() < (base_quality + noise))
        labels.append(label)

    return features, labels


# --- Test Feature Extraction ---


class TestFeatureExtraction:
    """Test that feature extraction produces valid feature vectors."""

    def test_extract_features_produces_valid_vector(self):
        """Feature extraction from all stages produces a valid ConfidenceFeatures."""
        extractor = ConfidenceFeatureExtractor()

        features = extractor.extract_features(
            extraction_result=_make_extraction_result(),
            resolution_result=_make_resolution_result(),
            evidence_result=_make_evidence_result(),
            sentiment_result=_make_sentiment_result(),
            agreement_result=_make_agreement_result(),
            document_type="news",
        )

        assert isinstance(features, ConfidenceFeatures)
        assert 0.0 <= features.entity_span_score <= 1.0
        assert 0.0 <= features.alias_resolution_margin <= 1.0
        assert 0.0 <= features.numeric_parser_validity <= 1.0
        assert 0.0 <= features.evidence_coverage <= 1.0
        assert 0.0 <= features.relation_score <= 1.0
        assert 0.0 <= features.sentiment_calibration_confidence <= 1.0
        assert 0.0 <= features.cross_stage_agreement <= 1.0
        assert 0.0 <= features.duplicate_novelty_certainty <= 1.0
        assert 0.0 <= features.document_completeness <= 1.0
        assert 0.0 <= features.known_hard_case_patterns <= 1.0
        assert features.document_type == "news"

    def test_extract_features_without_agreement(self):
        """Feature extraction uses sensible defaults when agreement is not available."""
        extractor = ConfidenceFeatureExtractor()

        features = extractor.extract_features(
            extraction_result=_make_extraction_result(),
            resolution_result=_make_resolution_result(),
            evidence_result=_make_evidence_result(),
            sentiment_result=_make_sentiment_result(),
            agreement_result=None,
            document_type="filing",
        )

        assert features.cross_stage_agreement == 0.5
        assert features.duplicate_novelty_certainty == 0.5
        assert features.known_hard_case_patterns == 0.0

    def test_extract_features_empty_entities(self):
        """Feature extraction handles empty entity scores gracefully."""
        extractor = ConfidenceFeatureExtractor()

        features = extractor.extract_features(
            extraction_result=_make_extraction_result(entity_scores=[]),
            resolution_result=_make_resolution_result(),
            evidence_result=_make_evidence_result(),
            sentiment_result=_make_sentiment_result(),
        )

        assert features.entity_span_score == 0.0

    def test_extract_features_no_claims(self):
        """Feature extraction handles zero claims gracefully."""
        extractor = ConfidenceFeatureExtractor()

        features = extractor.extract_features(
            extraction_result=_make_extraction_result(),
            resolution_result=_make_resolution_result(),
            evidence_result=_make_evidence_result(total=0, supported=0),
            sentiment_result=_make_sentiment_result(),
        )

        assert features.evidence_coverage == 0.0

    def test_to_vector_produces_correct_length(self):
        """Feature vector has expected dimensionality."""
        features = _make_features()
        vector = features.to_vector()
        assert len(vector) == 11
        assert all(isinstance(v, float) for v in vector)

    def test_unknown_document_type_defaults(self):
        """Unknown document types are normalized to 'unknown'."""
        extractor = ConfidenceFeatureExtractor()

        features = extractor.extract_features(
            extraction_result=_make_extraction_result(),
            resolution_result=_make_resolution_result(),
            evidence_result=_make_evidence_result(),
            sentiment_result=_make_sentiment_result(),
            document_type="exotic_type",
        )

        assert features.document_type == "unknown"


# --- Test Calibrator ---


class TestCalibrator:
    """Test calibrator fit/predict roundtrip and method comparison."""

    def test_fit_predict_isotonic(self):
        """Isotonic calibrator can fit and produce predictions in [0, 1]."""
        features, labels = _generate_training_data(n_samples=50)
        cal = ConfidenceCalibrator(method="isotonic")

        cal.fit(features, labels, version="test-v1")

        assert cal.is_fitted
        assert cal.version == "test-v1"

        prediction = cal.predict(features[0])
        assert 0.0 <= prediction <= 1.0

    def test_fit_predict_platt(self):
        """Platt calibrator can fit and produce predictions in [0, 1]."""
        features, labels = _generate_training_data(n_samples=50)
        cal = ConfidenceCalibrator(method="platt")

        cal.fit(features, labels, version="test-v1")

        assert cal.is_fitted
        prediction = cal.predict(features[0])
        assert 0.0 <= prediction <= 1.0

    def test_unfitted_returns_neutral(self):
        """Unfitted calibrator returns 0.5 as neutral default."""
        cal = ConfidenceCalibrator()
        features = _make_features()

        prediction = cal.predict(features)
        assert prediction == 0.5

    def test_predict_batch(self):
        """Batch prediction returns correct number of results."""
        features, labels = _generate_training_data(n_samples=50)
        cal = ConfidenceCalibrator(method="isotonic")
        cal.fit(features, labels)

        batch_predictions = cal.predict_batch(features[:10])
        assert len(batch_predictions) == 10
        assert all(0.0 <= p <= 1.0 for p in batch_predictions)

    def test_fit_empty_raises(self):
        """Fitting with empty data raises ValueError."""
        cal = ConfidenceCalibrator()
        with pytest.raises(ValueError, match="must not be empty"):
            cal.fit([], [])

    def test_fit_mismatched_lengths_raises(self):
        """Fitting with mismatched lengths raises ValueError."""
        features, labels = _generate_training_data(n_samples=10)
        cal = ConfidenceCalibrator()
        with pytest.raises(ValueError, match="must have the same length"):
            cal.fit(features, labels[:5])

    def test_metadata_after_fit(self):
        """Metadata is populated after fitting."""
        features, labels = _generate_training_data(n_samples=50)
        cal = ConfidenceCalibrator(method="isotonic")
        cal.fit(features, labels, version="v1.0.0", training_range="2024-01-01 to 2024-06-30")

        assert cal.metadata is not None
        assert cal.metadata.version == "v1.0.0"
        assert cal.metadata.method == "isotonic"
        assert cal.metadata.training_count == 50
        assert cal.metadata.training_range == "2024-01-01 to 2024-06-30"
        assert 0.0 <= cal.metadata.ece <= 1.0
        assert 0.0 <= cal.metadata.brier_score <= 1.0

    def test_compare_methods(self):
        """Method comparison returns ECE and Brier for both methods."""
        features, labels = _generate_training_data(n_samples=50)

        results = compare_methods(features, labels, n_folds=3)

        assert "isotonic" in results
        assert "platt" in results
        assert "ece" in results["isotonic"]
        assert "brier" in results["isotonic"]
        assert "ece" in results["platt"]
        assert "brier" in results["platt"]


# --- Test ECE and Brier ---


class TestMetrics:
    """Test ECE and Brier score computation."""

    def test_ece_perfect_calibration(self):
        """ECE is 0 for perfectly calibrated predictions."""
        # Perfect: predict 1.0 for positives, 0.0 for negatives
        predictions = np.array([1.0, 1.0, 0.0, 0.0, 1.0])
        labels = np.array([1.0, 1.0, 0.0, 0.0, 1.0])

        ece = _compute_ece(predictions, labels)
        assert ece == pytest.approx(0.0, abs=1e-10)

    def test_ece_worst_calibration(self):
        """ECE is high for badly calibrated predictions."""
        # Predict 1.0 but all are actually 0
        predictions = np.array([0.9, 0.9, 0.9, 0.9, 0.9])
        labels = np.array([0.0, 0.0, 0.0, 0.0, 0.0])

        ece = _compute_ece(predictions, labels)
        assert ece > 0.5

    def test_brier_perfect_predictions(self):
        """Brier score is 0 for perfect predictions."""
        predictions = np.array([1.0, 0.0, 1.0, 0.0])
        labels = np.array([1.0, 0.0, 1.0, 0.0])

        brier = _compute_brier(predictions, labels)
        assert brier == pytest.approx(0.0, abs=1e-10)

    def test_brier_worst_predictions(self):
        """Brier score is 1 for worst possible predictions."""
        predictions = np.array([1.0, 1.0, 0.0, 0.0])
        labels = np.array([0.0, 0.0, 1.0, 1.0])

        brier = _compute_brier(predictions, labels)
        assert brier == pytest.approx(1.0, abs=1e-10)

    def test_brier_uniform_predictions(self):
        """Brier score for uniform 0.5 predictions against balanced labels is 0.25."""
        predictions = np.array([0.5, 0.5, 0.5, 0.5])
        labels = np.array([1.0, 0.0, 1.0, 0.0])

        brier = _compute_brier(predictions, labels)
        assert brier == pytest.approx(0.25, abs=1e-10)

    def test_ece_empty_returns_zero(self):
        """ECE of empty arrays is 0."""
        ece = _compute_ece(np.array([]), np.array([]))
        assert ece == 0.0

    def test_brier_empty_returns_zero(self):
        """Brier of empty arrays is 0."""
        brier = _compute_brier(np.array([]), np.array([]))
        assert brier == 0.0


# --- Test Conservative Defaults ---


class TestDefaults:
    """Test conservative defaults for underrepresented classes."""

    def test_known_document_type(self):
        """Known document types return conservative probabilities in [0.3, 0.5]."""
        result = get_default_confidence("news", "earnings_beat")

        assert isinstance(result, ConfidenceResult)
        assert 0.3 <= result.probability <= 0.5
        assert result.under_calibrated is True
        assert result.is_calibrated is False
        assert "conservative-default" in result.calibration_version

    def test_unknown_document_type(self):
        """Unknown document types return the most conservative default (0.3)."""
        result = get_default_confidence("exotic_type", "unknown_event")

        assert result.probability == 0.3
        assert result.under_calibrated is True

    def test_unknown_event_class(self):
        """Unknown event classes use the lowest default."""
        result = get_default_confidence("news", "never_seen_before")

        assert result.probability == 0.30
        assert result.under_calibrated is True

    def test_all_document_types_conservative(self):
        """All defined document types have defaults in [0.3, 0.5]."""
        doc_types = ["news", "filing", "transcript", "press_release", "macro_event", "unknown"]
        for dt in doc_types:
            result = get_default_confidence(dt, "earnings_beat")
            assert 0.3 <= result.probability <= 0.5, f"Failed for {dt}"

    def test_is_underrepresented_no_counts(self):
        """Without known counts, unknown types are underrepresented."""
        assert is_underrepresented("exotic", "unknown_event") is True
        assert is_underrepresented("news", "earnings_beat") is False

    def test_is_underrepresented_with_counts(self):
        """With known counts, low-count classes are underrepresented."""
        counts = {("news", "earnings_beat"): 100, ("filing", "merger"): 5}
        assert is_underrepresented("news", "earnings_beat", known_counts=counts) is False
        assert is_underrepresented("filing", "merger", known_counts=counts) is True
        assert is_underrepresented("news", "unknown", known_counts=counts) is True


# --- Test Artifact Save/Load ---


class TestArtifacts:
    """Test calibration artifact persistence."""

    def test_save_load_roundtrip(self):
        """Save and load preserves calibrator state."""
        features, labels = _generate_training_data(n_samples=50)
        cal = ConfidenceCalibrator(method="isotonic")
        cal.fit(features, labels, version="v1.0.0", training_range="test")

        with tempfile.TemporaryDirectory() as tmpdir:
            save_artifact(cal, "v1.0.0", tmpdir)

            loaded = load_artifact(Path(tmpdir) / "v1.0.0")

            assert loaded.is_fitted
            assert loaded.version == "v1.0.0"
            assert loaded.method == "isotonic"

            # Predictions should match
            test_features = _make_features()
            original_pred = cal.predict(test_features)
            loaded_pred = loaded.predict(test_features)
            assert original_pred == pytest.approx(loaded_pred, abs=1e-10)

    def test_save_unfitted_raises(self):
        """Saving an unfitted calibrator raises ValueError."""
        cal = ConfidenceCalibrator()
        with tempfile.TemporaryDirectory() as tmpdir:
            with pytest.raises(ValueError, match="unfitted"):
                save_artifact(cal, "v1.0.0", tmpdir)

    def test_load_nonexistent_raises(self):
        """Loading from a missing path raises FileNotFoundError."""
        with pytest.raises(FileNotFoundError):
            load_artifact("/nonexistent/path")

    def test_load_metadata(self):
        """Metadata can be loaded independently."""
        features, labels = _generate_training_data(n_samples=50)
        cal = ConfidenceCalibrator(method="platt")
        cal.fit(features, labels, version="v2.0.0", training_range="2024-01-01 to 2024-12-31")

        with tempfile.TemporaryDirectory() as tmpdir:
            save_artifact(cal, "v2.0.0", tmpdir)

            metadata = load_metadata(Path(tmpdir) / "v2.0.0")
            assert metadata.version == "v2.0.0"
            assert metadata.method == "platt"
            assert metadata.training_count == 50

    def test_list_versions(self):
        """list_versions finds all saved artifact versions."""
        features, labels = _generate_training_data(n_samples=50)

        with tempfile.TemporaryDirectory() as tmpdir:
            for version in ["v1.0.0", "v1.1.0", "v2.0.0"]:
                cal = ConfidenceCalibrator(method="isotonic")
                cal.fit(features, labels, version=version)
                save_artifact(cal, version, tmpdir)

            versions = list_versions(tmpdir)
            assert versions == ["v1.0.0", "v1.1.0", "v2.0.0"]

    def test_list_versions_empty_dir(self):
        """list_versions returns empty list for empty or missing directory."""
        with tempfile.TemporaryDirectory() as tmpdir:
            assert list_versions(tmpdir) == []
        assert list_versions("/nonexistent") == []
