#!/usr/bin/env python3
"""
Statistical Classifier for RF Device Identification

Lightweight Bayesian classifier using protocol database as ground truth.
No ML dependencies required - uses pure statistical methods.

Formula: P(device|features) ∝ P(features|device) * P(device)
"""

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

from src.matcher.protocol_database import ProtocolSignature, get_protocol_database
from src.matcher.timing_analyzer import TimingCharacteristics


@dataclass
class FeatureVector:
    """
    Feature representation of RF signal

    Attributes:
        timing_ratio: long_pulse / short_pulse ratio
        frequency_band: ISM band identifier (0=315MHz, 1=433MHz, 2=868MHz, 3=915MHz)
        preamble_type: Preamble pattern type (0=none, 1=long_burst, 2=alternating, 3=sync_word, 4=custom)
        bit_length: Number of bits in transmission
        pulse_count: Total number of pulses
        duty_cycle: Ratio of HIGH to total time
    """
    timing_ratio: float
    frequency_band: int
    preamble_type: int
    bit_length: int
    pulse_count: int
    duty_cycle: float

    def to_array(self) -> np.ndarray:
        """Convert to numpy array for distance calculations"""
        return np.array([
            self.timing_ratio,
            self.frequency_band,
            self.preamble_type,
            self.bit_length / 100.0,  # Normalize to 0-1 range
            self.pulse_count / 100.0,  # Normalize
            self.duty_cycle
        ])


@dataclass
class StatisticalScore:
    """Result from statistical classifier"""
    protocol_name: str
    likelihood: float  # P(features|device)
    prior: float  # P(device)
    posterior: float  # P(device|features)
    feature_distance: float  # Euclidean distance in feature space
    confidence: float  # Overall confidence (0-1)


class StatisticalClassifier:
    """
    Lightweight Bayesian classifier for RF device identification

    Uses protocol database as training data and computes:
    P(device|features) ∝ P(features|device) * P(device)
    """

    def __init__(self, protocol_db=None):
        self.protocol_db = protocol_db or get_protocol_database()

        # Frequency band mapping
        self.frequency_bands = {
            (314_000_000, 316_000_000): 0,  # 315 MHz
            (433_050_000, 434_790_000): 1,  # 433 MHz
            (868_000_000, 868_600_000): 2,  # 868 MHz
            (902_000_000, 928_000_000): 3,  # 915 MHz
        }

        # Preamble type mapping
        self.preamble_types = {
            'none': 0,
            'long_burst': 1,
            'alternating': 2,
            'sync_word': 3,
            'custom': 4,
        }

        # Build statistical model from protocol database
        self._build_model()

    def _build_model(self):
        """Build statistical model from protocol database"""
        protocols = self.protocol_db.get_all()

        # Count protocol occurrences (for prior probabilities)
        self.protocol_counts = Counter()
        self.protocol_features = {}

        for proto in protocols:
            # Prior probability (uniform for now, could weight by popularity)
            self.protocol_counts[proto.name] += 1

            # Extract features from protocol signature
            features = self._extract_protocol_features(proto)
            self.protocol_features[proto.name] = features

        # Calculate priors
        total_protocols = sum(self.protocol_counts.values())
        self.priors = {
            name: count / total_protocols
            for name, count in self.protocol_counts.items()
        }

    def _extract_protocol_features(self, proto: ProtocolSignature) -> FeatureVector:
        """Extract feature vector from protocol signature"""
        # Timing ratio
        timing_ratio = proto.long_pulse_us / proto.short_pulse_us if proto.short_pulse_us > 0 else 2.0

        # Frequency band
        frequency_band = self._map_frequency_to_band(proto.frequency)

        # Preamble type (infer from pattern)
        preamble_type = self._infer_preamble_type(proto.preamble_pattern)

        # Bit length (use midpoint of range)
        bit_length = (proto.min_bits + proto.max_bits) // 2

        # Pulse count estimate
        pulse_count = proto.typical_pulse_count

        # Duty cycle estimate (assume 0.5 for PWM)
        duty_cycle = 0.5

        return FeatureVector(
            timing_ratio=timing_ratio,
            frequency_band=frequency_band,
            preamble_type=preamble_type,
            bit_length=bit_length,
            pulse_count=pulse_count,
            duty_cycle=duty_cycle
        )

    def _map_frequency_to_band(self, frequency: int) -> int:
        """Map frequency to ISM band identifier"""
        for (low, high), band_id in self.frequency_bands.items():
            if low <= frequency <= high:
                return band_id
        return 1  # Default to 433 MHz

    def _infer_preamble_type(self, preamble_pattern: Optional[str]) -> int:
        """Infer preamble type from pattern string"""
        if not preamble_pattern:
            return 0  # none

        pattern = preamble_pattern.lower()

        if '1111' in pattern or 'burst' in pattern:
            return 1  # long_burst
        elif '1010' in pattern or '0101' in pattern:
            return 2  # alternating
        elif 'sync' in pattern or '1000' in pattern:
            return 3  # sync_word
        else:
            return 4  # custom

    def extract_signal_features(
        self,
        timing: TimingCharacteristics,
        frequency: int,
        preamble_type: str,
        bit_count: int,
        pulse_count: int
    ) -> FeatureVector:
        """
        Extract feature vector from signal data

        Args:
            timing: Extracted timing characteristics
            frequency: Signal frequency in Hz
            preamble_type: Detected preamble type
            bit_count: Number of decoded bits
            pulse_count: Total pulse count

        Returns:
            FeatureVector for classification
        """
        timing_ratio = timing.pulse_ratio
        frequency_band = self._map_frequency_to_band(frequency)
        preamble_id = self.preamble_types.get(preamble_type, 0)
        duty_cycle = timing.duty_cycle

        return FeatureVector(
            timing_ratio=timing_ratio,
            frequency_band=frequency_band,
            preamble_type=preamble_id,
            bit_length=bit_count,
            pulse_count=pulse_count,
            duty_cycle=duty_cycle
        )

    def classify(
        self,
        signal_features: FeatureVector,
        candidate_protocols: List[ProtocolSignature]
    ) -> List[StatisticalScore]:
        """
        Classify signal using Bayesian approach

        Args:
            signal_features: Extracted features from signal
            candidate_protocols: Pre-filtered protocol candidates

        Returns:
            List of StatisticalScore sorted by posterior probability
        """
        scores = []
        signal_array = signal_features.to_array()

        for proto in candidate_protocols:
            # Get protocol features
            if proto.name not in self.protocol_features:
                continue

            proto_features = self.protocol_features[proto.name]
            proto_array = proto_features.to_array()

            # Calculate feature distance (Euclidean)
            distance = np.linalg.norm(signal_array - proto_array)

            # Likelihood: P(features|device) using Gaussian
            # Closer features = higher likelihood
            sigma = 2.0  # Standard deviation for Gaussian kernel
            likelihood = np.exp(-(distance ** 2) / (2 * sigma ** 2))

            # Prior: P(device)
            prior = self.priors.get(proto.name, 1.0 / len(self.protocol_features))

            # Posterior: P(device|features) ∝ P(features|device) * P(device)
            posterior = likelihood * prior

            # Confidence (normalized posterior)
            confidence = min(1.0, posterior * 10)  # Scale for reasonable range

            scores.append(StatisticalScore(
                protocol_name=proto.name,
                likelihood=likelihood,
                prior=prior,
                posterior=posterior,
                feature_distance=distance,
                confidence=confidence
            ))

        # Sort by posterior probability (highest first)
        return sorted(scores, key=lambda s: s.posterior, reverse=True)

    def get_feature_similarity(
        self,
        features1: FeatureVector,
        features2: FeatureVector
    ) -> float:
        """
        Calculate similarity between two feature vectors

        Returns:
            Similarity score (0-1, higher = more similar)
        """
        arr1 = features1.to_array()
        arr2 = features2.to_array()

        distance = np.linalg.norm(arr1 - arr2)

        # Convert distance to similarity (0-1)
        sigma = 2.0
        similarity = np.exp(-(distance ** 2) / (2 * sigma ** 2))

        return float(similarity)


# Singleton instance
_classifier: Optional[StatisticalClassifier] = None


def get_statistical_classifier() -> StatisticalClassifier:
    """Get singleton statistical classifier instance"""
    global _classifier
    if _classifier is None:
        _classifier = StatisticalClassifier()
    return _classifier


if __name__ == '__main__':
    # Test statistical classifier
    from src.matcher.timing_analyzer import TimingCharacteristics

    print("=== Statistical Classifier Test ===")
    print()

    classifier = get_statistical_classifier()

    print(f"Loaded {len(classifier.protocol_features)} protocol signatures")
    print()

    # Test feature extraction
    test_timing = TimingCharacteristics(
        short_pulse_us=500,
        long_pulse_us=1000,
        short_gap_us=500,
        long_gap_us=500,
        pulse_ratio=2.0,
        gap_ratio=1.0,
        duty_cycle=0.5,
        pulse_count=80,
        confidence=0.9,
        method='kmeans'
    )

    signal_features = classifier.extract_signal_features(
        timing=test_timing,
        frequency=433920000,
        preamble_type='alternating',
        bit_count=40,
        pulse_count=80
    )

    print("Test Signal Features:")
    print(f"  Timing Ratio: {signal_features.timing_ratio:.2f}")
    print(f"  Frequency Band: {signal_features.frequency_band} (433 MHz)")
    print(f"  Preamble Type: {signal_features.preamble_type} (alternating)")
    print(f"  Bit Length: {signal_features.bit_length}")
    print(f"  Pulse Count: {signal_features.pulse_count}")
    print(f"  Duty Cycle: {signal_features.duty_cycle:.2%}")
    print()

    # Get candidate protocols (433 MHz weather sensors)
    from src.matcher.protocol_database import get_protocol_database
    db = get_protocol_database()
    candidates = db.find_by_frequency(433920000)[:20]  # Top 20

    print(f"Testing against {len(candidates)} candidate protocols...")

    # Classify
    scores = classifier.classify(signal_features, candidates)

    print()
    print("Top 5 Statistical Matches:")
    for i, score in enumerate(scores[:5], 1):
        print(f"{i}. {score.protocol_name}")
        print(f"   Likelihood: {score.likelihood:.4f}")
        print(f"   Prior: {score.prior:.6f}")
        print(f"   Posterior: {score.posterior:.6f}")
        print(f"   Distance: {score.feature_distance:.3f}")
        print(f"   Confidence: {score.confidence:.1%}")
        print()
