#!/usr/bin/env python3
"""
Exp19: YOLO-Detection EDA (Version 1: Large-scale Analysis)

Purpose: Visualize YOLO-Det object detection results on grass images with multiple thresholds

Thresholds: conf=0.001 (0.1%), 0.003 (0.3%), 0.005 (0.5%)

Version 1 Features:
- ✅ Large sample size: 500 images
- ✅ 100 visualizations
- ✅ Detection count classification
- ✅ Multiple confidence thresholds
- ✅ High resolution: 300dpi
- ✅ Bounding boxes with class labels
"""

# Install ultralytics offline
import os
import subprocess
subprocess.run([
    'pip', 'install', '--no-index',
    '--find-links=/kaggle/input/ultralytics810/ultralytics-minimal/wheels',
    '--no-deps', 'ultralytics', 'hub-sdk', 'thop'
], check=True)
import cv2
import numpy as np
import pandas as pd
from tqdm import tqdm
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
from pathlib import Path
import torch

plt.rcParams['figure.figsize'] = (20, 16)

# Color palette for different classes
np.random.seed(42)
COLOR_PALETTE = np.random.randint(0, 255, (80, 3), dtype=np.uint8)


def draw_detections(img, boxes, conf_threshold=0.3):
    """
    Draw bounding boxes on image

    Args:
        img: BGR image
        boxes: bounding boxes [num_det, 6] (x1, y1, x2, y2, conf, cls)
        conf_threshold: confidence threshold

    Returns:
        img_with_boxes: image with boxes drawn
    """
    img_out = img.copy()

    if boxes is None or len(boxes) == 0:
        return img_out

    for box in boxes:
        x1, y1, x2, y2, conf, cls = box
        cls = int(cls)

        if conf < conf_threshold:
            continue

        # Get color for this class
        color = COLOR_PALETTE[cls % len(COLOR_PALETTE)].tolist()

        # Draw bounding box
        x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
        cv2.rectangle(img_out, (x1, y1), (x2, y2), color, 3)

        # Add label
        label = f"Cls{cls}: {conf:.2f}"
        (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
        cv2.rectangle(img_out, (x1, y1-h-10), (x1+w+10, y1), color, -1)
        cv2.putText(img_out, label, (x1+5, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

    return img_out


# PyTorch 2.6 patch
def patch_ultralytics():
    try:
        from ultralytics.nn import tasks
        def patched_torch_safe_load(weight):
            from ultralytics.utils.downloads import attempt_download_asset
            weight = attempt_download_asset(str(weight))
            file = Path(weight)
            ckpt = torch.load(file, map_location='cpu', weights_only=False)
            return ckpt, file
        tasks.torch_safe_load = patched_torch_safe_load
        print("✓ Patch applied")
    except: pass

patch_ultralytics()

# Load YOLO-Det from OFFLINE source
model = YOLO('/kaggle/input/yolo8-detection/yolov8n.pt')
print(f"✓ YOLO-Det loaded from offline dataset!")
print(f"Task: {model.task}")

# Load data
ROOT = "/kaggle/input/csiro-biomass/"
train_df = pd.read_csv(os.path.join(ROOT, "train.csv"))
unique_images = train_df['image_path'].unique()

# Sample 500 images
num_samples = 500
np.random.seed(42)
samples = np.random.choice(unique_images, size=min(num_samples, len(unique_images)), replace=False)
print(f"✓ Sampled {len(samples)} images (total: {len(unique_images)})")

# =============================================================================
# Visualization: Display detections for first 100 images
# =============================================================================

thresholds = [0.001, 0.003, 0.005]
threshold_names = ['0.1%', '0.3%', '0.5%']

num_viz = min(100, len(samples))

for idx in range(num_viz):
    img_path = samples[idx]
    full_path = os.path.join(ROOT, img_path)

    print(f"\n{'='*80}")
    print(f"Sample {idx+1}/{num_viz}: {img_path}")
    print(f"{'='*80}")

    # Load original image
    img_bgr = cv2.imread(full_path)
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    # Process each threshold
    fig, axes = plt.subplots(1, 4, figsize=(40, 10))

    # Original
    axes[0].imshow(img_rgb)
    axes[0].set_title('Original', fontsize=20, fontweight='bold')
    axes[0].axis('off')

    for i, (conf, conf_name) in enumerate(zip(thresholds, threshold_names)):
        # Run YOLO-Det
        results = model(full_path, verbose=False, conf=conf)
        result = results[0]

        # Get boxes
        boxes_data = None
        num_det = 0

        if result.boxes is not None and result.boxes.data is not None:
            boxes_data = result.boxes.data.cpu().numpy()
            num_det = boxes_data.shape[0]

        # Draw detections
        img_det = draw_detections(img_bgr, boxes_data, conf_threshold=0.3)
        img_det_rgb = cv2.cvtColor(img_det, cv2.COLOR_BGR2RGB)

        # Display
        axes[i+1].imshow(img_det_rgb)
        axes[i+1].set_title(
            f'Conf={conf_name} | Det={num_det}',
            fontsize=20,
            fontweight='bold',
            color='red' if num_det > 0 else 'gray'
        )
        axes[i+1].axis('off')

        # Print detection details
        if num_det > 0 and boxes_data is not None:
            avg_conf = boxes_data[:, 4].mean()
            print(f"  Conf={conf_name:5s}: {num_det:2d} detections, avg_conf={avg_conf:.3f}")
        else:
            print(f"  Conf={conf_name:5s}: {num_det:2d} detections")

    plt.tight_layout()
    plt.savefig(f'det_viz_{idx+1}.png', dpi=300, bbox_inches='tight')
    plt.show()
    plt.close()

print(f"\n✓ Saved {num_viz} detection visualization images (high resolution: 300dpi)")

# =============================================================================
# Detailed Statistics: Analyze detection rates across all samples
# =============================================================================

stats_list = []
class_distribution = []

for img_path in tqdm(samples, desc="Collecting statistics"):
    full_path = os.path.join(ROOT, img_path)

    for conf, conf_name in zip(thresholds, threshold_names):
        results = model(full_path, verbose=False, conf=conf)
        result = results[0]

        num_det = 0
        avg_conf = 0.0
        max_conf = 0.0
        total_box_area = 0

        if result.boxes is not None and result.boxes.data is not None:
            boxes_data = result.boxes.data.cpu().numpy()
            num_det = boxes_data.shape[0]

            if num_det > 0:
                avg_conf = boxes_data[:, 4].mean()
                max_conf = boxes_data[:, 4].max()

                # Calculate total box area
                for box in boxes_data:
                    x1, y1, x2, y2 = box[:4]
                    area = (x2 - x1) * (y2 - y1)
                    total_box_area += area

                # Record class distribution
                for box in boxes_data:
                    cls = int(box[5])
                    class_distribution.append({
                        'threshold': conf_name,
                        'image': img_path,
                        'class': cls,
                        'confidence': box[4]
                    })

        stats_list.append({
            'threshold': conf_name,
            'image': img_path,
            'num_detections': num_det,
            'avg_conf': avg_conf,
            'max_conf': max_conf,
            'total_box_area': total_box_area
        })

# Convert to DataFrame
stats_df = pd.DataFrame(stats_list)
class_df = pd.DataFrame(class_distribution)

print("\n" + "="*80)
print("=== Detection Statistics by Threshold (All Samples) ===")
print("="*80)

summary = stats_df.groupby('threshold').agg({
    'num_detections': ['count', 'sum', 'mean', 'std', 'min', 'max'],
    'avg_conf': ['mean', 'std'],
    'total_box_area': ['mean', 'std', 'max']
}).round(3)

print(summary)

print("\n" + "="*80)
print("=== Detection Rate ===")
print("="*80)

for conf_name in threshold_names:
    subset = stats_df[stats_df['threshold'] == conf_name]
    detection_rate = (subset['num_detections'] > 0).sum() / len(subset) * 100
    avg_det = subset['num_detections'].mean()
    max_det = subset['num_detections'].max()

    print(f"{conf_name:5s}: {detection_rate:5.1f}% ({int((subset['num_detections'] > 0).sum())}/{len(subset)} images)")
    print(f"       Avg detections: {avg_det:.2f}, Max detections: {max_det:.0f}")

# Class distribution
if len(class_df) > 0:
    print("\n" + "="*80)
    print("=== Class Distribution ===")
    print("="*80)
    class_summary = class_df.groupby(['threshold', 'class']).size().reset_index(name='count')
    print(class_summary.pivot(index='class', columns='threshold', values='count').fillna(0))

# Save files
stats_df.to_csv('det_detection_stats.csv', index=False)
if len(class_df) > 0:
    class_df.to_csv('det_class_distribution.csv', index=False)

print("\n✓ Saved: det_detection_stats.csv")
if len(class_df) > 0:
    print("✓ Saved: det_class_distribution.csv")

# =============================================================================
# Classification by Detection Count
# =============================================================================

threshold_for_classification = '0.5%'
classification_data = stats_df[stats_df['threshold'] == threshold_for_classification].copy()

def classify_detection_count(num_det):
    if num_det == 0:
        return '0 (None)'
    elif num_det == 1:
        return '1'
    elif num_det <= 5:
        return '2-5'
    elif num_det <= 10:
        return '6-10'
    else:
        return '11+'

classification_data['category'] = classification_data['num_detections'].apply(classify_detection_count)

print("\n" + "="*80)
print(f"=== Classification by Detection Count (Threshold: {threshold_for_classification}) ===")
print("="*80)
category_counts = classification_data['category'].value_counts()
print(category_counts)

print("\n✓ Exp19 YOLO-Det EDA completed!")