// RUN: %target-run-simple-swiftgyb(-Xfrontend -enable-experimental-forward-mode-differentiation)
// REQUIRES: executable_test

#if canImport(Darwin)
  import Darwin.C.tgmath
#elseif canImport(Glibc)
  import Glibc
#elseif os(Windows)
  import CRT
#else
#error("Unsupported platform")
#endif

#if !(os(Windows) || os(Android)) && (arch(i386) || arch(x86_64))
  typealias TestLiteralType = Float80
#else
  typealias TestLiteralType = Double
#endif

import StdlibUnittest
import _Differentiation

let DerivativeTests = TestSuite("TGMath")

func expectEqualWithTolerance<T>(_ expected: TestLiteralType, _ actual: T,
                                 ulps allowed: T = 3,
                                 file: String = #file, line: UInt = #line)
                                 where T: BinaryFloatingPoint {
  if actual == T(expected) || actual.isNaN && expected.isNaN || actual.isInfinite && expected.isInfinite {
    return
  }
  //  Compute error in ulp, compare to tolerance.
  let absoluteError = T(abs(TestLiteralType(actual) - expected))
  let ulpError = absoluteError / T(expected).ulp
  expectTrue(ulpError <= allowed,
             "\(actual) != \(expected) as \(T.self)" +
             "\n  \(ulpError)-ulp error exceeds \(allowed)-ulp tolerance.",
             file: file, line: line)
}

func computeDividedDifference<T: BinaryFloatingPoint> (
  _ f: (T, T) -> T,
  _ x: T,
  _ y: T,
  eps: T = 0.01
) -> (dfdx: T, dfdy: T) {
  let dfdx = (f(x + eps, y) - f(x, y)) / eps
  let dfdy = (f(x, y + eps) - f(x, y)) / eps
  return (dfdx, dfdy)
}

func checkGradient<T: BinaryFloatingPoint & Differentiable>(
  _ f: @differentiable(reverse) (T, T) -> T,
  _ x: T,
  _ y: T,
  ulps: T = 192)
where T == T.TangentVector {
  let eps = T(0.01)
  let grad = gradient(at: x, y, of: f)
  let (dfdx, dfdy) = computeDividedDifference(f, x, y, eps: eps)
  expectEqualWithTolerance(TestLiteralType(dfdx), grad.0, ulps: ulps)
  expectEqualWithTolerance(TestLiteralType(dfdy), grad.1, ulps: ulps)
}

func checkDerivative<T: BinaryFloatingPoint & Differentiable>(
  _ f: @differentiable(reverse) (T, T) -> T,
  _ x: T,
  _ y: T,
  ulps: T = 192)
where T == T.TangentVector {
  let eps = T(0.01)
  let deriv = derivative(at: x, y, of: f)
  let (dfdx, dfdy) = computeDividedDifference(f, x, y, eps: eps)
  expectEqualWithTolerance(TestLiteralType(dfdx + dfdy), deriv, ulps: ulps)
}

%for op in ['derivative', 'gradient']:
%for T in ['Float', 'Float80']:

%if T == 'Float80':
#if !(os(Windows) || os(Android)) && (arch(i386) || arch(x86_64))
%end

DerivativeTests.test("${op}_${T}") {
  expectEqualWithTolerance(7.3890560989306502274, ${op}(at: 2 as ${T}, of: exp))
  expectEqualWithTolerance(2.772588722239781145, ${op}(at: 2 as ${T}, of: exp2))
  expectEqualWithTolerance(7.3890560989306502274, ${op}(at: 2 as ${T}, of: expm1))
  expectEqualWithTolerance(0.5, ${op}(at: 2 as ${T}, of: log))
  expectEqualWithTolerance(0.21714724095162590833, ${op}(at: 2 as ${T}, of: log10))
  expectEqualWithTolerance(0.7213475204444817278, ${op}(at: 2 as ${T}, of: log2))
  expectEqualWithTolerance(0.33333333333333333334, ${op}(at: 2 as ${T}, of: log1p))
  expectEqualWithTolerance(5.774399204041917612, ${op}(at: 2 as ${T}, of: tan))
  expectEqualWithTolerance(-0.9092974268256816954, ${op}(at: 2 as ${T}, of: cos))
  expectEqualWithTolerance(-0.416146836547142387, ${op}(at: 2 as ${T}, of: sin))
  expectEqualWithTolerance(1.154700538379251529, ${op}(at: 0.5 as ${T}, of: asin))
  expectEqualWithTolerance(-1.154700538379251529, ${op}(at: 0.5 as ${T}, of: acos))
  expectEqualWithTolerance(0.8, ${op}(at: 0.5 as ${T}, of: atan))
  expectEqualWithTolerance(3.7621956910836314597, ${op}(at: 2 as ${T}, of: sinh))
  expectEqualWithTolerance(3.6268604078470187677, ${op}(at: 2 as ${T}, of: cosh))
  expectEqualWithTolerance(0.07065082485316446565, ${op}(at: 2 as ${T}, of: tanh))
  expectEqualWithTolerance(0.44721359549995793928, ${op}(at: 2 as ${T}, of: asinh))
  expectEqualWithTolerance(0.5773502691896257645, ${op}(at: 2 as ${T}, of: acosh))
  expectEqualWithTolerance(1.3333333333333333334, ${op}(at: 0.5 as ${T}, of: atanh))
  expectEqualWithTolerance(0.020666985354092053575, ${op}(at: 2 as ${T}, of: erf))
  expectEqualWithTolerance(-0.020666985354092053575, ${op}(at: 2 as ${T}, of: erfc))
  expectEqualWithTolerance(0.35355339059327376222, ${op}(at: 2 as ${T}, of: { sqrt($0) }))
  expectEqualWithTolerance(0, ${op}(at: 2 as ${T}, of: { ceil($0) }))
  expectEqualWithTolerance(0, ${op}(at: 2 as ${T}, of: { floor($0) }))
  expectEqualWithTolerance(0, ${op}(at: 2 as ${T}, of: { round($0) }))
  expectEqualWithTolerance(0, ${op}(at: 2 as ${T}, of: { trunc($0) }))

  // Differential operator specific tests.

  // fma
  let dfma = ${op}(at: 4 as ${T}, 5 as ${T}, 6 as ${T}, of: fma)
%if op == 'gradient':
  expectEqualWithTolerance(5, dfma.0)
  expectEqualWithTolerance(4, dfma.1)
  expectEqualWithTolerance(1, dfma.2)
%else: # if op == 'derivative'
  expectEqualWithTolerance(10, dfma)
%end

  // remainder, fmod
  for a in -10...10 {
    let x = ${T}(a)
    for b in -10...10 {
      let y = ${T}(b)
      guard b != 0 && remainder(x, y).sign == remainder(x + ${T}(0.001), y).sign &&
        remainder(x, y).sign == remainder(x, y + ${T}(0.001)).sign
        else { continue }
%if op == 'gradient':
      checkGradient({ remainder($0, $1) }, x, y)
      checkGradient({ fmod($0, $1) }, x, y)
%else: # if op == 'derivative'
      // TODO(TF-1108): Implement JVPs for `remainder` and `fmod`.
%end
    }
  }

  // pow
  let eps:${T} = 0.01
  let ulps:${T} = eps/eps.ulp

  // Checks for negative base.
  for a in -3..<0 {
    let x = ${T}(a)
    for b in -3...3 {
      let y = ${T}(b)
      let expectedDx =  y * pow(x, y - 1)
      let expectedDy = ${T}.zero
      let dpow = ${op}(at: x, y, of: pow)
%if op == 'gradient':
      expectEqualWithTolerance(TestLiteralType(expectedDx), dpow.0)
      expectEqualWithTolerance(TestLiteralType(expectedDy), dpow.1)
%else: # if op == 'derivative'
      expectEqualWithTolerance(TestLiteralType(expectedDx + expectedDy), dpow)
%end
    }
  }

  // Checks for 0 base.
  for b in -3...3 {
    let y = ${T}(b)
    var expectedValues: (dx: ${T}, dy: ${T})?
    if y.isLess(than: 0) {
      expectedValues = (dx: ${T}.infinity, dy: ${T}.nan)
    } else if y.isZero {
      expectedValues = (dx: ${T}.nan, dy: ${T}.zero)
    } else if !y.isEqual(to: 1) {
      expectedValues = (dx: ${T}.zero, dy: ${T}.zero)
    }
    if let (expectedDx, expectedDy) = expectedValues {
      let dpow = ${op}(at: 0.0, y, of: pow)
%if op == 'gradient':
      expectEqualWithTolerance(TestLiteralType(expectedDx), dpow.0)
      expectEqualWithTolerance(TestLiteralType(expectedDy), dpow.1)
%else: # if op == 'derivative'
      expectEqualWithTolerance(TestLiteralType(expectedDx + expectedDy), dpow)
%end
    } else {
%if op == 'gradient':
      checkGradient({ pow($0, $1) }, 0.0, y, ulps: ulps)
%else: # if op == 'derivative'
      checkDerivative({ pow($0, $1) }, 0.0, y, ulps: ulps)
%end
    }
  }

  // Checks for positive base.
  for a in 1...3 {
    let x = ${T}(a)
    for b in -3...3 {
      let y = ${T}(b)
%if op == 'gradient':
      checkGradient({ pow($0, $1) }, x, y, ulps: ulps)
%else: # if op == 'derivative'
      checkDerivative({ pow($0, $1) }, x, y, ulps: ulps)
%end
    }
  }
}

%if T == 'Float80':
#endif
%end

%end # for T in ['Float', 'Float80']:
%end # for op in ['derivative', 'gradient']:

runAllTests()