# Copyright 2020 Google LLC
#
# 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.
# ==============================================================================
"""Tests for lit_nlp.components.gradient_maps."""

from absl.testing import absltest
from lit_nlp.api import dataset as lit_dataset
from lit_nlp.api import model as lit_model
from lit_nlp.api import types as lit_types
from lit_nlp.components import image_gradient_maps
from lit_nlp.lib import image_utils
import numpy as np
from PIL import Image as PILImage

JsonDict = lit_types.JsonDict


class ClassificationTestModel(lit_model.BatchedModel):

  LABELS = ['Dummy', 'Cat', 'Dog']
  GRADIENT_SHAPE = (60, 40, 3)

  def max_minibatch_size(self) -> int:
    return 10

  def predict_minibatch(self, inputs: list[JsonDict]) -> list[JsonDict]:
    result = []
    for i, _ in enumerate(inputs):
      result.append({
          'preds': np.random.rand(len(self.LABELS)),
          'grads': np.random.rand(*self.GRADIENT_SHAPE),
          'grad_target': self.LABELS[i],
      })
    return result

  def input_spec(self):
    return {
        'image':
            lit_types.ImageBytes(),
        'grad_target':
            lit_types.CategoryLabel(vocab=self.LABELS, required=False)
    }

  def output_spec(self):
    return {
        'preds':
            lit_types.MulticlassPreds(vocab=self.LABELS),
        'grads':
            lit_types.ImageGradients(
                align='image', grad_target_field_key='grad_target'),
    }


class RegressionTestModel(lit_model.BatchedModel):
  """A test model for testing the regression case."""

  GRADIENT_SHAPE = (40, 20, 3)

  def max_minibatch_size(self) -> int:
    return 10

  def predict_minibatch(self, inputs: list[JsonDict]) -> list[JsonDict]:
    """Simulates regression of x1 + 2 * x2 using elements of the image array."""
    result = []
    for example in inputs:
      img = example['image']
      if isinstance(img, str):
        img = image_utils.convert_image_str_to_array(
            img, shape=self.GRADIENT_SHAPE)
      pred = img[0, 0, 0] + img[1, 0, 0] * 2
      grad = np.zeros(shape=self.GRADIENT_SHAPE)
      grad[0, 0, 0] = 1
      grad[1, 0, 0] = 2
      result.append({'pred': pred, 'grads': grad})
    return result

  def input_spec(self):
    return {
        'image': lit_types.ImageBytes(),
    }

  def output_spec(self):
    return {
        'pred': lit_types.RegressionScore(),
        'grads': lit_types.ImageGradients(align='image'),
    }


class ImageGradientsMapsTest(absltest.TestCase):

  def test_interpreter(self):
    interpreter = image_gradient_maps.VanillaGradients()
    model = ClassificationTestModel()
    self.assertTrue(interpreter.is_compatible(
        model, lit_dataset.NoneDataset({'test': model})))

    pil_image = PILImage.new(mode='RGB', size=(300, 200))
    inp = {'image': image_utils.convert_pil_to_image_str(pil_image)}
    run_output = interpreter.run(
        inputs=[inp], model=model, dataset=lit_dataset.Dataset())[0]

    self.assertIn('grads', run_output)
    self.assertIsInstance(run_output['grads'], str)

  def test_regression(self):
    interpreter = image_gradient_maps.GuidedIG()
    model = RegressionTestModel()
    self.assertTrue(interpreter.is_compatible(
        model, lit_dataset.NoneDataset({'test': model})))

    input_image_array = np.zeros(shape=[20, 15, 3], dtype=np.uint8)
    input_image_array[0, 0, 0] = 10
    input_image_array[1, 0, 0] = 20

    pil_image = PILImage.fromarray(input_image_array, mode='RGB')

    inp = {'image': image_utils.convert_pil_to_image_str(pil_image)}
    run_output = interpreter.run(
        inputs=[inp], model=model, dataset=lit_dataset.Dataset())[0]
    self.assertIn('grads', run_output)
    overlay_str = run_output['grads']
    overlay_bytes = image_utils.convert_image_str_to_array(
        overlay_str, shape=RegressionTestModel.GRADIENT_SHAPE)
    self.assertIsNotNone(overlay_bytes)
    self.assertSequenceEqual(overlay_bytes.shape,
                             RegressionTestModel.GRADIENT_SHAPE)


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