#!/usr/bin/env python3
"""
Model Evaluation Script
Utvärderar tränad modell med detaljerade metrics per klass
"""

import os
import sys
import json
import argparse
from pathlib import Path
from typing import Dict, List
from collections import defaultdict

import numpy as np
import torch
from ultralytics import YOLO


def evaluate_per_class(model, dataset_path: str, split: str = "test") -> Dict:
    """Utvärdera modell per klass"""
    print(f"\n=== Per-Class Evaluation ({split}) ===\n")
    
    # Kör validering
    metrics = model.val(data=dataset_path, split=split)
    
    # Extrahera per-klass metrics
    results = {
        "overall": {
            "mAP50": float(metrics.box.map50),
            "mAP50_95": float(metrics.box.map),
            "precision": float(metrics.box.mp),
            "recall": float(metrics.box.mr),
        },
        "per_class": {}
    }
    
    # Försök extrahera per-klass AP
    if hasattr(metrics.box, 'ap50') and metrics.box.ap50 is not None:
        ap50_per_class = metrics.box.ap50
        
        # Ladda klassnamn från data.yaml
        data_dir = Path(dataset_path).parent
        data_yaml = data_dir / "data.yaml"
        
        class_names = []
        if data_yaml.exists():
            import yaml
            with open(data_yaml) as f:
                data = yaml.safe_load(f)
                class_names = data.get('names', [])
        
        for i, ap in enumerate(ap50_per_class):
            class_name = class_names[i] if i < len(class_names) else f"class_{i}"
            results["per_class"][class_name] = {
                "ap50": float(ap),
                "status": "PASS" if ap >= 0.70 else "FAIL"
            }
    
    # Skriv ut resultat
    print(f"{'Class':<25} {'AP50':>8} {'Status':>8}")
    print("-" * 45)
    
    for class_name, class_metrics in sorted(results["per_class"].items()):
        status = "✅" if class_metrics["status"] == "PASS" else "⚠️"
        print(f"{class_name:<25} {class_metrics['ap50']:>8.4f} {status:>8}")
    
    print("-" * 45)
    print(f"{'Overall':<25} {results['overall']['mAP50']:>8.4f}")
    
    return results


def generate_confusion_matrix(model, dataset_path: str) -> np.ndarray:
    """Generera confusion matrix"""
    print("\n=== Confusion Matrix ===\n")
    
    # Kör validering med confusion matrix
    metrics = model.val(data=dataset_path, split="test", plots=True)
    
    # Confusion matrix sparas automatiskt av Ultralytics
    print("  Confusion matrix saved to runs/detect/val/confusion_matrix.png")
    
    return metrics


def benchmark_inference(model, image_path: str, num_runs: int = 100) -> Dict:
    """Benchmark inference-hastighet"""
    print(f"\n=== Inference Benchmark ({num_runs} runs) ===\n")
    
    import time
    from PIL import Image
    
    # Ladda testbild
    if not os.path.exists(image_path):
        print(f"  ⚠ Test image not found: {image_path}")
        return {}
    
    img = Image.open(image_path).convert('RGB')
    
    # Warmup
    for _ in range(10):
        model(img, verbose=False)
    
    # Benchmark
    times = []
    for _ in range(num_runs):
        start = time.time()
        model(img, verbose=False)
        times.append(time.time() - start)
    
    times = np.array(times)
    
    results = {
        "mean_ms": float(np.mean(times) * 1000),
        "std_ms": float(np.std(times) * 1000),
        "min_ms": float(np.min(times) * 1000),
        "max_ms": float(np.max(times) * 1000),
        "median_ms": float(np.median(times) * 1000),
        "fps": float(1.0 / np.mean(times))
    }
    
    print(f"  Mean: {results['mean_ms']:.2f} ms")
    print(f"  Std:  {results['std_ms']:.2f} ms")
    print(f"  Min:  {results['min_ms']:.2f} ms")
    print(f"  Max:  {results['max_ms']:.2f} ms")
    print(f"  FPS:  {results['fps']:.1f}")
    
    return results


def compare_with_targets(results: Dict, targets: Dict) -> Dict:
    """Jämför resultat med mål"""
    print("\n=== Target Comparison ===\n")
    
    comparison = {}
    
    for metric, target in targets.items():
        actual = results["overall"].get(metric, 0)
        diff = actual - target
        status = "PASS" if actual >= target else "FAIL"
        
        comparison[metric] = {
            "target": target,
            "actual": actual,
            "diff": diff,
            "status": status
        }
        
        status_icon = "✅" if status == "PASS" else "⚠️"
        print(f"  {metric:<15} Target: {target:.4f}  Actual: {actual:.4f}  Diff: {diff:+.4f}  {status_icon}")
    
    return comparison


def generate_evaluation_report(
    results: Dict,
    benchmark: Dict,
    comparison: Dict,
    output_dir: str
) -> str:
    """Generera utvärderingsrapport"""
    
    report = {
        "evaluation": results,
        "benchmark": benchmark,
        "target_comparison": comparison,
        "summary": {
            "overall_pass": all(c["status"] == "PASS" for c in comparison.values()),
            "num_pass": sum(1 for c in comparison.values() if c["status"] == "PASS"),
            "num_fail": sum(1 for c in comparison.values() if c["status"] == "FAIL"),
        }
    }
    
    # Spara JSON
    report_path = Path(output_dir) / "evaluation_report.json"
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)
    
    # Spara Markdown
    md = f"""# Model Evaluation Report

## Overall Metrics

| Metric | Value |
|--------|-------|
| mAP50 | {results['overall']['mAP50']:.4f} |
| mAP50-95 | {results['overall']['mAP50_95']:.4f} |
| Precision | {results['overall']['precision']:.4f} |
| Recall | {results['overall']['recall']:.4f} |

## Target Comparison

| Metric | Target | Actual | Status |
|--------|--------|--------|--------|
"""
    
    for metric, comp in comparison.items():
        status = "✅ PASS" if comp["status"] == "PASS" else "⚠️ FAIL"
        md += f"| {metric} | {comp['target']:.4f} | {comp['actual']:.4f} | {status} |\n"
    
    md += f"""
## Per-Class AP50

| Class | AP50 | Status |
|-------|------|--------|
"""
    
    for class_name, class_metrics in sorted(results["per_class"].items()):
        status = "✅" if class_metrics["status"] == "PASS" else "⚠️"
        md += f"| {class_name} | {class_metrics['ap50']:.4f} | {status} |\n"
    
    if benchmark:
        md += f"""
## Inference Benchmark

| Metric | Value |
|--------|-------|
| Mean latency | {benchmark.get('mean_ms', 0):.2f} ms |
| Std latency | {benchmark.get('std_ms', 0):.2f} ms |
| FPS | {benchmark.get('fps', 0):.1f} |
"""
    
    md += f"""
## Summary

- **Overall:** {'✅ All targets met' if report['summary']['overall_pass'] else '⚠️ Some targets not met'}
- **Passed:** {report['summary']['num_pass']}/{len(comparison)} metrics
- **Failed:** {report['summary']['num_fail']}/{len(comparison)} metrics
"""
    
    md_path = Path(output_dir) / "EVALUATION_REPORT.md"
    with open(md_path, "w") as f:
        f.write(md)
    
    return str(md_path)


def main():
    parser = argparse.ArgumentParser(description="Evaluate AI Model")
    parser.add_argument("--model", required=True, help="Path to trained model")
    parser.add_argument("--dataset", required=True, help="Path to dataset YAML")
    parser.add_argument("--output", default="/tmp/quixzoom_evaluation", help="Output directory")
    parser.add_argument("--benchmark", action="store_true", help="Run inference benchmark")
    parser.add_argument("--benchmark-image", help="Image for benchmark")
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("AI Model Evaluation")
    print("=" * 60)
    
    # Ladda modell
    print(f"\nLoading model: {args.model}")
    model = YOLO(args.model)
    
    # Skapa output-katalog
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Utvärdera per klass
    results = evaluate_per_class(model, args.dataset)
    
    # Generera confusion matrix
    generate_confusion_matrix(model, args.dataset)
    
    # Benchmark
    benchmark = {}
    if args.benchmark:
        benchmark = benchmark_inference(model, args.benchmark_image or "")
    
    # Jämför med mål
    targets = {
        "mAP50": 0.85,
        "mAP50_95": 0.70,
        "precision": 0.88,
        "recall": 0.85
    }
    comparison = compare_with_targets(results, targets)
    
    # Generera rapport
    report_path = generate_evaluation_report(results, benchmark, comparison, output_dir)
    
    # Sammanfattning
    print("\n" + "=" * 60)
    print("Evaluation Complete!")
    print("=" * 60)
    print(f"\nReport: {report_path}")
    print(f"\nmAP50: {results['overall']['mAP50']:.4f}")
    print(f"mAP50-95: {results['overall']['mAP50_95']:.4f}")
    print(f"Precision: {results['overall']['precision']:.4f}")
    print(f"Recall: {results['overall']['recall']:.4f}")
    
    return 0


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