/**
 * YOLO Postprocessing - Process YOLO detection outputs
 */

import { DETECTION_CONFIG } from '../config.js';
import { cropImage } from '../utils/imageUtils.js';

/**
 * Postprocess YOLO detections
 */
export async function postprocessYOLO(results, originalImage) {
    const output = results.output0 || results[Object.keys(results)[0]];
    const data = output.data;
    const dims = output.dims; // [1, 5, 8400]
    
    const detections = [];
    const { confThreshold, iouThreshold } = DETECTION_CONFIG.yolo;
    
    // YOLOv8 output format: [batch, [x, y, w, h, conf], num_anchors]
    // dims = [1, 5, 8400]
    const numAnchors = dims[2]; // 8400
    
    // Extract boxes with confidence above threshold
    const boxes = [];
    for (let i = 0; i < numAnchors; i++) {
        // Data is organized as: [x_center, y_center, width, height, confidence] for each anchor
        const x_center = data[i];
        const y_center = data[numAnchors + i];
        const width = data[2 * numAnchors + i];
        const height = data[3 * numAnchors + i];
        const confidence = data[4 * numAnchors + i];
        
        if (confidence > confThreshold) {
            // Convert from center format to corner format
            const x1 = (x_center - width / 2) / 640 * originalImage.width;
            const y1 = (y_center - height / 2) / 640 * originalImage.height;
            const x2 = (x_center + width / 2) / 640 * originalImage.width;
            const y2 = (y_center + height / 2) / 640 * originalImage.height;
            
            boxes.push({
                x1: Math.max(0, x1),
                y1: Math.max(0, y1),
                x2: Math.min(originalImage.width, x2),
                y2: Math.min(originalImage.height, y2),
                confidence: confidence
            });
        }
    }
    
    // Apply Non-Maximum Suppression
    const selectedBoxes = nonMaximumSuppression(boxes, iouThreshold);
    
    // Crop each detected embryo
    for (const box of selectedBoxes) {
        const croppedData = await cropImage(originalImage, box.x1, box.y1, box.x2, box.y2);
        
        detections.push({
            box: box,
            confidence: box.confidence,
            imageData: croppedData
        });
    }
    
    return detections;
}

/**
 * Non-Maximum Suppression to remove overlapping boxes
 */
function nonMaximumSuppression(boxes, iouThreshold) {
    // Sort boxes by confidence (descending)
    boxes.sort((a, b) => b.confidence - a.confidence);
    
    const selected = [];
    const suppressed = new Set();
    
    for (let i = 0; i < boxes.length; i++) {
        if (suppressed.has(i)) continue;
        
        selected.push(boxes[i]);
        
        // Suppress overlapping boxes
        for (let j = i + 1; j < boxes.length; j++) {
            if (suppressed.has(j)) continue;
            
            const iou = calculateIoU(boxes[i], boxes[j]);
            if (iou > iouThreshold) {
                suppressed.add(j);
            }
        }
    }
    
    return selected;
}

/**
 * Calculate Intersection over Union (IoU) between two boxes
 */
function calculateIoU(box1, box2) {
    const x1 = Math.max(box1.x1, box2.x1);
    const y1 = Math.max(box1.y1, box2.y1);
    const x2 = Math.min(box1.x2, box2.x2);
    const y2 = Math.min(box1.y2, box2.y2);
    
    const intersectionArea = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
    
    const box1Area = (box1.x2 - box1.x1) * (box1.y2 - box1.y1);
    const box2Area = (box2.x2 - box2.x1) * (box2.y2 - box2.y1);
    
    const unionArea = box1Area + box2Area - intersectionArea;
    
    return intersectionArea / unionArea;
}