/* -----------------------------------------------------------------------------
Copyright 2021,2023 Kevin P. Barry

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.
----------------------------------------------------------------------------- */

// Author: Kevin P. Barry [ta0kira@gmail.com]

testcase "valid instantiations" {
  success TestChecker
}

unittest exponential {
  Generator<Float> random <- RandomExponential.new(1.0)
  Float value <- random.generate()
}

unittest gaussian {
  Generator<Float> random <- RandomGaussian.new(0.0, 1.0)
  Float value <- random.generate()
}

unittest uniform {
  Generator<Float> random <- RandomUniform.new(0.0, 1.0)
  Float value <- random.generate()
}

unittest probability {
  Generator<Float> random <- RandomUniform.probability()
  Float value <- random.generate()
}

unittest constant {
  Generator<Float> random <- GenerateConstant.new(1.0)
  Float value <- random.generate()
}

unittest categorical {
  CategoricalTree<Float> tree <- CategoricalTree<Float>.new()
      .setWeight(0.0, 1)
      .setWeight(1.0, 1)
      .setWeight(2.0, 1)
      .setWeight(3.0, 1)
      .setWeight(4.0, 1)
  RandomCategorical<Float> random <- tree `RandomCategorical:sampleWith` RandomUniform.probability()
  \ random.isEmpty() `Matches:with` CheckValue:equals(false)
  Float value <- random.generate()
}


testcase "negative lambda in exponential" {
  failure
  require "lambda"
  require "-1"
}

unittest test {
  Generator<Float> random <- RandomExponential.new(-1.0)
}


testcase "negative standard deviation in gaussian" {
  failure
  require "standard deviation"
  require "-1"
}

unittest test {
  Generator<Float> random <- RandomGaussian.new(0.0, -1.0)
}


testcase "empty range in uniform" {
  failure
  require "range"
  require "0,-1"
}

unittest test {
  Generator<Float> random <- RandomUniform.new(0.0, -1.0)
}


testcase "sampling with empty CategoricalTree" {
  failure TestChecker
  require "less than the total"
}

unittest categorical {
  CategoricalTree<Float> tree <- CategoricalTree<Float>.new()
  RandomCategorical<Float> random <- tree `RandomCategorical:sampleWith` RandomUniform.probability()
  \ random.isEmpty() `Matches:with` CheckValue:equals(true)
  Float value <- random.generate()
}


testcase "distribution sanity checks" {
  success TestChecker
}

unittest exponential { $DisableCoverage$
  Int count <- 10000
  Float lambda <- 10.0
  $ReadOnly[count, lambda]$

  Generator<Float> random <- RandomExponential.new(lambda)

  Float sum <- 0.0
  traverse (Counter.zeroIndexed(count) -> _) {
    Float value <- random.generate()
    \ value `Matches:with` CheckValue:greaterThanNotEquals(0.0)
    sum <- sum+value
  }

  \ sum/count.asFloat() `Matches:with` CheckValue:betweenEquals(0.9/lambda, 1.1/lambda)
}

unittest gaussian { $DisableCoverage$
  Int count <- 10000
  Float mean <- 100.0
  Float sd <- 1.0
  $ReadOnly[count, mean, sd]$

  Generator<Float> random <- RandomGaussian.new(mean, sd)

  Float sum <- 0.0
  traverse (Counter.zeroIndexed(count) -> _) {
    sum <- sum+random.generate()
  }

  \ sum/count.asFloat() `Matches:with` CheckValue:betweenEquals(mean-0.1*sd, mean+0.1*sd)
}

unittest uniform { $DisableCoverage$
  Int count <- 10000
  Float min <- -11.0
  Float max <- 5.0
  $ReadOnly[count, min, max]$

  Generator<Float> random <- RandomUniform.new(min, max)

  Float sum <- 0.0
  traverse (Counter.zeroIndexed(count) -> _) {
    Float value <- random.generate()
    \ value `Matches:with` CheckValue:betweenEquals(min, max)
    sum <- sum+value
  }

  \ sum/count.asFloat() `Matches:with` CheckValue:betweenEquals((min+max)/2.0-0.1*(max-min), (min+max)/2.0+0.1*(max-min))
}

unittest probability { $DisableCoverage$
  Int count <- 10000
  $ReadOnly[count]$

  Generator<Float> random <- RandomUniform.probability()

  Float sum <- 0.0
  traverse (Counter.zeroIndexed(count) -> _) {
    Float value <- random.generate()
    \ value `Matches:with` CheckValue:betweenEquals(0.0, 1.0)
    sum <- sum+value
  }

  \ sum/count.asFloat() `Matches:with` CheckValue:betweenEquals(0.4, 0.6)
}

unittest constant { $DisableCoverage$
  Int count <- 10000
  Float constant <- 12345.0
  $ReadOnly[count, constant]$

  Generator<Float> random <- GenerateConstant.new(constant)

  traverse (Counter.zeroIndexed(count) -> _) {
    Float value <- random.generate()
    \ constant `Matches:with` CheckValue:equals(value)
  }
}

unittest categorical { $DisableCoverage$
  Int count <- 10000
  CategoricalTree<Float> tree <- CategoricalTree<Float>.new()
      .setWeight(0.0, 1)
      .setWeight(1.0, 1)
      .setWeight(2.0, 1)
      .setWeight(3.0, 1)
      .setWeight(4.0, 1)
  $ReadOnly[count, tree]$

  Generator<Float> random <- tree `RandomCategorical:sampleWith` RandomUniform.probability()

  Float sum <- 0.0
  traverse (Counter.zeroIndexed(count) -> _) {
    Float value <- random.generate()
    \ value `Matches:with` CheckValue:betweenEquals(0.0, 4.0)
    sum <- sum+value
  }

  \ sum/count.asFloat() `Matches:with` CheckValue:betweenEquals(1.9, 2.1)
}


testcase "distribution seed checks" {
  success TestChecker
}

unittest exponential {
  Float lambda <- 10.0
  $ReadOnly[lambda]$

  RandomExponential random1 <- RandomExponential.new(lambda)
  RandomExponential random2 <- RandomExponential.new(lambda)

  Float v1 <- random1.generate()
  Float v2 <- random2.generate()
  random1 <- random1.setSeed(123)
  Float v3 <- random1.generate()
  random2 <- random2.setSeed(123)
  Float v4 <- random2.generate()

  \ v1 `Matches:with` CheckValue:notEquals(v2)
  \ v1 `Matches:with` CheckValue:notEquals(v3)
  \ v3 `Matches:with` CheckValue:equals(v4)
}

unittest gaussian {
  Float mean <- 100.0
  Float sd <- 1.0
  $ReadOnly[mean, sd]$

  RandomGaussian random1 <- RandomGaussian.new(mean, sd)
  RandomGaussian random2 <- RandomGaussian.new(mean, sd)

  Float v1 <- random1.generate()
  Float v2 <- random2.generate()
  random1 <- random1.setSeed(123)
  Float v3 <- random1.generate()
  random2 <- random2.setSeed(123)
  Float v4 <- random2.generate()

  \ v1 `Matches:with` CheckValue:notEquals(v2)
  \ v1 `Matches:with` CheckValue:notEquals(v3)
  \ v3 `Matches:with` CheckValue:equals(v4)
}

unittest uniform {
  Float min <- -11.0
  Float max <- 5.0
  $ReadOnly[min, max]$

  RandomUniform random1 <- RandomUniform.new(min, max)
  RandomUniform random2 <- RandomUniform.new(min, max)

  Float v1 <- random1.generate()
  Float v2 <- random2.generate()
  random1 <- random1.setSeed(123)
  Float v3 <- random1.generate()
  random2 <- random2.setSeed(123)
  Float v4 <- random2.generate()

  \ v1 `Matches:with` CheckValue:notEquals(v2)
  \ v1 `Matches:with` CheckValue:notEquals(v3)
  \ v3 `Matches:with` CheckValue:equals(v4)
}


testcase "Randomize tests" {
  success TestChecker
}

unittest permuteFrom {
  CategoricalTree<String> tree <- CategoricalTree<String>.new()
      .setWeight("a", 2)
      .setWeight("b", 3)
      .setWeight("c", 1)

  Vector<String> expected <- Vector<String>.new()
      .append("a")
      .append("a")
      .append("b")
      .append("b")
      .append("b")
      .append("c")

  Vector<String> output <- Vector<String>.new()
  \ Randomize:permuteFrom(tree, RandomUniform.new(0.0, 1.0), output)
  \ Sorting:sort(output)

  \ output.size() `Matches:with` CheckValue:equals(expected.size())
  traverse (Counter.zeroIndexed(output.size()) -> Int pos) {
    \ output.readAt(pos) `Matches:with` CheckValue:equals(expected.readAt(pos))
  }

  \ tree.getWeight("a") `Matches:with` CheckValue:equals(2)
  \ tree.getWeight("b") `Matches:with` CheckValue:equals(3)
  \ tree.getWeight("c") `Matches:with` CheckValue:equals(1)
}

unittest permuteFromWeight {
  CategoricalTree<String> tree <- CategoricalTree<String>.new()
      .setWeight("a", 2)
      .setWeight("b", 3)
      .setWeight("c", 1)

  Vector<String> expected <- Vector<String>.new()
      .append("a")
      .append("b")
      .append("c")

  Vector<String> output <- Vector<String>.new()
  \ Randomize:permuteFromWeight(tree, RandomUniform.new(0.0, 1.0), output)
  \ Sorting:sort(output)

  \ output.size() `Matches:with` CheckValue:equals(expected.size())
  traverse (Counter.zeroIndexed(output.size()) -> Int pos) {
    \ output.readAt(pos) `Matches:with` CheckValue:equals(expected.readAt(pos))
  }

  \ tree.getWeight("a") `Matches:with` CheckValue:equals(2)
  \ tree.getWeight("b") `Matches:with` CheckValue:equals(3)
  \ tree.getWeight("c") `Matches:with` CheckValue:equals(1)
}