annotation/active_learning/human_feedback_loop.py

"""Human feedback loop for active learning.

Manages the cycle of:
1. Model prediction
2. Uncertainty sampling
3. Human annotation
4. Model retraining
"""

import json
import logging
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import pandas as pd

logger = logging.getLogger(__name__)


@dataclass
class FeedbackRecord:
    """Record of human feedback for a sample."""
    sample_id: str
    text: str
    original_prediction: str
    human_label: str
    confidence_feedback: float  # 0-1, did model seem confident?
    notes: str = ""
    timestamp: str = ""
    
    def to_dict(self) -> Dict:
        return {
            "sample_id": self.sample_id,
            "text": self.text,
            "original_prediction": self.original_prediction,
            "human_label": self.human_label,
            "confidence_feedback": self.confidence_feedback,
            "notes": self.notes,
            "timestamp": self.timestamp or datetime.now().isoformat(),
        }


@dataclass
class FeedbackLoopConfig:
    """Configuration for feedback loop."""
    min_feedback_samples: int = 50
    max_feedback_samples: int = 500
    retrain_threshold: int = 100  # Retrain after this many new samples
    disagreement_threshold: float = 0.3  # Retrain if disagreement rate > this
    batch_size: int = 32
    
    
@dataclass 
class LoopState:
    """State of the feedback loop."""
    iteration: int = 0
    total_annotated: int = 0
    total_retrained: int = 0
    disagreement_rate: float = 0.0
    model_performance: Dict = field(default_factory=dict)
    history: List[Dict] = field(default_factory=list)


class HumanFeedbackLoop:
    """Manages the human-in-the-loop training cycle."""
    
    def __init__(
        self,
        config: Optional[FeedbackLoopConfig] = None,
        output_dir: str = "outputs/active_learning",
    ):
        self.config = config or FeedbackLoopConfig()
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        self.state = LoopState()
        self.feedback_records: List[FeedbackRecord] = []
        self.labeled_samples: List[Dict] = []
        
    def add_feedback(
        self,
        sample_id: str,
        text: str,
        original_prediction: str,
        human_label: str,
        confidence_feedback: float = 0.5,
        notes: str = "",
    ) -> None:
        """Add human feedback for a sample."""
        record = FeedbackRecord(
            sample_id=sample_id,
            text=text,
            original_prediction=original_prediction,
            human_label=human_label,
            confidence_feedback=confidence_feedback,
            notes=notes,
            timestamp=datetime.now().isoformat(),
        )
        
        self.feedback_records.append(record)
        
        # Add to labeled samples
        self.labeled_samples.append({
            "id": sample_id,
            "text": text,
            "label": human_label,
            "source": "human_feedback",
        })
        
        self.state.total_annotated += 1
        
        logger.info(
            f"Added feedback for {sample_id}: "
            f"{original_prediction} -> {human_label}"
        )
    
    def batch_add_feedback(
        self,
        feedback_list: List[Dict],
    ) -> None:
        """Add multiple feedback records at once."""
        for fb in feedback_list:
            self.add_feedback(
                sample_id=fb.get("sample_id", fb.get("id")),
                text=fb.get("text", ""),
                original_prediction=fb.get("original_prediction", "unknown"),
                human_label=fb.get("human_label", fb.get("label")),
                confidence_feedback=fb.get("confidence_feedback", 0.5),
                notes=fb.get("notes", ""),
            )
    
    def should_retrain(self) -> Tuple[bool, str]:
        """Check if model should be retrained.
        
        Returns:
            (should_retrain, reason)
        """
        n_new = len(self.feedback_records)
        
        # Check minimum samples
        if n_new < self.config.min_feedback_samples:
            return False, f"Only {n_new} samples (min: {self.config.min_feedback_samples})"
        
        # Check retrain threshold
        if n_new >= self.config.retrain_threshold:
            self._calculate_disagreement_rate()
            if self.state.disagreement_rate > self.config.disagreement_threshold:
                return True, f"High disagreement ({self.state.disagreement_rate:.1%})"
            return True, f"Reached {n_new} samples threshold"
        
        return False, f"Not enough samples: {n_new}"
    
    def _calculate_disagreement_rate(self) -> float:
        """Calculate disagreement rate between model and human."""
        if not self.feedback_records:
            self.state.disagreement_rate = 0.0
            return 0.0
        
        disagreements = sum(
            1 for r in self.feedback_records
            if r.original_prediction != r.human_label
        )
        
        self.state.disagreement_rate = disagreements / len(self.feedback_records)
        return self.state.disagreement_rate
    
    def get_training_data(
        self,
        include_previous: bool = True,
    ) -> List[Dict]:
        """Get accumulated training data.
        
        Args:
            include_previous: Include previously retrained data
        
        Returns:
            List of samples with labels
        """
        if include_previous:
            return self.labeled_samples
        else:
            # Only return new samples since last retrain
            return self.labeled_samples[-self.config.retrain_threshold:]
    
    def export_training_data(
        self,
        path: Optional[str] = None,
        format: str = "jsonl",
    ) -> str:
        """Export training data to file."""
        if path is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            path = self.output_dir / f"training_data_{timestamp}.{format}"
        
        if format == "jsonl":
            with open(path, "w", encoding="utf-8") as f:
                for sample in self.labeled_samples:
                    f.write(json.dumps(sample, ensure_ascii=False) + "\n")
        elif format == "csv":
            df = pd.DataFrame(self.labeled_samples)
            df.to_csv(path, index=False)
        
        logger.info(f"Exported {len(self.labeled_samples)} samples to {path}")
        return str(path)
    
    def mark_retrained(self, performance: Optional[Dict] = None) -> None:
        """Mark that retraining has occurred."""
        self.state.iteration += 1
        self.state.total_retrained += 1
        
        if performance:
            self.state.model_performance = performance
        
        # Record history
        self.history.append({
            "iteration": self.state.iteration,
            "timestamp": datetime.now().isoformat(),
            "total_annotated": self.state.total_annotated,
            "disagreement_rate": self.state.disagreement_rate,
            "performance": performance,
        })
        
        logger.info(
            f"Model retrained (iteration {self.state.iteration}). "
            f"Total annotated: {self.state.total_annotated}"
        )
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get loop statistics."""
        return {
            "iteration": self.state.iteration,
            "total_annotated": self.state.total_annotated,
            "total_retrained": self.state.total_retrained,
            "disagreement_rate": self.state.disagreement_rate,
            "should_retrain": self.should_retrain()[0],
            "pending_samples": len(self.feedback_records),
            "recent_history": self.history[-5:] if self.history else [],
        }
    
    def get_label_distribution(self) -> Dict[str, int]:
        """Get distribution of labels."""
        from collections import Counter
        labels = [r.human_label for r in self.feedback_records]
        return dict(Counter(labels))
    
    def analyze_errors(self) -> Dict[str, Any]:
        """Analyze patterns in model errors."""
        errors = [
            r for r in self.feedback_records
            if r.original_prediction != r.human_label
        ]
        
        if not errors:
            return {"total_errors": 0}
        
        # Group by confusion pairs
        confusion_pairs = {}
        for e in errors:
            pair = (e.original_prediction, e.human_label)
            confusion_pairs[pair] = confusion_pairs.get(pair, 0) + 1
        
        return {
            "total_errors": len(errors),
            "error_rate": len(errors) / len(self.feedback_records),
            "confusion_matrix": confusion_pairs,
            "most_common_error": max(
                confusion_pairs.items(),
                key=lambda x: x[1]
            ) if confusion_pairs else None,
        }
    
    def save_state(self, path: Optional[str] = None) -> str:
        """Save loop state to file."""
        if path is None:
            path = self.output_dir / "loop_state.json"
        
        state_data = {
            "config": {
                "min_feedback_samples": self.config.min_feedback_samples,
                "max_feedback_samples": self.config.max_feedback_samples,
                "retrain_threshold": self.config.retrain_threshold,
                "disagreement_threshold": self.config.disagreement_threshold,
            },
            "state": {
                "iteration": self.state.iteration,
                "total_annotated": self.state.total_annotated,
                "total_retrained": self.state.total_retrained,
                "disagreement_rate": self.state.disagreement_rate,
            },
            "history": self.history,
        }
        
        with open(path, "w", encoding="utf-8") as f:
            json.dump(state_data, f, indent=2)
        
        return str(path)
    
    def load_state(self, path: str) -> None:
        """Load loop state from file."""
        with open(path, "r", encoding="utf-8") as f:
            state_data = json.load(f)
        
        config_dict = state_data.get("config", {})
        self.config = FeedbackLoopConfig(**config_dict)
        
        state_dict = state_data.get("state", {})
        self.state = LoopState(**state_dict)
        
        self.history = state_data.get("history", [])


def create_feedback_loop(
    config: Optional[Dict] = None,
) -> HumanFeedbackLoop:
    """Factory function to create feedback loop."""
    loop_config = None
    if config:
        loop_config = FeedbackLoopConfig(**config)
    
    return HumanFeedbackLoop(config=loop_config)


if __name__ == "__main__":
    loop = create_feedback_loop()
    
    # Simulate feedback
    loop.add_feedback(
        sample_id="utt_001",
        text="ကျေးဇူးပါ",
        original_prediction="positive",
        human_label="sarcastic",
        notes="Voice tone suggests complaint",
    )
    
    print(f"Should retrain: {loop.should_retrain()}")
    print(f"Stats: {loop.get_statistics()}")