// Copyright (c) the JPEG XL Project Authors. All rights reserved.

//

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file.

//

// Author: Jyrki Alakuijala (jyrki.alakuijala@gmail.com)

#ifndef LIB_JXL_BUTTERAUGLI_BUTTERAUGLI_H_

#define LIB_JXL_BUTTERAUGLI_BUTTERAUGLI_H_

#include <jxl/memory_manager.h>

#include <atomic>

#include <cstddef>

#include <cstdlib>

#include <cstring>

#include <memory>

#include "lib/jxl/base/compiler_specific.h"

#include "lib/jxl/base/status.h"

#include "lib/jxl/image.h"

#if !defined(BUTTERAUGLI_ENABLE_CHECKS)

#define BUTTERAUGLI_ENABLE_CHECKS 0

#endif

#define BUTTERAUGLI_RESTRICT JXL_RESTRICT

// This is the main interface to butteraugli image similarity

// analysis function.

namespace jxl {

struct ButteraugliParams {

  // Multiplier for penalizing new HF artifacts more than blurring away

  // features. 1.0=neutral.

  float hf_asymmetry = 1.0f;

  // Multiplier for the psychovisual difference in the X channel.

  float xmul = 1.0f;

  // Number of nits that correspond to 1.0f input values.

  float intensity_target = 80.0f;

};

// ButteraugliInterface defines the public interface for butteraugli.

//

// It calculates the difference between rgb0 and rgb1.

//

// rgb0 and rgb1 contain the images. rgb0[c][px] and rgb1[c][px] contains

// the red image for c == 0, green for c == 1, blue for c == 2. Location index

// px is calculated as y * xsize + x.

//

// Value of pixels of images rgb0 and rgb1 need to be represented as raw

// intensity. Most image formats store gamma corrected intensity in pixel

// values. This gamma correction has to be removed, by applying the following

// function to values in the 0-1 range:

// butteraugli_val = pow(input_val, gamma);

// A typical value of gamma is 2.2. It is usually stored in the image header.

// Take care not to confuse that value with its inverse. The gamma value should

// be always greater than one.

// Butteraugli does not work as intended if the caller does not perform

// gamma correction.

//

// hf_asymmetry is a multiplier for penalizing new HF artifacts more than

// blurring away features (1.0 -> neutral).

//

// diffmap will contain an image of the size xsize * ysize, containing

// localized differences for values px (indexed with the px the same as rgb0

// and rgb1). diffvalue will give a global score of similarity.

//

// A diffvalue smaller than kButteraugliGood indicates that images can be

// observed as the same image.

// diffvalue larger than kButteraugliBad indicates that a difference between

// the images can be observed.

// A diffvalue between kButteraugliGood and kButteraugliBad indicates that

// a subtle difference can be observed between the images.

//

// Returns true on success.

bool ButteraugliInterface(const Image3F &rgb0, const Image3F &rgb1,

                          const ButteraugliParams &params, ImageF &diffmap,

                          double &diffvalue);

// Deprecated (calls the previous function)

bool ButteraugliInterface(const Image3F &rgb0, const Image3F &rgb1,

                          float hf_asymmetry, float xmul, ImageF &diffmap,

                          double &diffvalue);

// Same as ButteraugliInterface, but reuses rgb0 and rgb1 for other purposes

// inside the function after they are not needed any more, and it ignores

// params.xmul.

Status ButteraugliInterfaceInPlace(Image3F &&rgb0, Image3F &&rgb1,

                                   const ButteraugliParams &params,

                                   ImageF &diffmap, double &diffvalue);

// Converts the butteraugli score into fuzzy class values that are continuous

// at the class boundary. The class boundary location is based on human

// raters, but the slope is arbitrary. Particularly, it does not reflect

// the expectation value of probabilities of the human raters. It is just

// expected that a smoother class boundary will allow for higher-level

// optimization algorithms to work faster.

//

// Returns 2.0 for a perfect match, and 1.0 for 'ok', 0.0 for bad. Because the

// scoring is fuzzy, a butteraugli score of 0.96 would return a class of

// around 1.9.

double ButteraugliFuzzyClass(double score);

// Input values should be in range 0 (bad) to 2 (good). Use

// kButteraugliNormalization as normalization.

double ButteraugliFuzzyInverse(double seek);

// Implementation details, don't use anything below or your code will

// break in the future.

#ifdef _MSC_VER

#define BUTTERAUGLI_INLINE __forceinline

#else

#define BUTTERAUGLI_INLINE inline

#endif

#ifdef __clang__

// Early versions of Clang did not support __builtin_assume_aligned.

#define BUTTERAUGLI_HAS_ASSUME_ALIGNED __has_builtin(__builtin_assume_aligned)

#elif defined(__GNUC__)

#define BUTTERAUGLI_HAS_ASSUME_ALIGNED 1

#else

#define BUTTERAUGLI_HAS_ASSUME_ALIGNED 0

#endif

// Returns a void* pointer which the compiler then assumes is N-byte aligned.

// Example: float* JXL_RESTRICT aligned = (float*)JXL_ASSUME_ALIGNED(in, 32);

//

// The assignment semantics are required by GCC/Clang. ICC provides an in-place

// __assume_aligned, whereas MSVC's __assume appears unsuitable.

#if BUTTERAUGLI_HAS_ASSUME_ALIGNED

#define BUTTERAUGLI_ASSUME_ALIGNED(ptr, align) \

  __builtin_assume_aligned((ptr), (align))

#else

#define BUTTERAUGLI_ASSUME_ALIGNED(ptr, align) (ptr)

#endif  // BUTTERAUGLI_HAS_ASSUME_ALIGNED

struct PsychoImage {

  ImageF uhf[2];  // XY

  ImageF hf[2];   // XY

  Image3F mf;     // XYB

  Image3F lf;     // XYB

};

// Blur needs a transposed image.

// Hold it here and only allocate on demand to reduce memory usage.

struct BlurTemp {

  Status GetTransposed(const ImageF &in, ImageF **out) {

    JxlMemoryManager *memory_manager = in.memory_manager();

    if (transposed_temp.xsize() == 0) {

      JXL_ASSIGN_OR_RETURN(

          transposed_temp,

          ImageF::Create(memory_manager, in.ysize(), in.xsize()));

    }

    *out = &transposed_temp;

    return true;

  }

  ImageF transposed_temp;

};

class ButteraugliComparator {

 public:

  // Butteraugli is calibrated at xmul = 1.0. We add a multiplier here so that

  // we can test the hypothesis that a higher weighing of the X channel would

  // improve results at higher Butteraugli values.

  virtual ~ButteraugliComparator() = default;

  static StatusOr<std::unique_ptr<ButteraugliComparator>> Make(

      const Image3F &rgb0, const ButteraugliParams &params);

  // Computes the butteraugli map between the original image given in the

  // constructor and the distorted image give here.

  Status Diffmap(const Image3F &rgb1, ImageF &result) const;

  // Same as above, but OpsinDynamicsImage() was already applied.

  Status DiffmapOpsinDynamicsImage(const Image3F &xyb1, ImageF &result) const;

  // Same as above, but the frequency decomposition was already applied.

  Status DiffmapPsychoImage(const PsychoImage &pi1, ImageF &diffmap) const;

  Status Mask(ImageF *BUTTERAUGLI_RESTRICT mask) const;

 private:

  ButteraugliComparator(size_t xsize, size_t ysize,

                        const ButteraugliParams &params);

  Image3F *Temp() const;

  void ReleaseTemp() const;

  const size_t xsize_;

  const size_t ysize_;

  ButteraugliParams params_;

  PsychoImage pi0_;

  // Shared temporary image storage to reduce the number of allocations;

  // obtained via Temp(), must call ReleaseTemp when no longer needed.

  mutable Image3F temp_;

  mutable std::atomic_flag temp_in_use_ = ATOMIC_FLAG_INIT;

  mutable BlurTemp blur_temp_;

  std::unique_ptr<ButteraugliComparator> sub_;

};

// Deprecated.

Status ButteraugliDiffmap(const Image3F &rgb0, const Image3F &rgb1,

                          double hf_asymmetry, double xmul, ImageF &diffmap);

Status ButteraugliDiffmap(const Image3F &rgb0, const Image3F &rgb1,

                          const ButteraugliParams &params, ImageF &diffmap);

double ButteraugliScoreFromDiffmap(const ImageF &diffmap,

                                   const ButteraugliParams *params = nullptr);

// Generate rgb-representation of the distance between two images.

StatusOr<Image3F> CreateHeatMapImage(const ImageF &distmap,

                                     double good_threshold,

                                     double bad_threshold);

}  // namespace jxl

#endif  // LIB_JXL_BUTTERAUGLI_BUTTERAUGLI_H_