package ope

import (
	"fmt"
	"math"
	"math/big"
	"testing"

	"github.com/stretchr/testify/assert"
	"xdx.jelly/xgcl/grand"
)

func choose(y int64, x int64) float64 {
	a, _ := new(big.Int).MulRange(y-x+1, y).Float64()
	b, _ := new(big.Int).MulRange(2, x).Float64()
	return a / b
}

func hyperProb(x int64, M int64, N int64, y int64) float64 {
	return choose(M, x) * choose(N-M, y-x) / choose(N, y)
}

func TestUniform(t *testing.T) {
	stats := make([]int, 10)
	for i := 0; i < 1000000; i++ {
		f := uniformRand(bytesToCoins(grand.GetRandom(4)))
		stats[int(f*10)]++
	}

	fmt.Println(stats)
}

func TestHyperSmall(t *testing.T) {
	M := int64(4)
	N := int64(10)
	y := int64(5)
	stats := make([]float64, min(y, M)+1)
	wants := make([]float64, min(y, M)+1)
	count := 1000000
	for i := 0; i < count; i++ {
		x := hypergeometricSmall(M, N, y, bytesToCoins(grand.GetRandom(4)))
		stats[x] += 1
	}

	for i := 0; i < len(stats); i++ {
		stats[i] = stats[i] / float64(count)
		wants[i] = hyperProb(int64(i), M, N, y)
	}
	if true {
		for i := 0; i < len(stats); i++ {
			fmt.Printf("%.2f ", stats[i])
		}
		fmt.Println()
		for i := 0; i < len(stats); i++ {
			fmt.Printf("%.2f ", wants[i])
		}
		fmt.Println()
	}
	epsilon := 0.001
	for i := 0; i < len(stats); i++ {
		assert.Less(t, math.Abs(stats[i]-wants[i]), epsilon)
	}
}

func TestHyperBig(t *testing.T) {
	M := int64(30)
	N := int64(100)
	y := int64(51)
	stats := make([]float64, min(y, M)+1)
	wants := make([]float64, min(y, M)+1)
	count := 1000000
	for i := 0; i < count; i++ {
		x := hypergeometricBig(M, N, y, bytesToCoins(grand.GetRandom(4)))
		stats[x] += 1
	}

	for i := 0; i < len(stats); i++ {
		stats[i] = stats[i] / float64(count)
		wants[i] = hyperProb(int64(i), M, N, y)
	}

	epsilon := 0.001
	for i := 0; i < len(stats); i++ {
		assert.Less(t, math.Abs(stats[i]-wants[i]), epsilon)
	}

}