# Copyright 2025 The Android Open Source Project
#
# 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.
"""Verifies sensitivities on RAW images for AE ISO priority mode."""

import logging
import math
import os.path

import numpy as np
from matplotlib import pyplot as plt
from mobly import test_runner

import camera_properties_utils
import image_processing_utils
import its_base_test
import its_session_utils

_BLACK_LEVEL_RTOL = 0.005  # 0.5%
_GR_PLANE_IDX = 1  # GR plane index in RGGB data
_IMG_STATS_GRID = 9  # Center 11.11%
_NAME = os.path.splitext(os.path.basename(__file__))[0]
_NUM_FRAMES = 4
_NUM_SENS_STEPS = 5
_VAR_THRESH = 1.01  # Each shot must be 1% noisier than previous
# CameraMetadata enum value
_CONTROL_AE_PRIORITY_MODE_SENSOR_SENSITIVITY_PRIORITY = 1


def ae_iso_priority_capture_request(sensitivity):
  """Returns a capture request enabling ISO priority mode.

  Args:
   sensitivity: The sensitivity value to populate the request with.

  Returns:
    The capture request, ready to be passed to the
    its_session_utils.device.do_capture function.
  """
  req = {
      'android.control.mode': 1,  # CONTROL_MODE_AUTO
      'android.control.aeMode': 1,  # CONTROL_AE_MODE_ON
      'android.control.afMode': 0,  # CONTROL_AF_MODE_OFF
      'android.control.aePriorityMode': 1,
      # CONTROL_AE_PRIORITY_MODE_SENSOR_SENSITIVITY_PRIORITY
      'android.sensor.sensitivity': sensitivity,
  }

  return req


class SensorSensitivityPriorityTest(its_base_test.ItsBaseTest):
  """Capture a set of raw images with increasing gains & measure noise.
  """

  def test_sensor_sensitivity_priority(self):
    logging.debug('Starting %s', _NAME)
    with its_session_utils.ItsSession(
        device_id=self.dut.serial,
        camera_id=self.camera_id,
        hidden_physical_id=self.hidden_physical_id) as cam:
      props = cam.get_camera_properties()
      props = cam.override_with_hidden_physical_camera_props(props)
      camera_properties_utils.skip_unless(
          camera_properties_utils.raw16(props) and
          camera_properties_utils.read_3a(props) and
          camera_properties_utils.per_frame_control(props) and
          (_CONTROL_AE_PRIORITY_MODE_SENSOR_SENSITIVITY_PRIORITY in
           camera_properties_utils.ae_priority_mode(props)) and
          not camera_properties_utils.mono_camera(props))
      name_with_log_path = os.path.join(self.log_path, _NAME)

      # Load chart for scene
      its_session_utils.load_scene(
          cam, props, self.scene, self.tablet,
          its_session_utils.CHART_DISTANCE_NO_SCALING)

      # Expose for the scene with min sensitivity
      sens_min, _ = props['android.sensor.info.sensitivityRange']
      # Digital gains might not be visible on RAW data
      sens_max = props['android.sensor.maxAnalogSensitivity']
      sens_step = (sens_max - sens_min) // _NUM_SENS_STEPS

      # Intentionally blur images for noise measurements
      cam.do_3a(do_af=False, get_results=False)

      # Get property constants for captures
      cfa_idxs = image_processing_utils.get_canonical_cfa_order(props)
      black_levels = image_processing_utils.get_black_levels(props)
      white_level = props['android.sensor.info.whiteLevel']

      sensitivities = list(range(sens_min, sens_max, sens_step))
      variances = []
      capture_requests = []

      for s in sensitivities:
        req = ae_iso_priority_capture_request(s)
        capture_requests.extend([req] * _NUM_FRAMES)

      # Capture in rawStats to reduce test run time
      fmt = its_session_utils.define_raw_stats_fmt_sensor_sensitivity(
          props, _IMG_STATS_GRID)
      caps = cam.do_capture(capture_requests, fmt)

      image_processing_utils.assert_capture_width_and_height(
          caps[0], _IMG_STATS_GRID, _IMG_STATS_GRID
      )
      # Measure mean & variance
      for i, cap in enumerate(caps):
        mean_img, var_img = image_processing_utils.unpack_rawstats_capture(
            cap
        )
        mean = mean_img[_IMG_STATS_GRID//2, _IMG_STATS_GRID//2,
                        cfa_idxs[_GR_PLANE_IDX]]
        var = var_img[_IMG_STATS_GRID//2, _IMG_STATS_GRID//2,
                      cfa_idxs[_GR_PLANE_IDX]]/white_level**2
        logging.debug('cap: %d, mean: %.2f, var: %e', i, mean, var)
        variances.append(var)

      # Flag dark images
      if math.isclose(mean, max(black_levels), rel_tol=_BLACK_LEVEL_RTOL):
        raise AssertionError(f'Images are too dark! Center mean: {mean:.2f}')

      # Create plot
      sensitivities = np.repeat(sensitivities, _NUM_FRAMES)
      plt.figure(_NAME)
      plt.plot(sensitivities, variances, '-ro')
      plt.xticks(sensitivities)
      plt.xlabel('Sensitivities')
      plt.ylabel('Image Center Patch Variance')
      plt.ticklabel_format(axis='y', style='sci', scilimits=(-6, -6))
      plt.title(_NAME)
      plt.savefig(f'{name_with_log_path}_variances.png')

      # Find average variance at each step
      vars_step_means = []
      for i in range(_NUM_SENS_STEPS):
        vars_step = []
        for j in range(_NUM_FRAMES):
          vars_step.append(variances[_NUM_FRAMES * i + j])
        vars_step_means.append(np.mean(vars_step))
      logging.debug('averaged variances: %s', vars_step_means)

      # Assert each set of shots is noisier than previous
      for variance_idx, variance in enumerate(vars_step_means[:-1]):
        if variance >= vars_step_means[variance_idx+1] / _VAR_THRESH:
          raise AssertionError(
              f'variances [i]: {variances[variance_idx]:.5f}, '
              f'[i+1]: {variances[variance_idx+1]:.5f}, THRESH: {_VAR_THRESH}'
          )

if __name__ == '__main__':
  test_runner.main()