// Copyright 2022 The MediaPipe Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.google.mediapipe.tasks.vision.core;

import android.graphics.RectF;
import com.google.mediapipe.formats.proto.RectProto.NormalizedRect;
import com.google.mediapipe.framework.MediaPipeException;
import com.google.mediapipe.framework.Packet;
import com.google.mediapipe.framework.ProtoUtil;
import com.google.mediapipe.framework.image.MPImage;
import com.google.mediapipe.tasks.core.TaskResult;
import com.google.mediapipe.tasks.core.TaskRunner;
import java.util.HashMap;
import java.util.Map;

/** The base class of MediaPipe vision tasks. */
public class BaseVisionTaskApi implements AutoCloseable {
  private static final long MICROSECONDS_PER_MILLISECOND = 1000;
  private final TaskRunner runner;
  private final RunningMode runningMode;
  private final String imageStreamName;
  private final String normRectStreamName;

  static {
    System.loadLibrary("mediapipe_tasks_vision_jni");
    ProtoUtil.registerTypeName(NormalizedRect.class, "mediapipe.NormalizedRect");
  }

  /**
   * Constructor to initialize a {@link BaseVisionTaskApi}.
   *
   * @param runner a {@link TaskRunner}.
   * @param runningMode a mediapipe vision task {@link RunningMode}.
   * @param imageStreamName the name of the input image stream.
   * @param normRectStreamName the name of the input normalized rect image stream used to provide
   *     (mandatory) rotation and (optional) region-of-interest.
   */
  public BaseVisionTaskApi(
      TaskRunner runner,
      RunningMode runningMode,
      String imageStreamName,
      String normRectStreamName) {
    this.runner = runner;
    this.runningMode = runningMode;
    this.imageStreamName = imageStreamName;
    this.normRectStreamName = normRectStreamName;
  }

  /**
   * A synchronous method to process single image inputs. The call blocks the current thread until a
   * failure status or a successful result is returned.
   *
   * @param image a MediaPipe {@link MPImage} object for processing.
   * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
   *     input image before running inference.
   * @throws MediaPipeException if the task is not in the image mode.
   */
  protected TaskResult processImageData(
      MPImage image, ImageProcessingOptions imageProcessingOptions) {
    if (runningMode != RunningMode.IMAGE) {
      throw new MediaPipeException(
          MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(),
          "Task is not initialized with the image mode. Current running mode:"
              + runningMode.name());
    }
    Map<String, Packet> inputPackets = new HashMap<>();
    inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image));
    inputPackets.put(
        normRectStreamName,
        runner.getPacketCreator().createProto(convertToNormalizedRect(imageProcessingOptions)));
    return runner.process(inputPackets);
  }

  /**
   * A synchronous method to process continuous video frames. The call blocks the current thread
   * until a failure status or a successful result is returned.
   *
   * @param image a MediaPipe {@link MPImage} object for processing.
   * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
   *     input image before running inference.
   * @param timestampMs the corresponding timestamp of the input image in milliseconds.
   * @throws MediaPipeException if the task is not in the video mode.
   */
  protected TaskResult processVideoData(
      MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
    if (runningMode != RunningMode.VIDEO) {
      throw new MediaPipeException(
          MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(),
          "Task is not initialized with the video mode. Current running mode:"
              + runningMode.name());
    }
    Map<String, Packet> inputPackets = new HashMap<>();
    inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image));
    inputPackets.put(
        normRectStreamName,
        runner.getPacketCreator().createProto(convertToNormalizedRect(imageProcessingOptions)));
    return runner.process(inputPackets, timestampMs * MICROSECONDS_PER_MILLISECOND);
  }

  /**
   * An asynchronous method to send live stream data to the {@link TaskRunner}. The results will be
   * available in the user-defined result listener.
   *
   * @param image a MediaPipe {@link MPImage} object for processing.
   * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
   *     input image before running inference.
   * @param timestampMs the corresponding timestamp of the input image in milliseconds.
   * @throws MediaPipeException if the task is not in the stream mode.
   */
  protected void sendLiveStreamData(
      MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
    if (runningMode != RunningMode.LIVE_STREAM) {
      throw new MediaPipeException(
          MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(),
          "Task is not initialized with the live stream mode. Current running mode:"
              + runningMode.name());
    }
    Map<String, Packet> inputPackets = new HashMap<>();
    inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image));
    inputPackets.put(
        normRectStreamName,
        runner.getPacketCreator().createProto(convertToNormalizedRect(imageProcessingOptions)));
    runner.send(inputPackets, timestampMs * MICROSECONDS_PER_MILLISECOND);
  }

  /** Closes and cleans up the MediaPipe vision task. */
  @Override
  public void close() {
    runner.close();
  }

  /**
   * Converts an {@link ImageProcessingOptions} instance into a {@link NormalizedRect} protobuf
   * message.
   */
  private static NormalizedRect convertToNormalizedRect(
      ImageProcessingOptions imageProcessingOptions) {
    RectF regionOfInterest =
        imageProcessingOptions.regionOfInterest().isPresent()
            ? imageProcessingOptions.regionOfInterest().get()
            : new RectF(0, 0, 1, 1);
    return NormalizedRect.newBuilder()
        .setXCenter(regionOfInterest.centerX())
        .setYCenter(regionOfInterest.centerY())
        .setWidth(regionOfInterest.width())
        .setHeight(regionOfInterest.height())
        // Convert to radians anti-clockwise.
        .setRotation(-(float) Math.PI * imageProcessingOptions.rotationDegrees() / 180.0f)
        .build();
  }

  /**
   * Generates the timestamp of a vision task result object from the vision task running mode and
   * the output packet.
   *
   * @param runningMode MediaPipe Vision Tasks {@link RunningMode}.
   * @param packet the output {@link Packet}.
   */
  public static long generateResultTimestampMs(RunningMode runningMode, Packet packet) {
    if (runningMode == RunningMode.IMAGE) {
      return -1;
    }
    return packet.getTimestamp() / MICROSECONDS_PER_MILLISECOND;
  }
}
