#!/usr/bin/env python3
"""
RF Device Identification Benchmarking Suite

Measures accuracy, precision, recall, and confidence distribution
across known protocol test signals.
"""

import sys
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass, field
from collections import defaultdict
import json

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.parser.sub_parser import parse_sub_file
from src.matcher.pattern_decoder import get_pattern_decoder, DeviceMatch


@dataclass
class BenchmarkResult:
    """Result for a single test case"""
    test_file: str
    expected_protocol: str
    expected_metadata: Dict

    # Results
    top_match: Optional[str] = None
    top_confidence: float = 0.0
    all_matches: List[DeviceMatch] = field(default_factory=list)

    # Metrics
    correct: bool = False
    rank: int = -1  # Rank of correct protocol in results (1-indexed, -1 if not found)

    # Performance
    parse_time_ms: float = 0.0
    match_time_ms: float = 0.0


@dataclass
class ProtocolStats:
    """Statistics for a specific protocol"""
    protocol_name: str
    total_tests: int = 0
    correct_top1: int = 0
    correct_top3: int = 0
    correct_top5: int = 0
    not_found: int = 0

    avg_confidence_when_correct: float = 0.0
    avg_confidence_when_wrong: float = 0.0
    avg_rank: float = 0.0

    confidences_correct: List[float] = field(default_factory=list)
    confidences_wrong: List[float] = field(default_factory=list)


@dataclass
class BenchmarkSummary:
    """Overall benchmark statistics"""
    total_tests: int = 0

    # Top-K accuracy
    top1_accuracy: float = 0.0
    top3_accuracy: float = 0.0
    top5_accuracy: float = 0.0

    # Confidence distribution
    high_confidence_count: int = 0  # >80%
    medium_confidence_count: int = 0  # 50-80%
    low_confidence_count: int = 0  # <50%

    # Per-protocol stats
    protocol_stats: Dict[str, ProtocolStats] = field(default_factory=dict)

    # Performance
    avg_parse_time_ms: float = 0.0
    avg_match_time_ms: float = 0.0

    # All results
    results: List[BenchmarkResult] = field(default_factory=list)


class RFBenchmark:
    """RF device identification benchmark runner"""

    def __init__(self):
        self.decoder = get_pattern_decoder()

    def run_benchmark(
        self,
        test_cases: List[Tuple[Path, str, Dict]]
    ) -> BenchmarkSummary:
        """
        Run benchmark on test cases

        Args:
            test_cases: List of (filepath, expected_protocol, metadata)

        Returns:
            BenchmarkSummary with all metrics
        """
        summary = BenchmarkSummary()
        summary.total_tests = len(test_cases)

        print(f"Running benchmark on {len(test_cases)} test cases...")
        print()

        for i, (filepath, expected_protocol, metadata) in enumerate(test_cases, 1):
            print(f"[{i}/{len(test_cases)}] Testing {filepath.name}...", end=" ")

            result = self._test_single_file(filepath, expected_protocol, metadata)
            summary.results.append(result)

            # Update protocol stats
            if expected_protocol not in summary.protocol_stats:
                summary.protocol_stats[expected_protocol] = ProtocolStats(
                    protocol_name=expected_protocol
                )

            stats = summary.protocol_stats[expected_protocol]
            stats.total_tests += 1

            if result.correct and result.rank == 1:
                stats.correct_top1 += 1
                stats.confidences_correct.append(result.top_confidence)
                print(f"✓ CORRECT (confidence: {result.top_confidence:.1%})")
            elif result.rank > 0 and result.rank <= 3:
                stats.correct_top3 += 1
                print(f"⚠ Rank {result.rank} (expected in top 3)")
            elif result.rank > 0 and result.rank <= 5:
                stats.correct_top5 += 1
                print(f"⚠ Rank {result.rank} (expected in top 5)")
            else:
                stats.not_found += 1
                stats.confidences_wrong.append(result.top_confidence)
                print(f"✗ WRONG: {result.top_match} (confidence: {result.top_confidence:.1%})")

        print()

        # Calculate aggregate metrics
        self._calculate_summary_metrics(summary)

        return summary

    def _test_single_file(
        self,
        filepath: Path,
        expected_protocol: str,
        metadata: Dict
    ) -> BenchmarkResult:
        """Test a single .sub file"""
        import time

        result = BenchmarkResult(
            test_file=str(filepath),
            expected_protocol=expected_protocol,
            expected_metadata=metadata
        )

        # Parse file
        t0 = time.time()
        try:
            signal_metadata = parse_sub_file(str(filepath))
        except Exception as e:
            print(f"Parse error: {e}")
            return result
        result.parse_time_ms = (time.time() - t0) * 1000

        # Run matching
        t0 = time.time()
        try:
            matches = self.decoder.decode(signal_metadata)
            result.all_matches = matches
        except Exception as e:
            print(f"Match error: {e}")
            return result
        result.match_time_ms = (time.time() - t0) * 1000

        if not matches:
            return result

        # Extract top match
        result.top_match = matches[0].name
        result.top_confidence = matches[0].confidence

        # Check if correct protocol is in results
        for rank, match in enumerate(matches, 1):
            if self._protocol_matches(match.name, expected_protocol):
                result.correct = (rank == 1)
                result.rank = rank
                break

        return result

    def _protocol_matches(self, actual: str, expected: str) -> bool:
        """Check if protocol names match (fuzzy matching)"""
        actual_lower = actual.lower()
        expected_lower = expected.lower()

        # Exact match
        if actual_lower == expected_lower:
            return True

        # Partial match (handle variations like "LaCrosse TX141-BV2" vs "LaCrosse-TX141TH-BV2")
        actual_parts = set(actual_lower.replace('-', ' ').split())
        expected_parts = set(expected_lower.replace('-', ' ').split())

        # If 2+ common tokens, consider match
        common = actual_parts & expected_parts
        if len(common) >= 2:
            return True

        return False

    def _calculate_summary_metrics(self, summary: BenchmarkSummary):
        """Calculate aggregate metrics from individual results"""
        if summary.total_tests == 0:
            return

        # Top-K accuracy
        top1_correct = sum(1 for r in summary.results if r.correct and r.rank == 1)
        top3_correct = sum(1 for r in summary.results if r.rank > 0 and r.rank <= 3)
        top5_correct = sum(1 for r in summary.results if r.rank > 0 and r.rank <= 5)

        summary.top1_accuracy = top1_correct / summary.total_tests
        summary.top3_accuracy = top3_correct / summary.total_tests
        summary.top5_accuracy = top5_correct / summary.total_tests

        # Confidence distribution
        for result in summary.results:
            if result.top_confidence >= 0.8:
                summary.high_confidence_count += 1
            elif result.top_confidence >= 0.5:
                summary.medium_confidence_count += 1
            else:
                summary.low_confidence_count += 1

        # Performance
        if summary.results:
            summary.avg_parse_time_ms = sum(r.parse_time_ms for r in summary.results) / len(summary.results)
            summary.avg_match_time_ms = sum(r.match_time_ms for r in summary.results) / len(summary.results)

        # Per-protocol stats
        for protocol, stats in summary.protocol_stats.items():
            if stats.total_tests > 0:
                stats.correct_top3 = stats.correct_top1  # Already counted in loop
                stats.correct_top5 = stats.correct_top1

            if stats.confidences_correct:
                stats.avg_confidence_when_correct = sum(stats.confidences_correct) / len(stats.confidences_correct)

            if stats.confidences_wrong:
                stats.avg_confidence_when_wrong = sum(stats.confidences_wrong) / len(stats.confidences_wrong)

    def print_summary(self, summary: BenchmarkSummary):
        """Print human-readable summary"""
        print("=" * 80)
        print("BENCHMARK SUMMARY")
        print("=" * 80)
        print()

        print(f"Total Tests: {summary.total_tests}")
        print()

        print("TOP-K ACCURACY:")
        print(f"  Top-1: {summary.top1_accuracy:.1%} ({int(summary.top1_accuracy * summary.total_tests)}/{summary.total_tests})")
        print(f"  Top-3: {summary.top3_accuracy:.1%} ({int(summary.top3_accuracy * summary.total_tests)}/{summary.total_tests})")
        print(f"  Top-5: {summary.top5_accuracy:.1%} ({int(summary.top5_accuracy * summary.total_tests)}/{summary.total_tests})")
        print()

        print("CONFIDENCE DISTRIBUTION:")
        print(f"  High (>80%):   {summary.high_confidence_count} ({summary.high_confidence_count/summary.total_tests:.1%})")
        print(f"  Medium (50-80%): {summary.medium_confidence_count} ({summary.medium_confidence_count/summary.total_tests:.1%})")
        print(f"  Low (<50%):    {summary.low_confidence_count} ({summary.low_confidence_count/summary.total_tests:.1%})")
        print()

        print("PERFORMANCE:")
        print(f"  Avg Parse Time: {summary.avg_parse_time_ms:.2f} ms")
        print(f"  Avg Match Time: {summary.avg_match_time_ms:.2f} ms")
        print(f"  Total Time:     {summary.avg_parse_time_ms + summary.avg_match_time_ms:.2f} ms")
        print()

        print("PER-PROTOCOL ACCURACY:")
        print(f"{'Protocol':<40} {'Tests':>6} {'Accuracy':>10} {'Avg Conf':>10}")
        print("-" * 80)

        # Sort by accuracy (worst first for tuning focus)
        sorted_protocols = sorted(
            summary.protocol_stats.items(),
            key=lambda x: x[1].correct_top1 / x[1].total_tests if x[1].total_tests > 0 else 0
        )

        for protocol_name, stats in sorted_protocols:
            accuracy = stats.correct_top1 / stats.total_tests if stats.total_tests > 0 else 0
            avg_conf = stats.avg_confidence_when_correct if stats.avg_confidence_when_correct > 0 else stats.avg_confidence_when_wrong

            print(f"{protocol_name:<40} {stats.total_tests:>6} {accuracy:>9.1%} {avg_conf:>9.1%}")

        print()

        # Worst performers
        print("WORST PERFORMERS (need tuning):")
        worst = sorted_protocols[:5]
        for protocol_name, stats in worst:
            accuracy = stats.correct_top1 / stats.total_tests if stats.total_tests > 0 else 0
            if accuracy < 1.0:
                print(f"  - {protocol_name}: {accuracy:.1%} ({stats.correct_top1}/{stats.total_tests})")

        print()

    def save_report(self, summary: BenchmarkSummary, output_path: Path):
        """Save detailed report to markdown"""
        with open(output_path, 'w') as f:
            f.write("# RF Device Identification - Benchmark Results\n\n")

            f.write(f"**Date**: {Path(__file__).stat().st_mtime}\n")
            f.write(f"**Total Tests**: {summary.total_tests}\n\n")

            f.write("## Overall Metrics\n\n")
            f.write("### Top-K Accuracy\n\n")
            f.write(f"- **Top-1**: {summary.top1_accuracy:.1%}\n")
            f.write(f"- **Top-3**: {summary.top3_accuracy:.1%}\n")
            f.write(f"- **Top-5**: {summary.top5_accuracy:.1%}\n\n")

            f.write("### Confidence Distribution\n\n")
            f.write(f"- **High (>80%)**: {summary.high_confidence_count} ({summary.high_confidence_count/summary.total_tests:.1%})\n")
            f.write(f"- **Medium (50-80%)**: {summary.medium_confidence_count} ({summary.medium_confidence_count/summary.total_tests:.1%})\n")
            f.write(f"- **Low (<50%)**: {summary.low_confidence_count} ({summary.low_confidence_count/summary.total_tests:.1%})\n\n")

            f.write("### Performance\n\n")
            f.write(f"- **Avg Parse Time**: {summary.avg_parse_time_ms:.2f} ms\n")
            f.write(f"- **Avg Match Time**: {summary.avg_match_time_ms:.2f} ms\n")
            f.write(f"- **Total**: {summary.avg_parse_time_ms + summary.avg_match_time_ms:.2f} ms\n\n")

            f.write("## Per-Protocol Results\n\n")
            f.write("| Protocol | Tests | Top-1 Acc | Avg Confidence |\n")
            f.write("|----------|-------|-----------|----------------|\n")

            sorted_protocols = sorted(
                summary.protocol_stats.items(),
                key=lambda x: x[1].correct_top1 / x[1].total_tests if x[1].total_tests > 0 else 0
            )

            for protocol_name, stats in sorted_protocols:
                accuracy = stats.correct_top1 / stats.total_tests if stats.total_tests > 0 else 0
                avg_conf = stats.avg_confidence_when_correct if stats.avg_confidence_when_correct > 0 else stats.avg_confidence_when_wrong

                f.write(f"| {protocol_name} | {stats.total_tests} | {accuracy:.1%} | {avg_conf:.1%} |\n")

            f.write("\n## Detailed Results\n\n")

            for result in summary.results:
                status = "✓" if result.correct else "✗"
                f.write(f"### {status} {Path(result.test_file).name}\n\n")
                f.write(f"- **Expected**: {result.expected_protocol}\n")
                f.write(f"- **Got**: {result.top_match} (confidence: {result.top_confidence:.1%})\n")
                f.write(f"- **Rank**: {result.rank if result.rank > 0 else 'Not Found'}\n")

                if result.all_matches:
                    f.write(f"\n**Top 5 Matches**:\n")
                    for i, match in enumerate(result.all_matches[:5], 1):
                        f.write(f"{i}. {match.name} ({match.confidence:.1%})\n")

                f.write("\n")


def main():
    """Main benchmark runner"""
    # Generate synthetic test signals
    benchmark_dir = Path(__file__).parent.parent / "tests" / "benchmark"
    sys.path.insert(0, str(benchmark_dir))
    from test_data_generator import SyntheticSignalGenerator

    print("=" * 80)
    print("RF DEVICE IDENTIFICATION BENCHMARK")
    print("=" * 80)
    print()

    # Setup
    synthetic_dir = benchmark_dir / "synthetic_signals"

    generator = SyntheticSignalGenerator(synthetic_dir)

    print("Step 1: Generating synthetic test signals...")
    test_cases = generator.generate_test_suite()
    print(f"  Generated {len(test_cases)} test signals")
    print()

    # Run benchmark
    print("Step 2: Running benchmark...")
    benchmark = RFBenchmark()
    summary = benchmark.run_benchmark(test_cases)

    # Print results
    print()
    benchmark.print_summary(summary)

    # Save report
    report_path = Path(__file__).parent.parent / "TEST_RESULTS_SUMMARY.md"
    benchmark.save_report(summary, report_path)
    print(f"Detailed report saved to: {report_path}")
    print()

    # Return exit code based on accuracy
    if summary.top1_accuracy >= 0.25:  # 25% target
        print("✓ Benchmark PASSED (Top-1 accuracy >= 25%)")
        return 0
    else:
        print(f"✗ Benchmark FAILED (Top-1 accuracy {summary.top1_accuracy:.1%} < 25%)")
        return 1


if __name__ == '__main__':
    sys.exit(main())