# Copyright 2022 Google.
#
# 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.

"""Helper routines for recording various training metrics."""

from typing import Any
import jax.numpy as jnp


Array = Any


def compute_accuracy_sum(logits, targets, valid_loss_mask=None):
  """Compute accuracy for logits and targets.

  Args:
   logits: [batch, length, num_classes] float array.
   targets: categorical targets [batch, length] int array.
   valid_loss_mask: None or array of shape bool[batch, length]

  Returns:
    The number of correct tokens in the output.
  """
  if logits.shape[:-1] != targets.shape:
    raise ValueError("Incorrect shapes. Got shape %s logits and %s targets" %
                     logits.shape, targets.shape)
  if valid_loss_mask is not None and valid_loss_mask.shape != targets.shape:
    raise ValueError("Incorrect shapes. Got shape %s targets and %s mask" %
                     targets.shape, valid_loss_mask.shape)

  accuracy = jnp.equal(jnp.argmax(logits, axis=-1), targets)
  if valid_loss_mask is not None:
    accuracy = jnp.logical_and(accuracy, valid_loss_mask)
  return jnp.sum(accuracy)  # Sum of the number of True values.


def reshape_image(image):
  """Reshape image to something that tensorboard recognizes.

  Args:
    image: Array of shape [xsize, size] or [num_images, xsize, ysize]

  Returns:
    Array of shape [num_images, xsize, ysize, 1]
  """

  # Reshape to [num_images, xdim, ydim, rgb] for tensorboard.
  sh = image.shape
  if image.ndim == 2:
    return jnp.reshape(image, [1, sh[0], sh[1], 1]).astype(jnp.float32)
  elif image.ndim == 3:
    return jnp.reshape(image, [sh[0], sh[1], sh[2], 1]).astype(jnp.float32)
  else:
    return None  # Not an image.


def normalize_image(images: Array, as_group: bool = False) -> Array:
  """Rescale the values in images to between 0.0 and 1.0.

  Args:
    images:   Array of size [batch_size, xsize, ysize]
    as_group: Scale all images in the batch by the same amount if True.

  Returns:
    A rescaled image of the same shape.
  """

  images = images.astype(jnp.float32)  # Return images as float32.
  if as_group:
    # Normalize the batch of images as a group.
    min_img = jnp.min(images)
    max_img = jnp.max(images)
  else:
    # Normalize each image in the batch individually.
    min_img = jnp.min(images, axis=(-2, -1), keepdims=True)
    max_img = jnp.max(images, axis=(-2, -1), keepdims=True)
  norm_image = (images - min_img) / (max_img - min_img + 1e-6)
  return jnp.where(jnp.isfinite(norm_image), norm_image, 0.0)


def overlay_images(image1: Array, image2: Array) -> Array:
  """Place image1 on top of image2, broadcasting image2 if necessary.

  Args:
    image1: array of shape [num_images, xsize, ysize]
    image2: array of shape [num_images, xsize, ysize]

  Returns:
    A combined image.
  """

  assert image1.ndim == 3  # (num_images, xsize, ysize)
  assert image2.ndim == 3
  image2 = jnp.broadcast_to(image2, image1.shape)
  return jnp.concatenate([image1, image2], axis=1)


def make_histograms(viz_dicts):
  """Generate image histograms."""
  hist_dict = {}
  for (i, viz_dict) in enumerate(viz_dicts):
    for (k, images) in viz_dict.items():
      hist_dict["h_" + k + "_" + str(i)] = images
  return hist_dict