package main

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

type src = *rand.Rand
type IntProb struct {
	N int64
	p float64
}
type IntProbs = []IntProb
type IntProbsWeights struct {
	IntProbs IntProbs
	w        int64
}

func generatePeopleKnownDistribution(r src) IntProbs {

	var probabilities IntProbs
	sum := 0.0

	for i := 16; i <= 2048; i *= 2 {
		p := r.Float64()
		probabilities = append(probabilities, IntProb{N: int64(i), p: p})
		sum += p
	}

	for i := range probabilities {
		probabilities[i].p = probabilities[i].p / sum
	}

	return probabilities
}

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'
	}
}

func drawFromDistributionWithReplacement(d IntProbs, r src) int64 {
	pp := r.Float64()
	sum := 0.0
	for i := range d {
		sum += d[i].p
		if pp <= sum {
			return int64(d[i].N) // this introduces some non-determinism, as order of maps in go isn't guaranteed
		}

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

func aboutEq(a int64, b int64) bool {
	h := int64(3)
	return ((-h) <= (a - b)) && ((a - b) <= h)
}

func draw148PplFromDistributionAndCheck(d IntProbs, r src, show bool) 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]++
	}
	/*
	  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
	*/
	if aboutEq(count[0], 69) && aboutEq(count[1], 46) && aboutEq(count[2], 19) && aboutEq(count[3], 14) {
		return 1
	} else {
		return 0
	}
}

func getUnnormalizedBayesianUpdateForDistribution(d IntProbs, r src) int64 {
	var sum int64 = 0
	n := 10_000
	for i := 0; i < n; i++ {
		draw_result := draw148PplFromDistributionAndCheck(d, r, i == 0)
		sum += draw_result
	}
	return sum
}

func main() {

	n_dists := 100_000
	var dists = make([]IntProbsWeights, n_dists)

	// Prepare for concurrency
	num_threads := 32
	var wg sync.WaitGroup
	wg.Add(num_threads)

	for i := range num_threads {
		go func() {
			defer wg.Done()
			var r = rand.New(rand.NewPCG(uint64(i), uint64(i+1)))

			for j := i * (n_dists / num_threads); j < (i+1)*(n_dists/num_threads); j++ {

				// Get indivisual distributions
				people_known_distribution := generatePeopleKnownDistribution(r)
				result := getUnnormalizedBayesianUpdateForDistribution(people_known_distribution, r)
				if result > 0 {
					dists[j] = IntProbsWeights{IntProbs: people_known_distribution, w: result}
				} else {
					dists[j] = IntProbsWeights{IntProbs: nil, w: 0}
					// make explicit that this is initialized to nil; we won't compute it later becasue weight is 0
				}
			}
		}()
	}
	wg.Wait()

	sum_weights := int64(0)
	for _, dist := range dists {
		sum_weights += dist.w
	}

	fmt.Printf("\n| Num people known | %% |\n")
	fmt.Printf("| ---- | ----- |\n")
	for i := int64(16); i <= 2048; i *= 2 {
		p := 0.0
		for _, dist := range dists {
			for _, int_prob := range dist.IntProbs {
				if int_prob.N == i {
					p += float64(dist.w) * int_prob.p
				}
			}
		}
		p = p / float64(sum_weights)
		fmt.Printf("| %4d | %4.1f%% | \n", i, p*100)
	}
}