tests/test_accuracy.py

from __future__ import annotations

import unittest
from pathlib import Path
from unittest.mock import patch

import numpy as np

from src.aspect_sentiment.audio import WhisperTranscriber
from src.aspect_sentiment.diarization import (
    DiarizationResult,
    TranscriptTurn,
    _refine_turn_roles,
    diarize_audio_segments,
    diarize_text,
)
from src.aspect_sentiment.conversation_reconstruction import reconstruct_conversation
from src.aspect_sentiment.llama_extraction import merge_rule_features
from src.aspect_sentiment.privacy import extract_and_redact_pii
from src.aspect_sentiment.role_classifier import classify_role_hybrid
from src.aspect_sentiment.tracking import SpeakerTracker
from src.api.server import local_structured_entities, readiness
from scripts.benchmark_diarization import compute_embedding_metrics, compute_reference_metrics


class PrivacyAccuracyTests(unittest.TestCase):
    def test_only_extracts_customer_name_with_direct_evidence(self):
        result = extract_and_redact_pii(
            "My name is Alice Johnson. I need a Samsung phone. John reviewed the proposal."
        )

        names = [entity.value for entity in result.entities if entity.type == "customer_name"]
        self.assertEqual(names, ["Alice Johnson"])
        self.assertNotIn("John", names)
        self.assertIn("[CUSTOMER_NAME_REDACTED]", result.cleaned_text)

    def test_does_not_treat_product_or_intent_as_a_name(self):
        result = extract_and_redact_pii("I am Looking for Samsung Galaxy under 50000.")
        names = [entity.value for entity in result.entities if entity.type == "customer_name"]
        self.assertEqual(names, [])

    def test_agent_and_customer_names_follow_speaker_evidence(self):
        diarization = DiarizationResult(
            turns=[
                TranscriptTurn("Agent", "Hello Ravi. My name is Sarah."),
                TranscriptTurn("Customer", "Hi Sarah. My name is Ravi."),
            ],
            provider="test",
        )

        entities = local_structured_entities(diarization.formatted, diarization)
        grouped = {
            entity_type: {item["value"] for item in entities if item["type"] == entity_type}
            for entity_type in ("customer_name", "agent_name")
        }

        self.assertEqual(grouped["customer_name"], {"Ravi"})
        self.assertEqual(grouped["agent_name"], {"Sarah"})

    def test_nested_labels_keep_agent_and_customer_names_separate(self):
        text = (
            "Customer: [Agent]: Good morning. My name is Jennifer. "
            "Customer: [Customer]: Hi Jennifer. My name is Michael Thomas. "
            "I'm looking for a laptop."
        )
        diarization = diarize_text(text)
        privacy = extract_and_redact_pii(diarization.customer_text)
        entities = local_structured_entities(text, diarization)
        privacy_names = {
            entity.value for entity in privacy.entities if entity.type == "customer_name"
        }
        grouped = {
            entity_type: {item["value"] for item in entities if item["type"] == entity_type}
            for entity_type in ("customer_name", "agent_name")
        }

        self.assertEqual(privacy_names, {"Michael Thomas"})
        self.assertEqual(grouped["customer_name"], {"Michael Thomas"})
        self.assertEqual(grouped["agent_name"], {"Jennifer"})

    def test_model_number_is_not_extracted_as_budget(self):
        text = (
            "Agent: The price of the MacBook Pro M4 with 16GB RAM and 512GB SSD "
            "is approximately ₹1,89,000. Customer: That's higher than my budget. "
            "I'm hoping to stay around ₹1,60,000."
        )
        entities = local_structured_entities(text, diarize_text(text))
        values = {(item["type"], item["value"]) for item in entities}

        self.assertNotIn(("budget", "4"), values)
        self.assertIn(("product_price", "189000"), values)
        self.assertIn(("budget", "160000"), values)


class DiarizationAccuracyTests(unittest.TestCase):
    def test_acoustic_speaker_is_not_flipped_by_sentence_words(self):
        turns = [
            TranscriptTurn(
                "Agent",
                "I need your account details. I can suggest a Dell laptop.",
                raw_speaker="SPEAKER_0",
            )
        ]

        refined = _refine_turn_roles(turns, preserve_speakers=True)
        self.assertTrue(refined)
        self.assertEqual({turn.speaker for turn in refined}, {"Agent"})

    def test_speaker_tracker_returns_match_confidence(self):
        tracker = SpeakerTracker(threshold=0.75, max_speakers=2)

        first = tracker.track_speaker_with_confidence(np.array([1.0, 0.0], dtype=np.float32))
        repeat = tracker.track_speaker_with_confidence(np.array([0.96, 0.04], dtype=np.float32))
        second = tracker.track_speaker_with_confidence(np.array([0.0, 1.0], dtype=np.float32))
        capped = tracker.track_speaker_with_confidence(np.array([0.7, 0.7], dtype=np.float32))

        self.assertEqual(first.speaker, "Speaker_A")
        self.assertEqual(repeat.speaker, "Speaker_A")
        self.assertGreaterEqual(repeat.confidence, 0.95)
        self.assertEqual(second.speaker, "Speaker_B")
        self.assertTrue(second.is_new)
        self.assertIn(capped.speaker, {"Speaker_A", "Speaker_B"})
        self.assertEqual(tracker.speaker_names, ["Speaker_A", "Speaker_B"])

    def test_vad_ecapa_pipeline_reports_overlap_and_confidence(self):
        whisper_segments = [
            {"start": 0.7, "end": 1.1, "text": "How can I help you today?"},
            {"start": 1.2, "end": 1.8, "text": "I need a laptop under 60000."},
        ]
        vad_segments = [
            {"start": 0.0, "end": 1.2},
            {"start": 0.8, "end": 2.0},
        ]
        embeddings = [
            np.array([1.0, 0.0], dtype=np.float32),
            np.array([0.0, 1.0], dtype=np.float32),
        ]

        def classify(speaker: str, text: str, **_: object) -> dict:
            return {
                "speaker": speaker,
                "role": "Agent" if speaker == "Speaker_A" else "Customer",
                "confidence": 0.96,
                "method": "test",
            }

        with (
            patch.dict("os.environ", {
                "ENABLE_SPEAKER_TRACKING": "true",
                "USE_LLM_DIARIZATION": "false",
                "USE_GROQ_WHISPER": "false",
            }),
            patch("src.aspect_sentiment.vad.get_speech_segments", return_value=vad_segments),
            patch("src.aspect_sentiment.embeddings.get_speaker_embedding", side_effect=embeddings),
            patch("src.aspect_sentiment.diarization._load_audio_mono", return_value=(np.zeros(32000), 16000)),
            patch("src.aspect_sentiment.role_classifier.classify_role_hybrid", side_effect=classify),
            patch("src.aspect_sentiment.flow_validator.validate_and_correct_roles", side_effect=lambda _, c, threshold=0.85: c),
        ):
            result = diarize_audio_segments(Path("sample.wav"), whisper_segments)

        self.assertEqual(result.provider, "vad-ecapa-tracking")
        self.assertEqual(result.speaker_map, {"Speaker_A": "Agent", "Speaker_B": "Customer"})
        self.assertIn("Speaker_A", result.speaker_confidence)
        self.assertTrue(result.turns[0].overlap)
        self.assertIn("overlapping_speech_detected", result.warnings)
        self.assertGreater(result.turns[0].confidence or 0.0, 0.5)

    def test_diarization_quality_metrics_with_reference(self):
        predicted = [
            {"start": 0.0, "end": 1.0, "speaker": "Speaker_A"},
            {"start": 1.0, "end": 2.0, "speaker": "Speaker_B"},
            {"start": 2.0, "end": 3.0, "speaker": "Speaker_A"},
        ]
        reference = {
            "segments": [
                {"start": 0.0, "end": 1.0, "speaker": "Agent"},
                {"start": 1.0, "end": 2.0, "speaker": "Customer"},
                {"start": 2.0, "end": 3.0, "speaker": "Agent"},
            ]
        }

        metrics = compute_reference_metrics(predicted, reference)

        self.assertEqual(metrics["diarizationErrorRatePct"], 0.0)
        self.assertEqual(metrics["speakerPurityPct"], 100.0)
        self.assertEqual(metrics["speakerConsistencyPct"], 100.0)
        self.assertEqual(metrics["speakerSwitchingErrors"], 0)

    def test_conversation_reconstruction_merges_fragments_and_preserves_overlap(self):
        diarization = DiarizationResult(
            turns=[
                TranscriptTurn("Customer", "under 60000", start=1.0, end=1.4, confidence=0.9),
                TranscriptTurn("Agent", "What is your budget?", start=0.0, end=0.8, confidence=0.8),
                TranscriptTurn("Customer", "I need a laptop", start=0.9, end=1.0, confidence=0.9),
                TranscriptTurn("Agent", "I can suggest Dell.", start=1.5, end=2.0, confidence=0.85, overlap=True),
            ],
            warnings=["test_warning"],
        )

        result = reconstruct_conversation(diarization)

        self.assertEqual(result.turns[0].speaker, "Agent")
        self.assertIn("I need a laptop under 60000", result.formatted)
        self.assertEqual(result.metadata.merged_fragments, 1)
        self.assertEqual(result.metadata.overlap_turns, 1)
        self.assertIn("overlapping_speech_preserved", result.metadata.warnings)

    def test_embedding_quality_metrics_report_false_speaker_creation(self):
        rows = [
            {"speaker": "Speaker_A", "embedding": np.array([1.0, 0.0], dtype=np.float32)},
            {"speaker": "Speaker_A", "embedding": np.array([0.9, 0.1], dtype=np.float32)},
            {"speaker": "Speaker_B", "embedding": np.array([0.0, 1.0], dtype=np.float32)},
            {"speaker": "Speaker_C", "embedding": np.array([0.1, 0.9], dtype=np.float32)},
        ]

        metrics = compute_embedding_metrics(rows, expected_speakers=2)

        self.assertEqual(metrics["falseSpeakerCreation"], 1)
        self.assertIn("avgIntraSpeaker", metrics["embeddingSimilarity"])


class ExtractionAccuracyTests(unittest.TestCase):
    def test_rejects_ungrounded_llama_features(self):
        features = merge_rule_features(
            [
                {"value": "Samsung", "label": "BRAND"},
                {"value": "RTX 4090", "label": "FEATURE"},
                {"value": "Ready to Purchase", "label": "INTENT"},
            ],
            "I need a Dell laptop under 60000, but I am not sure yet.",
        )

        values = {str(feature["value"]).lower() for feature in features}
        self.assertNotIn("samsung", values)
        self.assertNotIn("rtx 4090", values)
        self.assertIn("dell", values)
        self.assertIn("60000", values)


class RoleClassificationAccuracyTests(unittest.TestCase):
    def test_multi_signal_role_classifier_returns_probabilities(self):
        agent = classify_role_hybrid(
            "Speaker_A",
            "Good morning, this is Sarah calling from TechNova. May I know your budget and brand preference?",
            speaker_word_count=16,
            total_word_count=32,
        )
        customer = classify_role_hybrid(
            "Speaker_B",
            "I am looking for a laptop under 60000 but I am not sure if I should buy today.",
            speaker_word_count=17,
            total_word_count=32,
        )

        self.assertEqual(agent["role"], "Agent")
        self.assertEqual(customer["role"], "Customer")
        self.assertIn("probability", agent)
        self.assertIn("signals", agent)
        self.assertGreater(agent["probability"]["Agent"], customer["probability"]["Agent"])


class WhisperAccuracyTests(unittest.TestCase):
    def test_filters_low_confidence_hallucinations_and_repeats(self):
        segments = WhisperTranscriber._clean_segments(
            [
                {
                    "start": 0,
                    "end": 1,
                    "text": "Thanks for watching",
                    "avg_logprob": -1.2,
                    "no_speech_prob": 0.8,
                    "compression_ratio": 1.0,
                },
                {
                    "start": 1,
                    "end": 2,
                    "text": "I need a laptop",
                    "avg_logprob": -0.2,
                    "no_speech_prob": 0.05,
                    "compression_ratio": 1.1,
                },
                {
                    "start": 2,
                    "end": 3,
                    "text": "I need a laptop",
                    "avg_logprob": -0.2,
                    "no_speech_prob": 0.05,
                    "compression_ratio": 1.1,
                },
            ]
        )

        self.assertEqual([segment["text"] for segment in segments], ["I need a laptop"])

    def test_preserves_word_timestamps_and_removes_repeated_words(self):
        segments = WhisperTranscriber._clean_segments(
            [
                {
                    "start": 0,
                    "end": 1,
                    "text": "I I need need a laptop",
                    "avg_logprob": -0.1,
                    "no_speech_prob": 0.05,
                    "compression_ratio": 1.1,
                    "words": [{"word": "I", "start": 0.0, "end": 0.1}],
                }
            ]
        )

        self.assertEqual(segments[0]["text"], "I need a laptop")
        self.assertEqual(segments[0]["words"][0]["word"], "I")


class ApiHealthTests(unittest.TestCase):
    def test_readiness_returns_a_payload(self):
        payload = readiness()
        self.assertIn(payload["status"], {"ready", "degraded"})
        self.assertIn("checks", payload)


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