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

Purpose: Visualize YOLO-OBB (Oriented Bounding Box) detection results on grass images

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
- ✅ Oriented bounding boxes with rotation angles
"""

# 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_obb_boxes(img, obb_boxes, conf_threshold=0.3):
    """
    Draw oriented bounding boxes on image

    Args:
        img: BGR image
        obb_boxes: oriented bounding boxes [num_det, 7+]
                   (x1, y1, x2, y2, x3, y3, x4, y4, conf, cls) for OBB
        conf_threshold: confidence threshold

    Returns:
        img_with_obb: image with OBB drawn
    """
    img_out = img.copy()

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

    for box in obb_boxes:
        # OBB format: 4 corner points + conf + class
        if len(box) >= 6:  # At minimum: x, y, w, h, conf, cls
            # Get last 2 values as conf and cls
            conf = box[-2]
            cls = int(box[-1])

            if conf < conf_threshold:
                continue

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

            # Try to extract OBB points
            if len(box) >= 10:  # Full OBB format: 4 corners (x1,y1,x2,y2,x3,y3,x4,y4) + conf + cls
                # Extract 4 corner points
                points = np.array([
                    [int(box[0]), int(box[1])],
                    [int(box[2]), int(box[3])],
                    [int(box[4]), int(box[5])],
                    [int(box[6]), int(box[7])]
                ], dtype=np.int32)

                # Draw oriented box
                cv2.polylines(img_out, [points], isClosed=True, color=color, thickness=3)

                # Calculate center for label
                center_x = int(np.mean([box[0], box[2], box[4], box[6]]))
                center_y = int(np.mean([box[1], box[3], box[5], box[7]]))
            else:
                # Fallback to regular bbox if OBB format not available
                x, y, w, h = map(int, box[:4])
                cv2.rectangle(img_out, (x, y), (x+w, y+h), color, 3)
                center_x, center_y = x, y

            # Add label
            label = f"Cls{cls}: {conf:.2f}"
            (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
            cv2.rectangle(img_out, (center_x, center_y-h-10), (center_x+w+10, center_y), color, -1)
            cv2.putText(img_out, label, (center_x+5, center_y-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-OBB from OFFLINE source
model = YOLO('/kaggle/input/yolo8-obb/yolov8n-obb.pt')
print(f"✓ YOLO-OBB 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 OBB 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-OBB
        results = model(full_path, verbose=False, conf=conf)
        result = results[0]

        # Get OBB boxes
        obb_data = None
        num_det = 0

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

        # Draw OBB boxes
        img_obb = draw_obb_boxes(img_bgr, obb_data, conf_threshold=0.3)
        img_obb_rgb = cv2.cvtColor(img_obb, cv2.COLOR_BGR2RGB)

        # Display
        axes[i+1].imshow(img_obb_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 obb_data is not None:
            avg_conf = obb_data[:, -2].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'obb_viz_{idx+1}.png', dpi=300, bbox_inches='tight')
    plt.show()
    plt.close()

print(f"\n✓ Saved {num_viz} OBB 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

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

            if num_det > 0:
                avg_conf = obb_data[:, -2].mean()
                max_conf = obb_data[:, -2].max()

                # Record class distribution
                for box in obb_data:
                    cls = int(box[-1])
                    class_distribution.append({
                        'threshold': conf_name,
                        'image': img_path,
                        'class': cls,
                        'confidence': box[-2]
                    })

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

# 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']
}).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('obb_detection_stats.csv', index=False)
if len(class_df) > 0:
    class_df.to_csv('obb_class_distribution.csv', index=False)

print("\n✓ Saved: obb_detection_stats.csv")
if len(class_df) > 0:
    print("✓ Saved: obb_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✓ Exp20 YOLO-OBB EDA completed!")