// Copyright (c) 2019-2022, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

/// Convert an integer into a Q57 fixed-point fraction.
pub const fn q57(v: i32) -> i64 {
  debug_assert!(v >= -64 && v <= 63);
  (v as i64) << 57
}

#[rustfmt::skip]
const ATANH_LOG2: &[i64; 32] = &[
  0x32B803473F7AD0F4, 0x2F2A71BD4E25E916, 0x2E68B244BB93BA06,
  0x2E39FB9198CE62E4, 0x2E2E683F68565C8F, 0x2E2B850BE2077FC1,
  0x2E2ACC58FE7B78DB, 0x2E2A9E2DE52FD5F2, 0x2E2A92A338D53EEC,
  0x2E2A8FC08F5E19B6, 0x2E2A8F07E51A485E, 0x2E2A8ED9BA8AF388,
  0x2E2A8ECE2FE7384A, 0x2E2A8ECB4D3E4B1A, 0x2E2A8ECA94940FE8,
  0x2E2A8ECA6669811D, 0x2E2A8ECA5ADEDD6A, 0x2E2A8ECA57FC347E,
  0x2E2A8ECA57438A43, 0x2E2A8ECA57155FB4, 0x2E2A8ECA5709D510,
  0x2E2A8ECA5706F267, 0x2E2A8ECA570639BD, 0x2E2A8ECA57060B92,
  0x2E2A8ECA57060008, 0x2E2A8ECA5705FD25, 0x2E2A8ECA5705FC6C,
  0x2E2A8ECA5705FC3E, 0x2E2A8ECA5705FC33, 0x2E2A8ECA5705FC30,
  0x2E2A8ECA5705FC2F, 0x2E2A8ECA5705FC2F
];

/// Computes the binary exponential of `logq57`.
/// `logq57`: a log base 2 in Q57 format.
/// Returns a 64 bit integer in Q0 (no fraction).
pub const fn bexp64(logq57: i64) -> i64 {
  let ipart = (logq57 >> 57) as i32;
  if ipart < 0 {
    return 0;
  }
  if ipart >= 63 {
    return 0x7FFFFFFFFFFFFFFF;
  }
  // z is the fractional part of the log in Q62 format.
  // We need 1 bit of headroom since the magnitude can get larger than 1
  //  during the iteration, and a sign bit.
  let mut z = logq57 - q57(ipart);
  let mut w: i64;
  if z != 0 {
    // Rust has 128 bit multiplies, so it should be possible to do this
    //  faster without losing accuracy.
    z <<= 5;
    // w is the exponential in Q61 format (since it also needs headroom and can
    //  get as large as 2.0); we could get another bit if we dropped the sign,
    //  but we'll recover that bit later anyway.
    // Ideally this should start out as
    //   \lim_{n->\infty} 2^{61}/\product_{i=1}^n \sqrt{1-2^{-2i}}
    //  but in order to guarantee convergence we have to repeat iterations 4,
    //  13 (=3*4+1), and 40 (=3*13+1, etc.), so it winds up somewhat larger.
    w = 0x26A3D0E401DD846D;
    let mut i: i64 = 0;
    loop {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z -= (ATANH_LOG2[i as usize] + mask) ^ mask;
      // Repeat iteration 4.
      if i >= 3 {
        break;
      }
      z *= 2;
      i += 1;
    }
    loop {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z -= (ATANH_LOG2[i as usize] + mask) ^ mask;
      // Repeat iteration 13.
      if i >= 12 {
        break;
      }
      z *= 2;
      i += 1;
    }
    while i < 32 {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z = (z - ((ATANH_LOG2[i as usize] + mask) ^ mask)) * 2;
      i += 1;
    }
    // Skip the remaining iterations unless we really require that much
    //  precision.
    // We could have bailed out earlier for smaller iparts, but that would
    //  require initializing w from a table, as the limit doesn't converge to
    //  61-bit precision until n=30.
    let mut wlo: i32 = 0;
    if ipart > 30 {
      // For these iterations, we just update the low bits, as the high bits
      //  can't possibly be affected.
      // OD_ATANH_LOG2 has also converged (it actually did so one iteration
      //  earlier, but that's no reason for an extra special case).
      loop {
        let mask = -((z < 0) as i64);
        wlo += (((w >> i) + mask) ^ mask) as i32;
        z -= (ATANH_LOG2[31] + mask) ^ mask;
        // Repeat iteration 40.
        if i >= 39 {
          break;
        }
        z *= 2;
        i += 1;
      }
      while i < 61 {
        let mask = -((z < 0) as i64);
        wlo += (((w >> i) + mask) ^ mask) as i32;
        z = (z - ((ATANH_LOG2[31] + mask) ^ mask)) * 2;
        i += 1;
      }
    }
    w = (w << 1) + (wlo as i64);
  } else {
    w = 1i64 << 62;
  }
  if ipart < 62 {
    w = ((w >> (61 - ipart)) + 1) >> 1;
  }
  w
}

/// Computes the binary log of `n`.
/// `n`: a 64-bit integer in Q0 (no fraction).
/// Returns a 64-bit log in Q57.
pub const fn blog64(n: i64) -> i64 {
  if n <= 0 {
    return -1;
  }
  let ipart = 63 - n.leading_zeros() as i32;
  let w = if ipart > 61 { n >> (ipart - 61) } else { n << (61 - ipart) };
  if (w & (w - 1)) == 0 {
    return q57(ipart);
  }
  // z is the fractional part of the log in Q61 format.
  let mut z: i64 = 0;
  // Rust has 128 bit multiplies, so it should be possible to do this
  //  faster without losing accuracy.
  // x and y are the cosh() and sinh(), respectively, in Q61 format.
  // We are computing z = 2*atanh(y/x) = 2*atanh((w - 1)/(w + 1)).
  let mut x = w + (1i64 << 61);
  let mut y = w - (1i64 << 61);
  // Repeat iteration 4.
  // Repeat iteration 13.
  // Repeat iteration 40.
  let bounds = [3, 12, 39, 61];
  let mut i = 0;
  let mut j = 0;
  loop {
    let end = bounds[j];
    loop {
      let mask = -((y < 0) as i64);
      // ATANH_LOG2 has converged at iteration 32.
      z += ((ATANH_LOG2[if i < 31 { i } else { 31 }] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
      if i == end {
        break;
      }
      i += 1;
    }
    j += 1;
    if j == bounds.len() {
      break;
    }
  }
  z = (z + 8) >> 4;
  q57(ipart) + z
}

/// Computes the binary log of `n`.
/// `n`: an unsigned 32-bit integer in Q0 (no fraction).
/// Returns a signed 32-bit log in Q24.
#[allow(unused)]
pub const fn blog32(n: u32) -> i32 {
  if n == 0 {
    return -1;
  }
  let ipart = 31 - n.leading_zeros() as i32;
  let n = n as i64;
  let w = if ipart > 61 { n >> (ipart - 61) } else { n << (61 - ipart) };
  if (w & (w - 1)) == 0 {
    return ipart << 24;
  }
  // z is the fractional part of the log in Q61 format.
  let mut z: i64 = 0;
  // Rust has 128 bit multiplies, so it should be possible to do this
  //  faster without losing accuracy.
  // x and y are the cosh() and sinh(), respectively, in Q61 format.
  // We are computing z = 2*atanh(y/x) = 2*atanh((w - 1)/(w + 1)).
  let mut x = w + (1i64 << 61);
  let mut y = w - (1i64 << 61);
  // Repeat iteration 4.
  // Repeat iteration 13.
  let bounds = [3, 12, 29];
  let mut i = 0;
  let mut j = 0;
  loop {
    let end = bounds[j];
    loop {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[i] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
      if i == end {
        break;
      }
      i += 1;
    }
    j += 1;
    if j == bounds.len() {
      break;
    }
  }
  const SHIFT: usize = 61 - 24;
  z = (z + (1 << SHIFT >> 1)) >> SHIFT;
  (ipart << 24) + z as i32
}

/// Converts a Q57 fixed-point fraction to Q24 by rounding.
pub const fn q57_to_q24(v: i64) -> i32 {
  (((v >> 32) + 1) >> 1) as i32
}

/// Converts a Q24 fixed-point fraction to Q57.
pub const fn q24_to_q57(v: i32) -> i64 {
  (v as i64) << 33
}

/// Binary exponentiation of a `log_scale` with 24-bit fractional precision and
///  saturation.
/// `log_scale`: A binary logarithm in Q24 format.
/// Returns the binary exponential in Q24 format, saturated to 2**47 - 1 if
///  `log_scale` was too large.
pub const fn bexp_q24(log_scale: i32) -> i64 {
  if log_scale < 23 << 24 {
    let ret = bexp64(((log_scale as i64) << 33) + q57(24));
    if ret < (1i64 << 47) - 1 {
      return ret;
    }
  }
  (1i64 << 47) - 1
}

/// Polynomial approximation of a binary exponential.
/// Q10 input, Q0 output.
#[allow(unused)]
pub const fn bexp32_q10(z: i32) -> u32 {
  let ipart = z >> 10;
  let mut n = ((z & ((1 << 10) - 1)) << 4) as u32;
  n = ({
    n * (((n * (((n * (((n * 3548) >> 15) + 6817)) >> 15) + 15823)) >> 15)
      + 22708)
  } >> 15)
    + 16384;
  if 14 - ipart > 0 {
    (n + (1 << (13 - ipart))) >> (14 - ipart)
  } else {
    n << (ipart - 14)
  }
}

/// Polynomial approximation of a binary logarithm.
/// Q0 input, Q11 output.
pub const fn blog32_q11(w: u32) -> i32 {
  if w == 0 {
    return -1;
  }
  let ipart = 32 - w.leading_zeros() as i32;
  let n = if ipart - 16 > 0 { w >> (ipart - 16) } else { w << (16 - ipart) }
    as i32
    - 32768
    - 16384;
  let fpart = ({
    n * (((n * (((n * (((n * -1402) >> 15) + 2546)) >> 15) - 5216)) >> 15)
      + 15745)
  } >> 15)
    - 6797;
  (ipart << 11) + (fpart >> 3)
}

#[cfg(test)]
mod test {
  use super::*;

  #[test]
  fn blog64_vectors() {
    assert!(blog64(1793) == 0x159dc71e24d32daf);
    assert!(blog64(0x678dde6e5fd29f05) == 0x7d6373ad151ca685);
  }

  #[test]
  fn bexp64_vectors() {
    assert!(bexp64(0x159dc71e24d32daf) == 1793);
    assert!((bexp64(0x7d6373ad151ca685) - 0x678dde6e5fd29f05).abs() < 29);
  }

  #[test]
  fn blog64_bexp64_round_trip() {
    for a in 1..=std::u16::MAX as i64 {
      let b = std::i64::MAX / a;
      let (log_a, log_b, log_ab) = (blog64(a), blog64(b), blog64(a * b));
      assert!((log_a + log_b - log_ab).abs() < 4);
      assert!(bexp64(log_a) == a);
      assert!((bexp64(log_b) - b).abs() < 128);
      assert!((bexp64(log_ab) - a * b).abs() < 128);
    }
  }

  #[test]
  fn blog32_vectors() {
    assert_eq!(blog32(0), -1);
    assert_eq!(blog32(1793), q57_to_q24(0x159dc71e24d32daf));
  }

  #[test]
  fn bexp_q24_vectors() {
    assert_eq!(bexp_q24(i32::MAX), (1i64 << 47) - 1);
    assert_eq!(
      (bexp_q24(q57_to_q24(0x159dc71e24d32daf)) + (1 << 24 >> 1)) >> 24,
      1793
    );
  }

  #[test]
  fn blog32_bexp_q24_round_trip() {
    for a in 1..=std::u16::MAX as u32 {
      let b = (std::u32::MAX >> 9) / a;
      let (log_a, log_b, log_ab) = (blog32(a), blog32(b), blog32(a * b));
      assert!((log_a + log_b - log_ab).abs() < 4);
      assert!((bexp_q24(log_a) - (i64::from(a) << 24)).abs() < (1 << 24 >> 1));
      assert!(((bexp_q24(log_b) >> 24) - i64::from(b)).abs() < 128);
      assert!(
        ((bexp_q24(log_ab) >> 24) - i64::from(a) * i64::from(b)).abs() < 128
      );
    }
  }

  #[test]
  fn blog32_q11_bexp32_q10_round_trip() {
    for a in 1..=std::i16::MAX as i32 {
      let b = std::i16::MAX as i32 / a;
      let (log_a, log_b, log_ab) = (
        blog32_q11(a as u32),
        blog32_q11(b as u32),
        blog32_q11(a as u32 * b as u32),
      );
      assert!((log_a + log_b - log_ab).abs() < 4);
      assert!((bexp32_q10((log_a + 1) >> 1) as i32 - a).abs() < 18);
      assert!((bexp32_q10((log_b + 1) >> 1) as i32 - b).abs() < 2);
      assert!((bexp32_q10((log_ab + 1) >> 1) as i32 - a * b).abs() < 18);
    }
  }
}