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

Purpose: Visualize YOLO-Seg segmentation 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
- ✅ Segmentation masks with different colors per class
"""

# 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 (random bright colors)
np.random.seed(42)
COLOR_PALETTE = np.random.randint(0, 255, (80, 3), dtype=np.uint8)


def draw_segmentation_masks(img, masks, boxes, conf_threshold=0.3):
    """
    Draw segmentation masks on image

    Args:
        img: BGR image
        masks: segmentation masks [num_det, H, W]
        boxes: bounding boxes with class and conf [num_det, 6] (x1, y1, x2, y2, conf, cls)
        conf_threshold: confidence threshold

    Returns:
        img_with_masks: image with masks overlaid
    """
    img_out = img.copy()

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

    # Create overlay for transparency
    overlay = img_out.copy()

    for i, (mask, box) in enumerate(zip(masks, boxes)):
        conf = box[4]
        cls = int(box[5])

        if conf < conf_threshold:
            continue

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

        # Resize mask to original image size
        mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]))
        mask_bool = mask_resized > 0.5

        # Apply colored mask
        overlay[mask_bool] = color

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

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

    # Blend overlay with original
    img_out = cv2.addWeighted(overlay, 0.4, img_out, 0.6, 0)

    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-Seg from OFFLINE source
model = YOLO('/kaggle/input/yolo8-segmentation/yolov8n-seg.pt')
print(f"✓ YOLO-Seg 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 segmentation masks 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-Seg
        results = model(full_path, verbose=False, conf=conf)
        result = results[0]

        # Get masks and boxes
        masks_data = None
        boxes_data = None
        num_det = 0

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

        if result.boxes is not None and result.boxes.data is not None:
            boxes_data = result.boxes.data.cpu().numpy()

        # Draw segmentation masks
        img_seg = draw_segmentation_masks(img_bgr, masks_data, boxes_data, conf_threshold=0.3)
        img_seg_rgb = cv2.cvtColor(img_seg, cv2.COLOR_BGR2RGB)

        # Display
        axes[i+1].imshow(img_seg_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'seg_viz_{idx+1}.png', dpi=300, bbox_inches='tight')
    plt.show()
    plt.close()

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

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

stats_list = []

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_mask_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()

        if result.masks is not None and result.masks.data is not None:
            masks_data = result.masks.data.cpu().numpy()
            for mask in masks_data:
                total_mask_area += (mask > 0.5).sum()

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

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

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_mask_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}")

# Save files
stats_df.to_csv('seg_detection_stats.csv', index=False)
print("\n✓ Saved: seg_detection_stats.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✓ Exp18 YOLO-Seg EDA completed!")