#!/usr/bin/env python3
"""
Robust Timing Analyzer for RF Device Identification

Extracts timing characteristics from raw RF pulse data and scores against
known protocol signatures. Handles noisy/imperfect captures with statistical
methods and tolerance windows.

Key features:
  - Multi-method clustering (K-means, histogram, percentile)
  - Outlier removal for noise tolerance
  - Timing ratio calculation (short/long pulse ratios)
  - Similarity scoring with configurable tolerance
  - Support for 2-4 level modulation schemes
"""

import numpy as np
from dataclasses import dataclass
from typing import List, Tuple, Optional, Dict
from collections import defaultdict


@dataclass
class TimingCharacteristics:
    """
    Extracted timing characteristics from raw signal

    Attributes:
        short_pulse_us: Short HIGH pulse duration (microseconds)
        long_pulse_us: Long HIGH pulse duration (microseconds)
        short_gap_us: Short LOW gap duration (microseconds)
        long_gap_us: Long LOW gap duration (microseconds)
        pulse_ratio: Ratio of long_pulse / short_pulse
        gap_ratio: Ratio of long_gap / short_gap
        duty_cycle: Fraction of time signal is HIGH
        pulse_count: Total number of pulses
        confidence: Confidence in timing extraction (0.0-1.0)
        method: Extraction method used ('kmeans', 'histogram', 'percentile')
    """
    short_pulse_us: int
    long_pulse_us: int
    short_gap_us: int
    long_gap_us: int
    pulse_ratio: float
    gap_ratio: float
    duty_cycle: float
    pulse_count: int
    confidence: float
    method: str


@dataclass
class TimingScore:
    """
    Similarity score between extracted timing and protocol signature

    Attributes:
        protocol_name: Name of the protocol
        overall_score: Overall similarity score (0.0-1.0)
        pulse_score: Similarity of pulse timings (0.0-1.0)
        gap_score: Similarity of gap timings (0.0-1.0)
        ratio_score: Similarity of timing ratios (0.0-1.0)
        frequency_score: Frequency match score (0.0-1.0)
        details: Additional scoring details
    """
    protocol_name: str
    overall_score: float
    pulse_score: float
    gap_score: float
    ratio_score: float
    frequency_score: float
    details: Dict


class RobustTimingAnalyzer:
    """
    Robust timing analyzer with noise tolerance and multi-method extraction

    Workflow:
      1. Clean pulse data (remove outliers, normalize)
      2. Extract timing characteristics (try multiple methods)
      3. Calculate timing ratios and statistics
      4. Score against protocol database
      5. Return ranked matches
    """

    def __init__(self, outlier_threshold: float = 1.5):
        """
        Initialize analyzer

        Args:
            outlier_threshold: IQR multiplier for outlier detection (default: 1.5)
        """
        self.outlier_threshold = outlier_threshold

    def extract_timing(self, pulses: List[int]) -> TimingCharacteristics:
        """
        Extract timing characteristics from pulse data

        Uses multi-method approach and returns best result:
          1. K-means clustering (preferred for clean signals)
          2. Histogram peak detection (good for multi-modal)
          3. Percentile-based (fallback for noisy signals)

        Args:
            pulses: List of pulse durations (positive=HIGH, negative=LOW)

        Returns:
            TimingCharacteristics with confidence score
        """
        if not pulses or len(pulses) < 4:
            return self._empty_timing()

        # Step 1: Separate HIGH and LOW pulses BEFORE outlier removal
        # (outlier removal can corrupt alternating pattern)
        high_pulses_raw = [abs(p) for p in pulses if p > 0]
        low_pulses_raw = [abs(p) for p in pulses if p < 0]

        if not high_pulses_raw:
            return self._empty_timing()

        # Step 2: Remove outliers from HIGH and LOW separately
        high_pulses = self._remove_outliers_from_values(high_pulses_raw)
        low_pulses = self._remove_outliers_from_values(low_pulses_raw) if low_pulses_raw else []

        if not high_pulses:
            return self._empty_timing()

        # Step 3: Try multiple extraction methods
        methods = [
            ('kmeans', self._extract_kmeans),
            ('histogram', self._extract_histogram),
            ('percentile', self._extract_percentile)
        ]

        best_timing = None
        best_confidence = 0.0

        for method_name, method_func in methods:
            try:
                timing = method_func(high_pulses, low_pulses)
                if timing and timing.confidence > best_confidence:
                    best_timing = timing
                    best_confidence = timing.confidence
                    best_timing.method = method_name
            except Exception:
                continue

        if best_timing:
            return best_timing
        else:
            return self._empty_timing()

    def score_against_protocol(
        self,
        timing: TimingCharacteristics,
        protocol_short_us: int,
        protocol_long_us: int,
        protocol_frequency: int,
        signal_frequency: int,
        timing_tolerance: float = 0.25
    ) -> TimingScore:
        """
        Score timing characteristics against a protocol signature

        Args:
            timing: Extracted timing characteristics
            protocol_short_us: Protocol's short pulse width
            protocol_long_us: Protocol's long pulse width
            protocol_frequency: Protocol's expected frequency
            signal_frequency: Actual signal frequency
            timing_tolerance: Tolerance for timing match (default: ±25%)

        Returns:
            TimingScore with component scores and overall score
        """
        # Pulse timing similarity
        short_error = abs(timing.short_pulse_us - protocol_short_us) / protocol_short_us
        long_error = abs(timing.long_pulse_us - protocol_long_us) / protocol_long_us

        pulse_score = self._error_to_score(
            max(short_error, long_error),
            tolerance=timing_tolerance
        )

        # Gap timing similarity (if available)
        gap_score = 1.0  # Default: don't penalize gaps

        # Ratio similarity
        protocol_ratio = protocol_long_us / protocol_short_us if protocol_short_us > 0 else 0
        ratio_error = abs(timing.pulse_ratio - protocol_ratio) / max(protocol_ratio, 0.01)
        ratio_score = self._error_to_score(ratio_error, tolerance=0.3)

        # Frequency similarity
        freq_error = abs(signal_frequency - protocol_frequency) / protocol_frequency
        frequency_score = self._error_to_score(freq_error, tolerance=0.01)  # ±1%

        # Overall score (weighted average)
        overall_score = (
            pulse_score * 0.50 +      # Pulse timing most important
            ratio_score * 0.25 +       # Ratio consistency
            frequency_score * 0.20 +   # Frequency match
            gap_score * 0.05          # Gaps less important
        ) * timing.confidence  # Scale by extraction confidence

        return TimingScore(
            protocol_name="",  # Will be set by caller
            overall_score=overall_score,
            pulse_score=pulse_score,
            gap_score=gap_score,
            ratio_score=ratio_score,
            frequency_score=frequency_score,
            details={
                'short_error_pct': f"{short_error:.1%}",
                'long_error_pct': f"{long_error:.1%}",
                'ratio_error_pct': f"{ratio_error:.1%}",
                'freq_error_pct': f"{freq_error:.1%}",
                'extraction_method': timing.method
            }
        )

    # === Internal Methods ===

    def _remove_outliers_from_values(self, values: List[int]) -> List[int]:
        """
        Remove statistical outliers from a list of values using IQR method

        Args:
            values: List of positive values (already abs())

        Returns:
            Cleaned values
        """
        if len(values) < 10:
            return values  # Too few samples for outlier detection

        q1 = np.percentile(values, 25)
        q3 = np.percentile(values, 75)
        iqr = q3 - q1

        # Use more lenient threshold (2.5 instead of 1.5)
        lower_bound = q1 - 2.5 * iqr
        upper_bound = q3 + 2.5 * iqr

        cleaned = [v for v in values if lower_bound <= v <= upper_bound]

        # Only apply if doesn't remove too much
        if len(cleaned) >= len(values) * 0.7:
            return cleaned
        else:
            return values

    def _extract_kmeans(
        self,
        high_pulses: List[int],
        low_pulses: List[int]
    ) -> Optional[TimingCharacteristics]:
        """Extract timing using K-means clustering (k=2)"""
        try:
            from sklearn.cluster import KMeans
            from sklearn.metrics import silhouette_score
        except ImportError:
            return None

        if len(high_pulses) < 10:
            return None

        X = np.array(high_pulses).reshape(-1, 1)
        kmeans = KMeans(n_clusters=2, random_state=42, n_init=10)
        labels = kmeans.fit_predict(X)

        # Silhouette score: quality of clustering
        confidence = silhouette_score(X, labels)

        # Get cluster centers
        centers = sorted(kmeans.cluster_centers_.flatten())
        short_pulse = int(centers[0])
        long_pulse = int(centers[1])

        # Extract gap timings
        short_gap, long_gap = self._extract_gaps_percentile(low_pulses)

        # Calculate metrics
        pulse_ratio = long_pulse / short_pulse if short_pulse > 0 else 0
        gap_ratio = long_gap / short_gap if short_gap > 0 else 0
        duty_cycle = self._calculate_duty_cycle(high_pulses, low_pulses)

        return TimingCharacteristics(
            short_pulse_us=short_pulse,
            long_pulse_us=long_pulse,
            short_gap_us=short_gap,
            long_gap_us=long_gap,
            pulse_ratio=pulse_ratio,
            gap_ratio=gap_ratio,
            duty_cycle=duty_cycle,
            pulse_count=len(high_pulses),
            confidence=max(0.0, min(1.0, confidence)),  # Clamp to [0, 1]
            method='kmeans'
        )

    def _extract_histogram(
        self,
        high_pulses: List[int],
        low_pulses: List[int]
    ) -> Optional[TimingCharacteristics]:
        """Extract timing using histogram peak detection"""
        if len(high_pulses) < 10:
            return None

        # Create histogram
        hist, bin_edges = np.histogram(high_pulses, bins=min(50, len(high_pulses) // 5))

        # Find peaks
        try:
            from scipy.signal import find_peaks
            peaks, properties = find_peaks(hist, height=len(high_pulses) * 0.05)
        except ImportError:
            # Fallback: simple peak detection
            peaks = []
            for i in range(1, len(hist) - 1):
                if hist[i] > hist[i-1] and hist[i] > hist[i+1]:
                    peaks.append(i)

        if len(peaks) < 2:
            return None

        # Get pulse widths at peak locations
        peak_widths = [int(bin_edges[p]) for p in peaks]
        peak_widths_sorted = sorted(peak_widths)

        short_pulse = peak_widths_sorted[0]
        long_pulse = peak_widths_sorted[1]

        # Extract gaps
        short_gap, long_gap = self._extract_gaps_percentile(low_pulses)

        # Confidence: ratio of peak heights to baseline
        peak_heights = [hist[p] for p in peaks[:2]]
        baseline = np.median(hist)
        snr = np.mean(peak_heights) / (baseline + 1)
        confidence = min(1.0, snr / 10)

        # Calculate metrics
        pulse_ratio = long_pulse / short_pulse if short_pulse > 0 else 0
        gap_ratio = long_gap / short_gap if short_gap > 0 else 0
        duty_cycle = self._calculate_duty_cycle(high_pulses, low_pulses)

        return TimingCharacteristics(
            short_pulse_us=short_pulse,
            long_pulse_us=long_pulse,
            short_gap_us=short_gap,
            long_gap_us=long_gap,
            pulse_ratio=pulse_ratio,
            gap_ratio=gap_ratio,
            duty_cycle=duty_cycle,
            pulse_count=len(high_pulses),
            confidence=confidence,
            method='histogram'
        )

    def _extract_percentile(
        self,
        high_pulses: List[int],
        low_pulses: List[int]
    ) -> TimingCharacteristics:
        """Extract timing using percentile-based approach (fallback)"""
        # Use percentiles that are robust to outliers (15th/85th instead of min/max)
        short_pulse = int(np.percentile(high_pulses, 15))
        long_pulse = int(np.percentile(high_pulses, 85))

        # Extract gaps
        short_gap, long_gap = self._extract_gaps_percentile(low_pulses)

        # Lower confidence for percentile method
        confidence = 0.6

        # Calculate metrics
        pulse_ratio = long_pulse / short_pulse if short_pulse > 0 else 0
        gap_ratio = long_gap / short_gap if short_gap > 0 else 0
        duty_cycle = self._calculate_duty_cycle(high_pulses, low_pulses)

        return TimingCharacteristics(
            short_pulse_us=short_pulse,
            long_pulse_us=long_pulse,
            short_gap_us=short_gap,
            long_gap_us=long_gap,
            pulse_ratio=pulse_ratio,
            gap_ratio=gap_ratio,
            duty_cycle=duty_cycle,
            pulse_count=len(high_pulses),
            confidence=confidence,
            method='percentile'
        )

    def _extract_gaps_percentile(self, low_pulses: List[int]) -> Tuple[int, int]:
        """Extract gap timings using percentile method"""
        if not low_pulses:
            return (0, 0)

        short_gap = int(np.percentile(low_pulses, 25))
        long_gap = int(np.percentile(low_pulses, 75))

        return (short_gap, long_gap)

    def _calculate_duty_cycle(
        self,
        high_pulses: List[int],
        low_pulses: List[int]
    ) -> float:
        """Calculate duty cycle (fraction of time signal is HIGH)"""
        total_high = sum(high_pulses)
        total_low = sum(low_pulses)
        total_time = total_high + total_low

        if total_time == 0:
            return 0.0

        return total_high / total_time

    def _error_to_score(self, error: float, tolerance: float) -> float:
        """
        Convert error percentage to similarity score

        Uses smooth falloff within tolerance window:
          - error = 0 → score = 1.0
          - error = tolerance → score = 0.5
          - error = 2*tolerance → score = 0.0

        Args:
            error: Error as fraction (0.2 = 20% error)
            tolerance: Tolerance threshold

        Returns:
            Score between 0.0 and 1.0
        """
        if error <= tolerance:
            # Linear falloff within tolerance
            return 1.0 - (error / tolerance) * 0.5
        elif error <= 2 * tolerance:
            # Continue falloff to 0
            normalized = (error - tolerance) / tolerance
            return 0.5 * (1.0 - normalized)
        else:
            return 0.0

    def _empty_timing(self) -> TimingCharacteristics:
        """Return empty timing characteristics"""
        return TimingCharacteristics(
            short_pulse_us=0,
            long_pulse_us=0,
            short_gap_us=0,
            long_gap_us=0,
            pulse_ratio=0.0,
            gap_ratio=0.0,
            duty_cycle=0.0,
            pulse_count=0,
            confidence=0.0,
            method='none'
        )


# === Strategy Integration ===

class TimingMatchStrategy:
    """
    Matching strategy that uses timing analyzer

    Integrates with engine.py's MatchStrategy interface
    """

    def __init__(self, protocol_db):
        """
        Initialize strategy

        Args:
            protocol_db: Protocol database instance
        """
        self.protocol_db = protocol_db
        self.analyzer = RobustTimingAnalyzer()

    def match(self, metadata, db) -> List:
        """
        Match signal metadata using timing analysis

        Args:
            metadata: SignalMetadata object
            db: Database connection (unused, uses protocol_db)

        Returns:
            List of MatchResult objects
        """
        from ..matcher.engine import MatchResult

        if not metadata.has_raw_data:
            return []

        # Extract timing from signal
        timing = self.analyzer.extract_timing(metadata.raw_data)

        if timing.confidence < 0.3:  # Too noisy
            return []

        # Score against all protocols
        matches = []
        for protocol in self.protocol_db.get_all():
            score = self.analyzer.score_against_protocol(
                timing=timing,
                protocol_short_us=protocol.short_pulse_us,
                protocol_long_us=protocol.long_pulse_us,
                protocol_frequency=protocol.frequency,
                signal_frequency=metadata.frequency,
                timing_tolerance=protocol.timing_tolerance
            )

            # Only keep reasonable matches
            if score.overall_score >= protocol.min_confidence:
                matches.append(MatchResult(
                    device_id=0,  # Will be filled by database lookup
                    device_name=protocol.name,
                    manufacturer=protocol.manufacturer or "Unknown",
                    confidence=score.overall_score,
                    match_method='timing',
                    match_details={
                        'pulse_score': score.pulse_score,
                        'ratio_score': score.ratio_score,
                        'frequency_score': score.frequency_score,
                        'timing_method': timing.method,
                        'short_pulse_us': timing.short_pulse_us,
                        'long_pulse_us': timing.long_pulse_us,
                        'pulse_ratio': f"{timing.pulse_ratio:.2f}",
                        **score.details
                    }
                ))

        # Sort by confidence
        matches.sort(key=lambda m: m.confidence, reverse=True)

        return matches[:20]  # Return top 20


# Singleton instance
_analyzer: Optional[RobustTimingAnalyzer] = None


def get_timing_analyzer() -> RobustTimingAnalyzer:
    """Get singleton timing analyzer instance"""
    global _analyzer
    if _analyzer is None:
        _analyzer = RobustTimingAnalyzer()
    return _analyzer
