impact/lib/services/yolo_impact_detection_service.dart

import 'dart:io';
import 'dart:math' as math;
import 'dart:typed_data';
import 'package:tflite_flutter/tflite_flutter.dart';
import 'package:image/image.dart' as img;
import 'target_detection_service.dart';

class YOLOImpactDetectionService {
  Interpreter? _interpreter;
  int _inputSize = 640;
  List<int> _outputShape = [1, 17, 8400];
  int _numClasses = 13;

  static const String modelPath = 'assets/models/yolov8n_32.tflite';
  static const String labelsPath = 'assets/models/labels.txt';

  Future<void> init() async {
    if (_interpreter != null) return;

    try {
      _interpreter = await Interpreter.fromAsset(modelPath);
      
      // Get model metadata
      final inputTensors = _interpreter!.getInputTensors();
      if (inputTensors.isNotEmpty) {
        // [1, 640, 640, 3] or [1, 3, 640, 640]
        final shape = inputTensors[0].shape;
        if (shape.length == 4) {
          _inputSize = shape[1] == 3 ? shape[2] : shape[1];
        }
      }

      final outputTensors = _interpreter!.getOutputTensors();
      if (outputTensors.isNotEmpty) {
        _outputShape = outputTensors[0].shape;
        // Output is usually [1, 4 + num_classes, num_boxes]
        if (_outputShape.length == 3) {
          _numClasses = _outputShape[1] - 4;
        }
      }

      print('YOLO Interpreter loaded successfully');
      print('Model Input Size: $_inputSize');
      print('Model Output Shape: $_outputShape (Classes: $_numClasses)');
    } catch (e) {
      print('Error loading YOLO model: $e');
    }
  }

  Future<List<DetectedImpactResult>> detectImpacts(String imagePath) async {
    if (_interpreter == null) await init();
    if (_interpreter == null) return [];

    try {
      final bytes = File(imagePath).readAsBytesSync();
      final decodedImage = img.decodeImage(bytes);
      if (decodedImage == null) return [];

      final originalImage = img.bakeOrientation(decodedImage);

      final resizedImage = img.copyResize(
        originalImage,
        width: _inputSize,
        height: _inputSize,
      );

      var input = _imageToByteListFloat32(resizedImage, _inputSize);
      
      // Allocate output buffer dynamically
      var output = List<double>.filled(
        _outputShape.fold(1, (a, b) => a * b), 
        0
      ).reshape(_outputShape);

      _interpreter!.run(input, output);

      final results = _processOutput(
        output[0],
        originalImage.width,
        originalImage.height,
      );

      print('YOLO Detection result count: ${results.length}');
      return results;
    } catch (e) {
      print('Error during YOLO inference: $e');
      return [];
    }
  }

  List<DetectedImpactResult> _processOutput(
    List<List<double>> output,
    int imgWidth,
    int imgHeight,
  ) {
    final List<_Detection> candidates = [];
    const double threshold = 0.25;

    // output is typically [4 + numClasses, numBoxes]
    final int rows = output.length;
    final int numBoxes = output[0].length;

    for (int i = 0; i < numBoxes; i++) {
      double maxConfidence = 0;
      int bestClass = 4;
      
      // Find best class among all classes (indices 4 to rows-1)
      for (int c = 4; c < rows; c++) {
        if (output[c][i] > maxConfidence) {
          maxConfidence = output[c][i];
          bestClass = c;
        }
      }

      if (maxConfidence > threshold) {
        final int classIndex = bestClass - 4;
        
        candidates.add(
          _Detection(
            // Use dynamic mapping: 0,1,2,3 are typically x,y,w,h
            // We'll keep the current mapping for now as it matches user's previous model
            y: output[0][i],
            x: output[1][i],
            w: output[2][i],
            h: output[3][i],
            confidence: maxConfidence,
            classIndex: classIndex,
          ),
        );
      }
    }

    final List<_Detection> suppressed = _nms(candidates);

    return suppressed
        .map(
          (det) {
            // Score = ClassIndex (0 to 10) for impact models
            int score = det.classIndex;
            
            return DetectedImpactResult(
              x: det.x, 
              y: det.y,
              radius: 5.0,
              suggestedScore: score,
            );
          },
        )
        .toList();
  }

  List<_Detection> _nms(List<_Detection> detections) {
    if (detections.isEmpty) return [];

    detections.sort((a, b) => b.confidence.compareTo(a.confidence));

    final List<_Detection> selected = [];
    final List<bool> active = List.filled(detections.length, true);

    for (int i = 0; i < detections.length; i++) {
      if (!active[i]) continue;

      selected.add(detections[i]);

      for (int j = i + 1; j < detections.length; j++) {
        if (!active[j]) continue;

        if (_iou(detections[i], detections[j]) > 0.45) {
          active[j] = false;
        }
      }
    }

    return selected;
  }

  double _iou(_Detection a, _Detection b) {
    final double areaA = a.w * a.h;
    final double areaB = b.w * b.h;

    final double x1 = math.max(a.x - a.w / 2, b.x - b.w / 2);
    final double y1 = math.max(a.y - a.h / 2, b.y - b.h / 2);
    final double x2 = math.min(a.x + a.w / 2, b.x + b.w / 2);
    final double y2 = math.min(a.y + a.h / 2, b.y + b.h / 2);

    final double intersection = math.max(0.0, x2 - x1) * math.max(0.0, y2 - y1);
    return intersection / (areaA + areaB - intersection);
  }

  Uint8List _imageToByteListFloat32(img.Image image, int inputSize) {
    var convertedBytes = Float32List(1 * inputSize * inputSize * 3);
    var buffer = Float32List.view(convertedBytes.buffer);
    int pixelIndex = 0;
    for (int i = 0; i < inputSize; i++) {
      for (int j = 0; j < inputSize; j++) {
        var pixel = image.getPixel(j, i);
        buffer[pixelIndex++] = (pixel.r / 255.0);
        buffer[pixelIndex++] = (pixel.g / 255.0);
        buffer[pixelIndex++] = (pixel.b / 255.0);
      }
    }
    return convertedBytes.buffer.asUint8List();
  }
}

class _Detection {
  final double x, y, w, h, confidence;
  final int classIndex;
  _Detection({
    required this.x,
    required this.y,
    required this.w,
    required this.h,
    required this.confidence,
    required this.classIndex,
  });
}