package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"runtime"
	"sort"
	"strconv"
	"sync"
)

type Result struct {
	N             int
	K             int
	Trial         int
	KthValue      float64
	NHat          float64
	RelativeError float64
}

func runTrial(n, k, trial int, rng *rand.Rand) Result {
	// Generate n random numbers
	samples := make([]float64, n)
	for i := 0; i < n; i++ {
		samples[i] = rng.Float64()
	}

	// Sort them
	sort.Float64s(samples)

	// k-th element (1-indexed, so index k-1)
	kthValue := samples[k-1]

	// Estimate: n_hat = k / kth_value - 1
	var nHat float64
	if kthValue > 0 {
		nHat = float64(k)/kthValue - 1
	} else {
		nHat = float64(n) * 1000
	}

	relativeError := (nHat - float64(n)) / float64(n)

	return Result{
		N:             n,
		K:             k,
		Trial:         trial,
		KthValue:      kthValue,
		NHat:          nHat,
		RelativeError: relativeError,
	}
}

func main() {
	// Parse command line arguments
	maxN := flag.Int("max-n", 1_000_000_000, "maximum value of n")
	maxK := flag.Int("max-k", 16384, "maximum value of k")
	numTrials := flag.Int("trials", 100, "number of trials per (n,k) pair")
	flag.Parse()

	fmt.Printf("Config: max_n=%d, max_k=%d, trials=%d\n", *maxN, *maxK, *numTrials)

	// n values: 10000 doubling up to max_n
	nValues := []int{}
	for n := 10000; n <= *maxN; n *= 2 {
		nValues = append(nValues, n)
	}

	// k values: 1, 2, 4, ... up to max_k
	kValues := []int{}
	for k := 1; k <= *maxK; k *= 2 {
		kValues = append(kValues, k)
	}

	// Use all available cores
	numWorkers := runtime.NumCPU()
	fmt.Printf("Using %d workers\n", numWorkers)

	// Create work items
	type WorkItem struct {
		N     int
		K     int
		Trial int
	}

	workItems := []WorkItem{}
	for _, n := range nValues {
		for _, k := range kValues {
			if k > n {
				continue
			}
			for trial := 0; trial < *numTrials; trial++ {
				workItems = append(workItems, WorkItem{N: n, K: k, Trial: trial})
			}
		}
	}

	fmt.Printf("Total work items: %d\n", len(workItems))

	// Results channel
	results := make(chan Result, 10000)

	// Work channel
	work := make(chan WorkItem, numWorkers*2)

	// Start workers
	var wg sync.WaitGroup
	for w := 0; w < numWorkers; w++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			// Each worker has its own RNG seeded differently
			rng := rand.New(rand.NewSource(int64(workerID * 1000000)))
			for item := range work {
				result := runTrial(item.N, item.K, item.Trial, rng)
				results <- result
			}
		}(w)
	}

	// Start result writer goroutine
	var writerWg sync.WaitGroup
	writerWg.Add(1)
	go func() {
		defer writerWg.Done()

		file, err := os.Create("results.csv")
		if err != nil {
			panic(err)
		}
		defer file.Close()

		writer := csv.NewWriter(file)
		defer writer.Flush()

		// Header
		writer.Write([]string{"n", "k", "trial", "kth_value", "n_hat", "relative_error"})

		count := 0
		for result := range results {
			writer.Write([]string{
				strconv.Itoa(result.N),
				strconv.Itoa(result.K),
				strconv.Itoa(result.Trial),
				fmt.Sprintf("%.15f", result.KthValue),
				fmt.Sprintf("%.2f", result.NHat),
				fmt.Sprintf("%.6f", result.RelativeError),
			})
			count++
			if count%10000 == 0 {
				fmt.Printf("Written %d results\n", count)
				writer.Flush()
			}
		}
		fmt.Printf("Total written: %d results\n", count)
	}()

	// Feed work to workers
	totalItems := len(workItems)
	for i, item := range workItems {
		work <- item
		if (i+1)%10000 == 0 {
			fmt.Printf("Queued %d/%d (%.1f%%) - current n=%d\n", i+1, totalItems, float64(i+1)/float64(totalItems)*100, item.N)
		}
	}
	close(work)

	// Wait for workers to finish
	wg.Wait()
	close(results)

	// Wait for writer to finish
	writerWg.Wait()

	fmt.Println("Done! Results written to results.csv")
}