package main

import "fmt"
import "git.nunosempere.com/NunoSempere/probppl/choose"
import "math"
import rand "math/rand/v2"

type src = *rand.Rand
type pplKnownDistrib = map[int64]float64

func generatePeopleKnownDistribution(r src) map[int64]float64 {
	mapping := make(map[int64]float64)
	sum := 0.0

	// Consider zero case separately
	p0 := r.Float64()
	mapping[0.0] = p0
	sum += p0

	// Consider successive exponents of 1.5
	l := 1.0
	m := 1.5
	for i := 1; i < 22; i++ {
		l = l * m
		p := r.Float64()
		mapping[int64(l)] = p
		sum += p
	}

	for key, value := range mapping {
		mapping[key] = value / sum
	}

	return mapping
}

func chooseWrapper(n int64, k int64) int64 {
	if n < k {
		return 0
	} else {
		return choose.Choose(n, k)
	}
}

func getProbabilityOfKBirthdayMatchesGivenNPeopleKnown(n int64, k int64) float64 {
	return float64(chooseWrapper(n, k)) * math.Pow(1.0/365.0, float64(k)) * math.Pow(1.0-(1.0/365.0), float64(n-k))
}

func getMatchesDrawGivenNPeopleKnown(n int64, r src) int64 {
	p0 := getProbabilityOfKBirthdayMatchesGivenNPeopleKnown(n, 0)
	p1 := getProbabilityOfKBirthdayMatchesGivenNPeopleKnown(n, 1)
	p2 := getProbabilityOfKBirthdayMatchesGivenNPeopleKnown(n, 2)

	p := r.Float64()
	if p < p0 {
		return 0
	} else if p < (p0 + p1) {
		return 1
	} else if p < (p0 + p1 + p2) {
		return 2
	} else {
		return 3 // stands for 'greater than 3'
	}
}

/*
  Draw 148 times
  How many people do you know that were born in the same day of the year as you?
	0: 46.6% | 69
	1: 31.1% | 46
	2: 12.8% | 19
	≥3: 9.5% | 14
*/

func drawFromDistributionWithReplacement(d pplKnownDistrib, r src) int64 {
	pp := r.Float64()
	sum := 0.0
	for i, p := range d {
		sum += p
		if p <= sum {
			return int64(i)
		}

	}
	fmt.Printf("%f, %f\n", sum, pp)
	fmt.Println(d)
	panic("Probabilities should sum up to 1")
}

func draw148PplFromDistributionAndCheck(d pplKnownDistrib, r src) int64 {

	count := make(map[int64]int64)
	count[0] = 0
	count[1] = 0
	count[2] = 0
	count[3] = 0
	for i := 0; i < 148; i++ {
		person_i_ppl_known := drawFromDistributionWithReplacement(d, r)
		person_i_num_birthday_matches := getMatchesDrawGivenNPeopleKnown(person_i_ppl_known, r)
		count[person_i_num_birthday_matches]++
	}
	if (count[0] == 69) && (count[1] == 46) && (count[2] == 19) && (count[3] == 14) {
		return 1
	} else {
		return 0
	}
}

func getUnnormalizedBayesianUpdateForDistribution(d pplKnownDistrib, r src) float64 {
	var sum int64 = 0
	n := 1_000
	for i := 0; i < n; i++ {
		sum += draw148PplFromDistributionAndCheck(d, r)
	}
	return float64(sum) / float64(n)
}

func main() {
	fmt.Println("Hello world")

	var r = rand.New(rand.NewPCG(uint64(1), uint64(2)))
	people_known_distribution := generatePeopleKnownDistribution(r)
	/* for i, p := range people_known_distribution {
		fmt.Print64f("%d: %.4f\n", i, p)
	} */
	result := getUnnormalizedBayesianUpdateForDistribution(people_known_distribution, r)
	fmt.Println(result)

}