import React, { useRef, useState } from 'react';
import Webcam from 'react-webcam';
import * as ort from 'onnxruntime-web';

function ObjectDetection() {
  const [results, setResults] = useState([]);
  const [loading, setLoading] = useState(false);
  const webcamRef = useRef(null);

  const runInference = async () => {
    if (!webcamRef.current) return;
    setLoading(true);

    try {
      // Capture image from webcam
      const imageSrc = webcamRef.current.getScreenshot();

      // Load the ONNX model
      const model = await ort.InferenceSession.create('./model.onnx');

      // Preprocess the image
      const inputTensor = await preprocessImage(imageSrc);

      // Define model input
      const feeds = { input: inputTensor };

      // Run inference
      const output = await model.run(feeds);

      // Postprocess the output
      const detections = postprocessOutput(output);
      setResults(detections);
    } catch (error) {
      console.error('Error running inference:', error);
    }

    setLoading(false);
  };

  const preprocessImage = async (imageSrc) => {
    const img = new Image();
    img.src = imageSrc;
    await new Promise((resolve) => (img.onload = resolve));

    const canvas = document.createElement('canvas');
    const context = canvas.getContext('2d');

    // Resize to model input size
    const modelInputWidth = 300; // Replace with your model's input width
    const modelInputHeight = 300; // Replace with your model's input height
    canvas.width = modelInputWidth;
    canvas.height = modelInputHeight;

    context.drawImage(img, 0, 0, modelInputWidth, modelInputHeight);

    const imageData = context.getImageData(0, 0, modelInputWidth, modelInputHeight);
    const inputTensor = new Float32Array(imageData.data.length / 4 * 3);

    // Normalize pixel values and prepare input tensor
    for (let i = 0, j = 0; i < imageData.data.length; i += 4) {
      inputTensor[j++] = imageData.data[i] / 255;     // R
      inputTensor[j++] = imageData.data[i + 1] / 255; // G
      inputTensor[j++] = imageData.data[i + 2] / 255; // B
    }

    return new ort.Tensor('float32', inputTensor, [1, 3, modelInputHeight, modelInputWidth]);
  };

  const postprocessOutput = (output) => {
    const boxes = output['boxes'].data; // Replace 'boxes' with your model's output name
    const scores = output['scores'].data; // Replace 'scores' with your model's output name
    const classes = output['classes'].data; // Replace 'classes' with your model's output name

    const detections = [];
    for (let i = 0; i < scores.length; i++) {
      if (scores[i] > 0.5) { // Confidence threshold
        detections.push({
          box: boxes.slice(i * 4, i * 4 + 4),
          score: scores[i],
          class: classes[i],
        });
      }
    }

    return detections;
  };

  return React.createElement(
    'div',
    null,
    React.createElement('h1', null, 'Object Detection with Webcam'),
    React.createElement(Webcam, {
      audio: false,
      ref: webcamRef,
      screenshotFormat: 'image/jpeg',
      width: 300,
      height: 300,
    }),
    React.createElement(
      'button',
      { onClick: runInference, disabled: loading },
      loading ? 'Detecting...' : 'Capture & Detect'
    ),
    React.createElement(
      'div',
      null,
      React.createElement('h2', null, 'Results:'),
      React.createElement(
        'ul',
        null,
        results.map((result, index) =>
          React.createElement(
            'li',
            { key: index },
            `Class: ${result.class}, Score: ${result.score.toFixed(2)}, Box: ${result.box.join(', ')}`
          )
        )
      )
    )
  );
}

export default ObjectDetection;