annotation/automatic_verifier.py

"""Automatic verification of annotations using rule-based checks."""

import json
import logging
import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import pandas as pd
import yaml

logger = logging.getLogger(__name__)


class VerificationRule:
    """Base class for verification rules."""
    
    def __init__(self, rule_id: str, severity: str = "error"):
        self.rule_id = rule_id
        self.severity = severity  # error, warning, info
    
    def verify(self, sample: Dict) -> Tuple[bool, Optional[str]]:
        """Verify a sample. Returns (passed, error_message)."""
        raise NotImplementedError


class TextLengthRule(VerificationRule):
    """Check if text length is within acceptable range."""
    
    def __init__(self, min_length: int = 1, max_length: int = 500):
        super().__init__("text_length")
        self.min_length = min_length
        self.max_length = max_length
    
    def verify(self, sample: Dict) -> Tuple[bool, Optional[str]]:
        text = sample.get("text", "")
        length = len(text.strip())
        
        if length < self.min_length:
            return False, f"Text too short: {length} chars (min: {self.min_length})"
        if length > self.max_length:
            return False, f"Text too long: {length} chars (max: {self.max_length})"
        
        return True, None


class SentimentConsistencyRule(VerificationRule):
    """Check sentiment consistency between text and prosody."""
    
    # Sentiment keywords mapping
    POSITIVE_KEYWORDS = ["ကျေးဇူး", "ပါး", "မင်္ဂလာ", "ဝမ်းသာ", "ပျော်"]
    NEGATIVE_KEYWORDS = ["မကျေနပ်", "ဒေါသ", "စိတ်ဓာတ်ကျ", "ပူ", "ဆူ"]
    
    def __init__(self):
        super().__init__("sentiment_consistency")
    
    def verify(self, sample: Dict) -> Tuple[bool, Optional[str]]:
        text = sample.get("text", "")
        prosody = sample.get("prosody", {})
        
        text_positive = any(kw in text for kw in self.POSITIVE_KEYWORDS)
        text_negative = any(kw in text for kw in self.NEGATIVE_KEYWORDS)
        
        sentiment = sample.get("sentiment", "neutral")
        
        # Warning only (not error) - prosody and text can diverge
        if sentiment == "positive" and text_negative and not text_positive:
            return True, "Warning: Text has negative keywords but labeled positive"
        if sentiment == "negative" and text_positive and not text_negative:
            return True, "Warning: Text has positive keywords but labeled negative"
        
        return True, None


class LabelDistributionRule(VerificationRule):
    """Check if label distribution is reasonable."""
    
    def __init__(
        self,
        min_samples_per_class: int = 10,
        max_imbalance_ratio: float = 10.0,
    ):
        super().__init__("label_distribution", severity="warning")
        self.min_samples_per_class = min_samples_per_class
        self.max_imbalance_ratio = max_imbalance_ratio
    
    def verify(self, samples: List[Dict]) -> Tuple[bool, Optional[str]]:
        if not samples:
            return False, "No samples provided"
        
        from collections import Counter
        sentiments = [s.get("sentiment", "unknown") for s in samples]
        counts = Counter(sentiments)
        
        if len(counts) == 0:
            return False, "No sentiment labels found"
        
        # Check minimum samples per class
        for sentiment, count in counts.items():
            if count < self.min_samples_per_class:
                return False, f"Class '{sentiment}' has only {count} samples (min: {self.min_samples_per_class})"
        
        # Check imbalance
        max_count = max(counts.values())
        min_count = min(counts.values())
        if max_count / min_count > self.max_imbalance_ratio:
            return True, f"Warning: Imbalance ratio {max_count/min_count:.1f} > {self.max_imbalance_ratio}"
        
        return True, None


class DuplicateTextRule(VerificationRule):
    """Check for duplicate text entries."""
    
    def __init__(self, threshold: float = 0.9):
        super().__init__("duplicate_text", severity="warning")
        self.threshold = threshold
    
    def verify(self, samples: List[Dict]) -> Tuple[bool, Optional[str]]:
        texts = [s.get("text", "").strip().lower() for s in samples]
        
        duplicates = []
        seen = {}
        
        for i, text in enumerate(texts):
            if text in seen:
                duplicates.append((seen[text], i))
            else:
                seen[text] = i
        
        if duplicates:
            return True, f"Found {len(duplicates)} potential duplicates"
        
        return True, None


class AutomaticVerifier:
    """Verify annotations using rule-based checks."""
    
    def __init__(self, rules_config: Optional[str] = None):
        self.rules: List[VerificationRule] = []
        
        if rules_config and Path(rules_config).exists():
            self._load_config(rules_config)
        else:
            self._setup_default_rules()
    
    def _setup_default_rules(self) -> None:
        """Set up default verification rules."""
        self.rules = [
            TextLengthRule(min_length=1, max_length=500),
            SentimentConsistencyRule(),
            LabelDistributionRule(),
            DuplicateTextRule(),
        ]
    
    def _load_config(self, config_path: str) -> None:
        """Load rules from config file."""
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
        
        self.rules = []
        
        for rule_def in config.get("rules", []):
            rule_type = rule_def.get("type")
            
            if rule_type == "text_length":
                self.rules.append(TextLengthRule(
                    min_length=rule_def.get("min_length", 1),
                    max_length=rule_def.get("max_length", 500),
                ))
            elif rule_type == "sentiment_consistency":
                self.rules.append(SentimentConsistencyRule())
            elif rule_type == "label_distribution":
                self.rules.append(LabelDistributionRule(
                    min_samples_per_class=rule_def.get("min_samples_per_class", 10),
                    max_imbalance_ratio=rule_def.get("max_imbalance_ratio", 10.0),
                ))
            elif rule_type == "duplicate_text":
                self.rules.append(DuplicateTextRule(
                    threshold=rule_def.get("threshold", 0.9),
                ))
    
    def verify_sample(self, sample: Dict) -> Dict[str, Any]:
        """Verify a single sample against all rules."""
        results = {
            "sample_id": sample.get("id", "unknown"),
            "passed": True,
            "errors": [],
            "warnings": [],
        }
        
        for rule in self.rules:
            if isinstance(rule, LabelDistributionRule) or isinstance(rule, DuplicateTextRule):
                # These rules need full dataset
                continue
            
            passed, message = rule.verify(sample)
            
            if not passed:
                results["passed"] = False
                if rule.severity == "error":
                    results["errors"].append({
                        "rule_id": rule.rule_id,
                        "message": message,
                    })
                else:
                    results["warnings"].append({
                        "rule_id": rule.rule_id,
                        "message": message,
                    })
            elif message:
                results["warnings"].append({
                    "rule_id": rule.rule_id,
                    "message": message,
                })
        
        return results
    
    def verify_dataset(
        self,
        samples: List[Dict],
    ) -> Dict[str, Any]:
        """Verify entire dataset."""
        results = {
            "total_samples": len(samples),
            "sample_results": [],
            "dataset_errors": [],
            "statistics": {},
        }
        
        # Check sample-level rules
        for sample in samples:
            sample_result = self.verify_sample(sample)
            results["sample_results"].append(sample_result)
        
        # Check dataset-level rules
        for rule in self.rules:
            if isinstance(rule, (LabelDistributionRule, DuplicateTextRule)):
                passed, message = rule.verify(samples)
                if not passed:
                    results["dataset_errors"].append({
                        "rule_id": rule.rule_id,
                        "severity": rule.severity,
                        "message": message,
                    })
        
        # Calculate statistics
        total_errors = sum(
            len(r.get("errors", []))
            for r in results["sample_results"]
        )
        total_warnings = sum(
            len(r.get("warnings", []))
            for r in results["sample_results"]
        )
        
        results["statistics"] = {
            "total_errors": total_errors,
            "total_warnings": total_warnings,
            "samples_passed": sum(
                1 for r in results["sample_results"] if r["passed"]
            ),
        }
        
        return results
    
    def filter_samples(
        self,
        samples: List[Dict],
        remove_errors: bool = True,
        remove_warnings: bool = False,
    ) -> Tuple[List[Dict], List[Dict]]:
        """Filter samples based on verification results."""
        results = self.verify_dataset(samples)
        
        kept = []
        removed = []
        
        for i, (sample, result) in enumerate(zip(samples, results["sample_results"])):
            should_remove = False
            
            if remove_errors and result["errors"]:
                should_remove = True
            if remove_warnings and result["warnings"]:
                should_remove = True
            
            if should_remove:
                removed.append({
                    "sample": sample,
                    "reason": result,
                })
            else:
                kept.append(sample)
        
        logger.info(
            f"Filtered: {len(kept)} kept, {len(removed)} removed"
        )
        
        return kept, removed
    
    def generate_report(
        self,
        samples: List[Dict],
        output_path: Optional[str] = None,
    ) -> str:
        """Generate verification report."""
        results = self.verify_dataset(samples)
        
        report_lines = [
            "=" * 60,
            "AUTOMATIC ANNOTATION VERIFICATION REPORT",
            "=" * 60,
            f"Total Samples: {results['total_samples']}",
            f"Samples Passed: {results['statistics']['samples_passed']}",
            f"Total Errors: {results['statistics']['total_errors']}",
            f"Total Warnings: {results['statistics']['total_warnings']}",
            "",
            "-" * 60,
            "DATASET-LEVEL ISSUES",
            "-" * 60,
        ]
        
        for error in results.get("dataset_errors", []):
            report_lines.append(
                f"[{error['severity'].upper()}] {error['rule_id']}: {error['message']}"
            )
        
        if not results.get("dataset_errors"):
            report_lines.append("No dataset-level issues found.")
        
        report_lines.extend([
            "",
            "-" * 60,
            "SAMPLE-LEVEL ISSUES",
            "-" * 60,
        ])
        
        error_count = 0
        for result in results["sample_results"]:
            if result["errors"] or result["warnings"]:
                error_count += 1
                report_lines.append(f"\nSample: {result['sample_id']}")
                for error in result["errors"]:
                    report_lines.append(f"  ERROR: {error['message']}")
                for warning in result["warnings"]:
                    report_lines.append(f"  WARNING: {warning['message']}")
                
                if error_count >= 20:
                    report_lines.append("\n... (showing first 20 samples with issues)")
                    break
        
        report = "\n".join(report_lines)
        
        if output_path:
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(report)
            logger.info(f"Report saved to {output_path}")
        
        return report


def create_verifier(config_path: Optional[str] = None) -> AutomaticVerifier:
    """Factory function to create verifier."""
    return AutomaticVerifier(rules_config=config_path)


if __name__ == "__main__":
    verifier = create_verifier()
    
    # Test samples
    test_samples = [
        {
            "id": "utt_001",
            "text": "ကျေးဇူးပါ",
            "sentiment": "positive",
            "prosody": {"mean_pitch": 150},
        },
        {
            "id": "utt_002",
            "text": "",
            "sentiment": "negative",
        },
    ]
    
    results = verifier.verify_dataset(test_samples)
    print(f"Verification results: {results['statistics']}")