// Copyright 2016 The Draco Authors.
//
// 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.
//
// Utilities for converting unit vectors to octahedral coordinates and back.
// For more details about octahedral coordinates, see for example Cigolle
// et al.'14 “A Survey of Efficient Representations for Independent Unit
// Vectors”.
//
// In short this is motivated by an octahedron inscribed into a sphere. The
// direction of the normal vector can be defined by a point on the octahedron.
// On the right hemisphere (x > 0) this point is projected onto the x = 0 plane,
// that is, the right side of the octahedron forms a diamond like shape. The
// left side of the octahedron is also projected onto the x = 0 plane, however,
// in this case we flap the triangles of the diamond outward. Afterwards we
// shift the resulting square such that all values are positive.
//
// Important values in this file:
// * q: number of quantization bits
// * max_quantized_value: the max value representable with q bits (odd)
// * max_value: max value of the diamond = max_quantized_value - 1  (even)
// * center_value: center of the diamond after shift
//
// Note that the parameter space is somewhat periodic, e.g. (0, 0) ==
// (max_value, max_value), which is also why the diamond is one smaller than the
// maximal representable value in order to have an odd range of values.

#ifndef DRACO_COMPRESSION_ATTRIBUTES_NORMAL_COMPRESSION_UTILS_H_
#define DRACO_COMPRESSION_ATTRIBUTES_NORMAL_COMPRESSION_UTILS_H_

#include <inttypes.h>
#include <algorithm>
#include <cmath>

#include "draco/core/macros.h"

namespace draco {

class OctahedronToolBox {
 public:
  OctahedronToolBox()
      : quantization_bits_(-1),
        max_quantized_value_(-1),
        max_value_(-1),
        center_value_(-1) {}

  bool SetQuantizationBits(int32_t q) {
    if (q < 2 || q > 30)
      return false;
    quantization_bits_ = q;
    max_quantized_value_ = (1 << quantization_bits_) - 1;
    max_value_ = max_quantized_value_ - 1;
    center_value_ = max_value_ / 2;
    return true;
  }
  bool IsInitialized() const { return quantization_bits_ != -1; }

  // Convert all edge points in the top left and bottom right quadrants to
  // their corresponding position in the bottom left and top right quadrants.
  // Convert all corner edge points to the top right corner.
  inline void CanonicalizeOctahedralCoords(int32_t s, int32_t t, int32_t *out_s,
                                           int32_t *out_t) const {
    if ((s == 0 && t == 0) || (s == 0 && t == max_value_) ||
        (s == max_value_ && t == 0)) {
      s = max_value_;
      t = max_value_;
    } else if (s == 0 && t > center_value_) {
      t = center_value_ - (t - center_value_);
    } else if (s == max_value_ && t < center_value_) {
      t = center_value_ + (center_value_ - t);
    } else if (t == max_value_ && s < center_value_) {
      s = center_value_ + (center_value_ - s);
    } else if (t == 0 && s > center_value_) {
      s = center_value_ - (s - center_value_);
    }

    *out_s = s;
    *out_t = t;
  }

  // Converts an integer vector to octahedral coordinates.
  // Precondition: |int_vec| abs sum must equal center value.
  inline void IntegerVectorToQuantizedOctahedralCoords(const int32_t *int_vec,
                                                       int32_t *out_s,
                                                       int32_t *out_t) const {
    DRACO_DCHECK_EQ(
        std::abs(int_vec[0]) + std::abs(int_vec[1]) + std::abs(int_vec[2]),
        center_value_);
    int32_t s, t;
    if (int_vec[0] >= 0) {
      // Right hemisphere.
      s = (int_vec[1] + center_value_);
      t = (int_vec[2] + center_value_);
    } else {
      // Left hemisphere.
      if (int_vec[1] < 0) {
        s = std::abs(int_vec[2]);
      } else {
        s = (max_value_ - std::abs(int_vec[2]));
      }
      if (int_vec[2] < 0) {
        t = std::abs(int_vec[1]);
      } else {
        t = (max_value_ - std::abs(int_vec[1]));
      }
    }
    CanonicalizeOctahedralCoords(s, t, out_s, out_t);
  }

  template <class T>
  void FloatVectorToQuantizedOctahedralCoords(const T *vector, int32_t *out_s,
                                              int32_t *out_t) const {
    const double abs_sum = std::abs(static_cast<double>(vector[0])) +
                           std::abs(static_cast<double>(vector[1])) +
                           std::abs(static_cast<double>(vector[2]));

    // Adjust values such that abs sum equals 1.
    double scaled_vector[3];
    if (abs_sum > 1e-6) {
      // Scale needed to project the vector to the surface of an octahedron.
      const double scale = 1.0 / abs_sum;
      scaled_vector[0] = vector[0] * scale;
      scaled_vector[1] = vector[1] * scale;
      scaled_vector[2] = vector[2] * scale;
    } else {
      scaled_vector[0] = 1.0;
      scaled_vector[1] = 0;
      scaled_vector[2] = 0;
    }

    // Scale vector such that the sum equals the center value.
    int32_t int_vec[3];
    int_vec[0] =
        static_cast<int32_t>(floor(scaled_vector[0] * center_value_ + 0.5));
    int_vec[1] =
        static_cast<int32_t>(floor(scaled_vector[1] * center_value_ + 0.5));
    // Make sure the sum is exactly the center value.
    int_vec[2] = center_value_ - std::abs(int_vec[0]) - std::abs(int_vec[1]);
    if (int_vec[2] < 0) {
      // If the sum of first two coordinates is too large, we need to decrease
      // the length of one of the coordinates.
      if (int_vec[1] > 0) {
        int_vec[1] += int_vec[2];
      } else {
        int_vec[1] -= int_vec[2];
      }
      int_vec[2] = 0;
    }
    // Take care of the sign.
    if (scaled_vector[2] < 0)
      int_vec[2] *= -1;

    IntegerVectorToQuantizedOctahedralCoords(int_vec, out_s, out_t);
  }

  // Normalize |vec| such that its abs sum is equal to the center value;
  template <class T>
  void CanonicalizeIntegerVector(T *vec) const {
    static_assert(std::is_integral<T>::value, "T must be an integral type.");
    static_assert(std::is_signed<T>::value, "T must be a signed type.");
    const int64_t abs_sum = static_cast<int64_t>(std::abs(vec[0])) +
                            static_cast<int64_t>(std::abs(vec[1])) +
                            static_cast<int64_t>(std::abs(vec[2]));

    if (abs_sum == 0) {
      vec[0] = center_value_;  // vec[1] == v[2] == 0
    } else {
      vec[0] =
          (static_cast<int64_t>(vec[0]) * static_cast<int64_t>(center_value_)) /
          abs_sum;
      vec[1] =
          (static_cast<int64_t>(vec[1]) * static_cast<int64_t>(center_value_)) /
          abs_sum;
      if (vec[2] >= 0) {
        vec[2] = center_value_ - std::abs(vec[0]) - std::abs(vec[1]);
      } else {
        vec[2] = -(center_value_ - std::abs(vec[0]) - std::abs(vec[1]));
      }
    }
  }

  template <typename T>
  void OctaherdalCoordsToUnitVector(T in_s, T in_t, T *out_vector) const {
    DRACO_DCHECK_GE(in_s, 0);
    DRACO_DCHECK_GE(in_t, 0);
    DRACO_DCHECK_LE(in_s, 1);
    DRACO_DCHECK_LE(in_t, 1);
    T s = in_s;
    T t = in_t;
    T spt = s + t;
    T smt = s - t;
    T x_sign = 1.0;
    if (spt >= 0.5 && spt <= 1.5 && smt >= -0.5 && smt <= 0.5) {
      // Right hemisphere. Don't do anything.
    } else {
      // Left hemisphere.
      x_sign = -1.0;
      if (spt <= 0.5) {
        s = 0.5 - in_t;
        t = 0.5 - in_s;
      } else if (spt >= 1.5) {
        s = 1.5 - in_t;
        t = 1.5 - in_s;
      } else if (smt <= -0.5) {
        s = in_t - 0.5;
        t = in_s + 0.5;
      } else {
        s = in_t + 0.5;
        t = in_s - 0.5;
      }
      spt = s + t;
      smt = s - t;
    }
    const T y = 2.0 * s - 1.0;
    const T z = 2.0 * t - 1.0;
    const T x = std::min(std::min(2.0 * spt - 1.0, 3.0 - 2.0 * spt),
                         std::min(2.0 * smt + 1.0, 1.0 - 2.0 * smt)) *
                x_sign;
    // Normalize the computed vector.
    const T normSquared = x * x + y * y + z * z;
    if (normSquared < 1e-6) {
      out_vector[0] = 0;
      out_vector[1] = 0;
      out_vector[2] = 0;
    } else {
      const T d = 1.0 / std::sqrt(normSquared);
      out_vector[0] = x * d;
      out_vector[1] = y * d;
      out_vector[2] = z * d;
    }
  }

  template <typename T>
  void QuantizedOctaherdalCoordsToUnitVector(int32_t in_s, int32_t in_t,
                                             T *out_vector) const {
    T scale = 1.0 / static_cast<T>(max_value_);
    OctaherdalCoordsToUnitVector(in_s * scale, in_t * scale, out_vector);
  }

  // |s| and |t| are expected to be signed values.
  inline bool IsInDiamond(const int32_t &s, const int32_t &t) const {
    // Expect center already at origin.
    DRACO_DCHECK_LE(s, center_value_);
    DRACO_DCHECK_LE(t, center_value_);
    DRACO_DCHECK_GE(s, -center_value_);
    DRACO_DCHECK_GE(t, -center_value_);
    return std::abs(s) + std::abs(t) <= center_value_;
  }

  void InvertDiamond(int32_t *s, int32_t *t) const {
    // Expect center already at origin.
    DRACO_DCHECK_LE(*s, center_value_);
    DRACO_DCHECK_LE(*t, center_value_);
    DRACO_DCHECK_GE(*s, -center_value_);
    DRACO_DCHECK_GE(*t, -center_value_);
    int32_t sign_s = 0;
    int32_t sign_t = 0;
    if (*s >= 0 && *t >= 0) {
      sign_s = 1;
      sign_t = 1;
    } else if (*s <= 0 && *t <= 0) {
      sign_s = -1;
      sign_t = -1;
    } else {
      sign_s = (*s > 0) ? 1 : -1;
      sign_t = (*t > 0) ? 1 : -1;
    }

    const int32_t corner_point_s = sign_s * center_value_;
    const int32_t corner_point_t = sign_t * center_value_;
    *s = 2 * *s - corner_point_s;
    *t = 2 * *t - corner_point_t;
    if (sign_s * sign_t >= 0) {
      int32_t temp = *s;
      *s = -*t;
      *t = -temp;
    } else {
      std::swap(*s, *t);
    }
    *s = (*s + corner_point_s) / 2;
    *t = (*t + corner_point_t) / 2;
  }

  void InvertDirection(int32_t *s, int32_t *t) const {
    // Expect center already at origin.
    DRACO_DCHECK_LE(*s, center_value_);
    DRACO_DCHECK_LE(*t, center_value_);
    DRACO_DCHECK_GE(*s, -center_value_);
    DRACO_DCHECK_GE(*t, -center_value_);
    *s *= -1;
    *t *= -1;
    this->InvertDiamond(s, t);
  }

  // For correction values.
  int32_t ModMax(int32_t x) const {
    if (x > this->center_value())
      return x - this->max_quantized_value();
    if (x < -this->center_value())
      return x + this->max_quantized_value();
    return x;
  }

  // For correction values.
  int32_t MakePositive(int32_t x) const {
    DRACO_DCHECK_LE(x, this->center_value() * 2);
    if (x < 0)
      return x + this->max_quantized_value();
    return x;
  }

  int32_t quantization_bits() const { return quantization_bits_; }
  int32_t max_quantized_value() const { return max_quantized_value_; }
  int32_t max_value() const { return max_value_; }
  int32_t center_value() const { return center_value_; }

 private:
  int32_t quantization_bits_;
  int32_t max_quantized_value_;
  int32_t max_value_;
  int32_t center_value_;
};
}  // namespace draco

#endif  // DRACO_COMPRESSION_ATTRIBUTES_NORMAL_COMPRESSION_UTILS_H_
