import json
import sys
import types
import unittest
from unittest.mock import MagicMock, patch

from owmi.analysis import calibration_scores
from owmi.backends import HFBackend
from owmi.benchmarks.base import BenchmarkExample
from owmi.benchmarks.evaluate import (
    characterization_scores,
    evaluate_probe_output,
    reconstruction_scores,
    semantic_similarity,
)
from owmi.benchmarks.runner import BenchmarkRunner
from owmi.benchmarks.suite import (
    SuiteConfig,
    build_job_manifest,
    build_manifest,
)


class RevisionRoundTripTests(unittest.TestCase):
    def _config(self):
        return SuiteConfig.from_dict({
            'suite_name': 'revision-test',
            'models': [{
                'name': 'org/model',
                'revision': '0123456789abcdef0123456789abcdef01234567',
                'load_in_4bit': False,
            }],
            'benchmarks': [{
                'name': 'toy', 'schema': 'short_answer', 'dataset_name': 'unused',
            }],
            'objects': [{'kind': 'residual_stream', 'layer_index': 1}],
            'interventions': [{'kind': 'zero', 'mode': 'zero'}],
            'probes': [{'task': 'detection'}],
        })

    def test_revision_round_trips_from_suite_through_hf_loading(self):
        example = BenchmarkExample('toy', 'one', 'short_answer', 'q', {'answer': 'a'})
        revision = self._config().models[0].revision
        with patch('owmi.benchmarks.suite._load_examples_for_spec', return_value=[example]):
            expanded = build_manifest(self._config()).rows[0].to_dict()
            compact = build_job_manifest(self._config())[0].to_dict()
        expanded = json.loads(json.dumps(expanded))
        compact = json.loads(json.dumps(compact))
        self.assertEqual(expanded['model']['revision'], revision)
        self.assertEqual(compact['model']['revision'], revision)

        config = BenchmarkRunner(expanded)._build_experiment_config()
        self.assertEqual(config.revision, revision)
        tokenizer = MagicMock()
        tokenizer.pad_token = '<pad>'
        model = MagicMock()
        auto_tokenizer = MagicMock()
        auto_tokenizer.from_pretrained.return_value = tokenizer
        auto_model = MagicMock()
        auto_model.from_pretrained.return_value = model
        transformers = types.ModuleType('transformers')
        transformers.AutoTokenizer = auto_tokenizer
        transformers.AutoModelForCausalLM = auto_model
        with patch.dict(sys.modules, {'transformers': transformers}):
            HFBackend(config.model_name, config).load()
        self.assertEqual(auto_tokenizer.from_pretrained.call_args.kwargs['revision'], revision)
        self.assertEqual(auto_model.from_pretrained.call_args.kwargs['revision'], revision)


class ConfidenceFamilyTests(unittest.TestCase):
    def test_selective_prediction_auroc_uses_confidence_for_correctness(self):
        scores = calibration_scores(
            confidence=[0.1, 0.9, 0.2, 0.8],
            correctness=[0.0, 1.0, 0.0, 1.0],
        )
        self.assertEqual(scores['selective_prediction_auroc'], 1.0)
        self.assertEqual(scores['selective_auroc'], 1.0)
        tied = calibration_scores([0.5, 0.5], [0.0, 1.0])
        self.assertEqual(tied['selective_prediction_auroc'], 0.5)


class ProbeScoringTests(unittest.TestCase):
    def test_characterization_is_closed_set_and_reports_macro_f1(self):
        metrics = characterization_scores(
            ['zeroing', 'noise', 'masked', 'noise happened'],
            ['zero', 'noise', 'zero', 'noise'],
            ('zero', 'noise'),
        )
        self.assertEqual(metrics['category_accuracy'], 0.75)
        self.assertAlmostEqual(metrics['macro_f1'], 5.0 / 6.0)
        example = BenchmarkExample('toy', 'one', 'short_answer', 'q', {'answer': 'a'})
        self.assertEqual(
            evaluate_probe_output(
                '{"change_type": "noise happened"}', 'characterization', example,
                {'mode': 'noise', 'categories': ['zero', 'noise']},
            ),
            0.0,
        )

    @patch('owmi.benchmarks.evaluate.semantic_similarity', return_value=(0.95, 'test-embedding'))
    def test_reconstruction_combines_exact_category_and_embedding(self, _similarity):
        categories = ('animal', 'vehicle')
        gamed = reconstruction_scores(
            'A domesticated canine', 'A household dog',
            predicted_category='vehicle', target_category='animal', categories=categories,
        )
        self.assertGreaterEqual(gamed['embedding_similarity'], gamed['similarity_threshold'])
        self.assertEqual(gamed['category_match'], 0.0)
        self.assertEqual(gamed['combined_match'], 0.0)

        valid_paraphrase = reconstruction_scores(
            'A domesticated canine', 'A household dog',
            predicted_category='animal', target_category='animal', categories=categories,
        )
        self.assertEqual(valid_paraphrase['combined_match'], 1.0)
        exact = reconstruction_scores(
            'A household dog', 'the household dog',
            predicted_category=None, target_category='animal', categories=categories,
        )
        self.assertEqual(exact['exact_match'], 1.0)
        self.assertEqual(exact['combined_match'], 1.0)

    def test_reconstruction_similarity_has_dependency_free_tfidf_fallback(self):
        with patch('owmi.benchmarks.evaluate._SENTENCE_TRANSFORMER_UNAVAILABLE', True):
            similarity, backend = semantic_similarity('masked household dog', 'household dog')
        self.assertGreater(similarity, 0.0)
        self.assertEqual(backend, 'tfidf-word-1-2gram-cosine')

    def test_embedding_similarity_ranks_known_similar_above_known_dissimilar(self):
        # Item #19: the embedding-similarity metric must rank a known-similar
        # pair above a known-dissimilar pair, under whichever backend is
        # actually available in this environment (sentence-transformers if
        # cached locally, otherwise the dependency-free TF-IDF fallback).
        anchor = 'a domesticated canine that barks and wags its tail'
        similar, _ = semantic_similarity(anchor, 'a household dog that barks and wags its tail')
        dissimilar, _ = semantic_similarity(anchor, 'quarterly tax filing deadlines for small businesses')
        self.assertGreater(similar, dissimilar)
        identical, _ = semantic_similarity(anchor, anchor)
        self.assertGreaterEqual(identical, similar)


if __name__ == '__main__':
    unittest.main()