scoring.py

"""
Scoring Module
End-to-end pipeline for computing audio similarity scores.
"""

import logging

import torch
from typing import Optional, Dict
from audio_loader import AudioLoader
from wavlm_encoder import WavLMEncoder
from dtw_similarity import DTWSimilarity

logger = logging.getLogger(__name__)


class SimilarityScorer:
    """End-to-end similarity scoring pipeline."""
    
    def __init__(
        self,
        model_name: str = "./wavlm-base-plus",
        device: Optional[str] = None,
        distance_metric: str = "cosine",
        sakoe_chiba_ratio: float = 0.1,
        normalize_dtw: bool = True,
        score_midpoint: float = 0.3,
        score_steepness: float = 10.0
    ):
        """
        Initialize similarity scorer.
        
        Args:
            model_name: WavLM model name from HuggingFace
            device: Device to run model on
            distance_metric: Distance metric for DTW ("cosine" or "euclidean")
            sakoe_chiba_ratio: Band radius ratio for DTW constraint
            normalize_dtw: Whether to normalize DTW distance by path length
            score_midpoint: Distance value that maps to 50 for the logisitic score
            score_steepness: Logistic steepness for score mapping
        """
        self.audio_loader = AudioLoader(target_sr=16000)
        self.encoder = WavLMEncoder(model_name=model_name, device=device)
        self.dtw_similarity = DTWSimilarity(distance_metric=distance_metric, sakoe_chiba_ratio=sakoe_chiba_ratio)
        self.normalize_dtw = normalize_dtw
        self.score_midpoint = score_midpoint
        self.score_steepness = score_steepness
    
    def compute_similarity(
        self, 
        audio_path1: str, 
        audio_path2: str
    ) -> float:
        """
        Compute similarity score between two audio files.
        
        Args:
            audio_path1: Path to first audio file (participant's recitation)
            audio_path2: Path to second audio file (reference recitation)
            
        Returns:
            Similarity score (DTW distance - lower means more similar)
        """
        # Load and preprocess audio
        waveform1 = self.audio_loader.load_and_preprocess(audio_path1)
        waveform2 = self.audio_loader.load_and_preprocess(audio_path2)
        
        # Extract features
        features1 = self.encoder.encode(waveform1)
        features2 = self.encoder.encode(waveform2)
        
        # Compute DTW similarity
        dtw_distance = self.dtw_similarity.compute_similarity(
            features1, 
            features2,
            normalize=self.normalize_dtw
        )
        
        logger.info(
            "compute_similarity | audio1=%s | audio2=%s | dtw_distance=%.6f",
            audio_path1, audio_path2, dtw_distance,
        )
        
        return dtw_distance
    
    def compute_similarity_score_normalized(
        self, 
        audio_path1: str, 
        audio_path2: str
    ) -> float:
        """
        Compute normalized similarity score (0-100 scale, higher is better).
        
        Uses a logistic (sigmoid) mapping calibrated so that:
          - identical pairs  (d ≈ 0.05) → score ~ 92
          - very similar     (d ≈ 0.15) → score ~ 82
          * borderline       (d ≈ 0.30) → score = 50
          * different pairs  (d ≈ 0.35) → score ~ 38
          * noise/non-speech (d > 0.50) → score < 12
        
        Args:
            audio_path1: Path to first audio file
            audio_path2: Path to second audio file
            
        Returns:
            Normalized similarity score (0-100, where 100 is perfect match)
        """
        # Get DTW distance (already path-normalised when self.normalize_dtw)
        dtw_distance = self.compute_similarity(audio_path1, audio_path2)
        
        # Logistic mapping (calibrated for cosine-distance DTW)
        similarity_score = DTWSimilarity.distance_to_score(
            dtw_distance,
            midpoint=self.score_midpoint,
            steepness=self.score_steepness
        )
        
        logger.info(
            "score_normalized | dtw_distance=%.6f | score=%.2f",
            dtw_distance, similarity_score,
        )
        
        return similarity_score
    
    def compute_detailed_similarity(
        self, 
        audio_path1: str, 
        audio_path2: str,
        use_vad: bool = True,
        layer_indices: Optional[list] = None
    ) -> Dict:
        """
        Compute detailed similarity analysis (supports multiple layers).
        
        Args:
            audio_path1: Path to first audio file (reference)
            audio_path2: Path to second audio file (test)
            use_vad: Whether to use VAD endpoint trimming
            layer_indices: List of WavLM layers to extract (defaults to [12])
            
        Returns:
            Dictionary containing detailed metrics, warping paths, DTW matrices, and original waveforms.
        """
        if layer_indices is None:
            layer_indices = [12]

        self.audio_loader.use_vad = use_vad

        # Process and retrieve both raw and VAD waveforms
        waveform_ref_raw, sr_ref = self.audio_loader.load_audio(audio_path1)
        waveform_test_raw, sr_test = self.audio_loader.load_audio(audio_path2)

        if use_vad:
            waveform_ref_vad = self.audio_loader.vad_trim_endpoints(waveform_ref_raw, sr_ref)
            waveform_test_vad = self.audio_loader.vad_trim_endpoints(waveform_test_raw, sr_test)
        else:
            waveform_ref_vad = waveform_ref_raw
            waveform_test_vad = waveform_test_raw

        waveform_ref = self.audio_loader.normalize_audio(waveform_ref_vad)
        waveform_test = self.audio_loader.normalize_audio(waveform_test_vad)
        
        # Extract features for all requested layers
        feat_ref_dict = self.encoder.encode(waveform_ref, extract_layers=layer_indices)
        feat_test_dict = self.encoder.encode(waveform_test, extract_layers=layer_indices)
        
        results = {}
        for layer in layer_indices:
            seq1 = feat_ref_dict[layer].cpu().numpy()
            seq2 = feat_test_dict[layer].cpu().numpy()
            
            # Compute DTW
            raw_distance, dtw_matrix = self.dtw_similarity.dtw(seq1, seq2)
            warping_path = self.dtw_similarity.compute_path(dtw_matrix)
            path_length = len(warping_path)
            
            # Normalize
            if self.normalize_dtw and path_length > 0:
                normalized_distance = raw_distance / path_length
            else:
                normalized_distance = raw_distance
            
            # Logistic score
            normalized_score = DTWSimilarity.distance_to_score(
                normalized_distance,
                midpoint=self.score_midpoint,
                steepness=self.score_steepness
            )

            num_frames_ref = seq1.shape[0]
            num_frames_test = seq2.shape[0]
            sr = 16000
            ref_duration_sec = round(num_frames_ref * 320 / sr, 3)
            test_duration_sec = round(num_frames_test * 320 / sr, 3)
            duration_ratio = round(test_duration_sec / ref_duration_sec, 4) if ref_duration_sec > 0 else 0.0
            
            diagnostics = {
                "raw_dtw_distance": float(raw_distance),
                "normalized_distance": float(normalized_distance),
                "path_length": int(path_length),
                "num_frames_ref": num_frames_ref,
                "num_frames_test": num_frames_test,
                "feature_dimension": seq1.shape[1],
                "ref_duration_sec": ref_duration_sec,
                "test_duration_sec": test_duration_sec,
                "duration_ratio": duration_ratio,
                "distance_metric": self.dtw_similarity.distance_metric,
                "normalized": self.normalize_dtw,
                "sakoe_chiba_ratio": self.dtw_similarity.sakoe_chiba_ratio,
            }

            results[layer] = {
                "score": normalized_score,
                "dtw_distance": float(normalized_distance),
                "raw_dtw_distance": float(raw_distance),
                "warping_path": warping_path,
                "dtw_matrix": dtw_matrix,
                "diagnostics": diagnostics
            }
        
        return {
            "results": results,
            "waveforms": {
                "ref_raw": waveform_ref_raw,
                "ref_vad": waveform_ref_vad,
                "ref_normalized": waveform_ref,
                "test_raw": waveform_test_raw,
                "test_vad": waveform_test_vad,
                "test_normalized": waveform_test,
            }
        }