packsim.go

package main

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

type UnpaddedPieceSize uint64
type PaddedPieceSize uint64

func (s PaddedPieceSize) Unpadded() UnpaddedPieceSize {
	return UnpaddedPieceSize(uint64(s) / 128 * 127)
}

func PaddedSize(size uint64) UnpaddedPieceSize {
	if size <= 127 {
		return UnpaddedPieceSize(127)
	}
	paddedPieceSize := (size + 126) / 127 * 128
	if bits.OnesCount64(paddedPieceSize) != 1 {
		paddedPieceSize = 1 << uint(64-bits.LeadingZeros64(paddedPieceSize))
	}
	return PaddedPieceSize(paddedPieceSize).Unpadded()
}

func alignedOffset(offset, paddedSize uint64) uint64 {
	if offset%paddedSize != 0 {
		return ((offset / paddedSize) + 1) * paddedSize
	}
	return offset
}

// Config
var (
	csvPath              string
	numRuns              int
	numSlicesPerPartition int
	numWorkers           int
	containerSizeGiB     float64
	partitionMultiplier  float64
	skipLognormal        bool
	pickN                int
)

var containerSize uint64

type StrategyStats struct {
	Name       string
	Totals     []float64
	Mean       float64
	Std        float64
	MeanPacked float64
}

type scoreFunc func(offset uint64, blockSize int64) int64

type strategyDef struct {
	name   string
	scorer scoreFunc
	padded bool
}

func getStrategies() []strategyDef {
	return []strategyDef{
		{"1. Padded sequential", nil, true},
		{"2. Unpadded sequential", nil, false},
		{fmt.Sprintf("3a. Pick-%d min align (unpadded)", pickN), scoreMinAlignmentWaste, false},
		{fmt.Sprintf("3b. Pick-%d largest (unpadded)", pickN), scoreLargestFirst, false},
		{fmt.Sprintf("3c. Pick-%d min waste (unpadded)", pickN), scoreMinTotalWaste, false},
		{fmt.Sprintf("3d. Pick-%d small first (unpadded)", pickN), scoreSmallestPaddedFirst, false},
		{fmt.Sprintf("4a. Pick-%d min align (padded)", pickN), scoreMinAlignmentWastePadded, true},
		{fmt.Sprintf("4b. Pick-%d largest (padded)", pickN), scoreLargestFirstPadded, true},
		{fmt.Sprintf("4c. Pick-%d min waste (padded)", pickN), scoreMinTotalWastePadded, true},
		{fmt.Sprintf("4d. Pick-%d small first (padded)", pickN), scoreSmallestPaddedFirstPadded, true},
	}
}

// VirtualPartition holds indices into the blocks array for one partition
type VirtualPartition struct {
	indices []int
}

func main() {
	flag.StringVar(&csvPath, "csv", "../data_dist_raw.csv", "Path to CSV file (num_occurrences,block_size)")
	flag.IntVar(&numRuns, "n", 190, "Number of test runs/partitions")
	flag.IntVar(&numSlicesPerPartition, "slices", 64, "Number of slices per virtual partition")
	flag.IntVar(&numWorkers, "workers", runtime.NumCPU(), "Number of parallel workers")
	flag.Float64Var(&containerSizeGiB, "container", 0, "Container size in GiB (default: (32<<30)/128*127 bytes)")
	flag.Float64Var(&partitionMultiplier, "partition-mult", 4.0, "Target partition size as multiple of container")
	flag.BoolVar(&skipLognormal, "skip-lognormal", false, "Skip log-normal distribution test")
	flag.IntVar(&pickN, "pick", 10, "Number of candidates to pick from in pick-N strategies")
	flag.Parse()

	if containerSizeGiB > 0 {
		containerSize = uint64(containerSizeGiB * (1 << 30))
	} else {
		containerSize = uint64((32 << 30) / 128 * 127)
	}

	runtime.GOMAXPROCS(numWorkers)

	fmt.Println("Loading data...")
	file, err := os.Open(csvPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error opening CSV: %v\n", err)
		os.Exit(1)
	}

	reader := csv.NewReader(file)
	records, err := reader.ReadAll()
	file.Close()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error reading CSV: %v\n", err)
		os.Exit(1)
	}

	var realBlocks []int64
	totalSize := int64(0)
	for _, record := range records {
		if len(record) != 2 {
			continue
		}
		numOccurrences, _ := strconv.ParseInt(record[0], 10, 64)
		blockSize, _ := strconv.ParseInt(record[1], 10, 64)
		if blockSize <= 0 {
			continue
		}
		for i := int64(0); i < numOccurrences; i++ {
			realBlocks = append(realBlocks, blockSize)
			totalSize += blockSize
		}
	}
	records = nil
	runtime.GC()

	meanSize := float64(totalSize) / float64(len(realBlocks))
	fmt.Printf("Real data: %d blocks, mean=%.2f bytes (%.2f KiB)\n", len(realBlocks), meanSize, meanSize/1024)

	// Build virtual partitions: each partition samples from slices across the array
	targetBytesPerPartition := float64(containerSize) * partitionMultiplier
	targetBlocksPerPartition := int(targetBytesPerPartition / meanSize)
	blocksPerSlice := targetBlocksPerPartition / numSlicesPerPartition

	fmt.Printf("Building %d virtual partitions from %d slices each...\n", numRuns, numSlicesPerPartition)
	fmt.Printf("Target blocks per partition: %d (~%.2f GiB)\n", targetBlocksPerPartition, targetBytesPerPartition/(1<<30))

	// Create virtual partitions
	realPartitions := buildVirtualPartitions(len(realBlocks), numRuns, numSlicesPerPartition, blocksPerSlice)

	// Verify partition sizes
	var partitionDataSum int64
	for _, idx := range realPartitions[0].indices {
		partitionDataSum += realBlocks[idx]
	}
	fmt.Printf("Actual partition 0: %d blocks, %.2f GiB\n", len(realPartitions[0].indices), float64(partitionDataSum)/(1<<30))

	fmt.Printf("Container size: %.2f GiB\n", float64(containerSize)/(1<<30))
	fmt.Printf("Running %d iterations with %d workers, pick-%d...\n\n", numRuns, numWorkers, pickN)

	fmt.Println("=== RUN 1: REAL DATA ===")
	runAllTests(realBlocks, realPartitions, numWorkers)

	if !skipLognormal {
		fmt.Println("\nGenerating log-normal data...")
		sigma := 2.0
		mu := math.Log(meanSize) - (sigma*sigma)/2
		lognormalBlocks := make([]int64, len(realBlocks))
		for i := range lognormalBlocks {
			sample := math.Exp(mu + sigma*rand.NormFloat64())
			lognormalBlocks[i] = int64(math.Max(1, sample))
		}

		var lnTotal int64
		for _, b := range lognormalBlocks {
			lnTotal += b
		}
		lnMean := float64(lnTotal) / float64(len(lognormalBlocks))
		fmt.Printf("Log-normal data: %d blocks, mean=%.2f bytes\n", len(lognormalBlocks), lnMean)

		targetBlocksLN := int(targetBytesPerPartition / lnMean)
		blocksPerSliceLN := targetBlocksLN / numSlicesPerPartition
		lnPartitions := buildVirtualPartitions(len(lognormalBlocks), numRuns, numSlicesPerPartition, blocksPerSliceLN)

		fmt.Println("\n=== RUN 2: LOG-NORMAL DATA ===")
		runAllTests(lognormalBlocks, lnPartitions, numWorkers)
	}
}

func buildVirtualPartitions(totalBlocks, numPartitions, slicesPerPartition, blocksPerSlice int) []VirtualPartition {
	partitions := make([]VirtualPartition, numPartitions)

	// Total slices needed
	totalSlices := numPartitions * slicesPerPartition

	// Distribute slices across the array
	// Each slice starts at: sliceIdx * (totalBlocks / totalSlices)
	sliceStride := totalBlocks / totalSlices

	for p := 0; p < numPartitions; p++ {
		partitions[p].indices = make([]int, 0, slicesPerPartition*blocksPerSlice)

		for s := 0; s < slicesPerPartition; s++ {
			// This partition's s-th slice comes from position (p + s*numPartitions)
			sliceIdx := p + s*numPartitions
			sliceStart := sliceIdx * sliceStride
			sliceEnd := sliceStart + blocksPerSlice
			if sliceEnd > totalBlocks {
				sliceEnd = totalBlocks
			}

			for i := sliceStart; i < sliceEnd; i++ {
				partitions[p].indices = append(partitions[p].indices, i)
			}
		}
	}

	return partitions
}

func runAllTests(blocks []int64, partitions []VirtualPartition, numWorkers int) {
	strategies := getStrategies()
	allStats := make([]StrategyStats, len(strategies))
	for i, s := range strategies {
		allStats[i].Name = s.name
		allStats[i].Totals = make([]float64, numRuns)
	}
	packedSums := make([]float64, len(strategies))

	for stratIdx, strat := range strategies {
		fmt.Printf("  Running %s...\n", strat.name)

		type result struct {
			runIdx int
			total  uint64
			packed int
		}

		results := make(chan result, numRuns)
		var wg sync.WaitGroup

		jobs := make(chan int, numRuns)
		for w := 0; w < numWorkers; w++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				var localIndices []int
				var localRemaining []int
				var rng *rand.Rand

				for runIdx := range jobs {
					partition := partitions[runIdx]
					partLen := len(partition.indices)

					if localIndices == nil || cap(localIndices) < partLen {
						localIndices = make([]int, partLen)
						localRemaining = make([]int, partLen)
					}
					localIndices = localIndices[:partLen]
					localRemaining = localRemaining[:partLen]

					if rng == nil {
						rng = rand.New(rand.NewSource(int64(runIdx)))
					}

					// Copy partition indices and shuffle
					copy(localIndices, partition.indices)
					rng.Shuffle(len(localIndices), func(i, j int) {
						localIndices[i], localIndices[j] = localIndices[j], localIndices[i]
					})

					var total uint64
					var packed int
					if strat.scorer == nil {
						total, packed = runSequential(blocks, localIndices, strat.padded)
					} else {
						total, packed = runPickN(blocks, localIndices, localRemaining, pickN, strat.scorer, strat.padded, rng)
					}

					results <- result{runIdx, total, packed}
				}
			}()
		}

		go func() {
			for i := 0; i < numRuns; i++ {
				jobs <- i
			}
			close(jobs)
		}()

		go func() {
			wg.Wait()
			close(results)
		}()

		for r := range results {
			allStats[stratIdx].Totals[r.runIdx] = float64(r.total)
			packedSums[stratIdx] += float64(r.packed)
		}
	}

	for i := range allStats {
		allStats[i].Mean, allStats[i].Std = meanStd(allStats[i].Totals)
		allStats[i].MeanPacked = packedSums[i] / float64(numRuns)
	}

	fmt.Printf("\n%-38s %12s %12s %10s %12s %12s\n", "Strategy", "Mean(GiB)", "Std(MiB)", "Util%", "Blocks", "95% CI")
	fmt.Println(strings.Repeat("-", 102))
	for _, s := range allStats {
		meanGiB := s.Mean / (1 << 30)
		stdMiB := s.Std / (1 << 20)
		util := s.Mean / float64(containerSize) * 100
		ci95 := 1.96 * s.Std / math.Sqrt(float64(numRuns)) / (1 << 20)
		fmt.Printf("%-38s %12.3f %12.3f %10.2f %12.0f %9.2f MiB\n",
			s.Name, meanGiB, stdMiB, util, s.MeanPacked, ci95)
	}

	fmt.Println("\nP-values (Welch t-test vs strategy 3a):")
	baseline := allStats[2].Totals
	for i, s := range allStats {
		if i == 2 {
			continue
		}
		p := tTestPValue(baseline, s.Totals)
		sig := ""
		if p < 0.001 {
			sig = "***"
		} else if p < 0.01 {
			sig = "**"
		} else if p < 0.05 {
			sig = "*"
		}
		diff := (allStats[2].Mean - s.Mean) / (1 << 30)
		fmt.Printf("  %-36s: p=%.2e %s (diff: %+.3f GiB)\n", s.Name, p, sig, diff)
	}
}

func runSequential(blocks []int64, indices []int, padded bool) (uint64, int) {
	total := uint64(0)
	offset := uint64(0)
	packed := 0
	for _, idx := range indices {
		blockSize := blocks[idx]
		paddedSize := uint64(PaddedSize(uint64(blockSize)))
		alignedOff := alignedOffset(offset, paddedSize)

		var sizeToStore uint64
		if padded {
			sizeToStore = paddedSize
		} else {
			sizeToStore = uint64(blockSize)
		}

		if alignedOff+sizeToStore <= containerSize {
			total += uint64(blockSize)
			offset = alignedOff + sizeToStore
			packed++
		}
	}
	return total, packed
}

func runPickN(blocks []int64, indices, remaining []int, n int, scorer scoreFunc, padded bool, rng *rand.Rand) (uint64, int) {
	copy(remaining[:len(indices)], indices)
	remainingLen := len(indices)

	total := uint64(0)
	packed := 0
	offset := uint64(0)

	for remainingLen > 0 {
		numCandidates := n
		if numCandidates > remainingLen {
			numCandidates = remainingLen
		}

		for i := 0; i < numCandidates; i++ {
			j := rng.Intn(remainingLen - i)
			remaining[remainingLen-1-i], remaining[j] = remaining[j], remaining[remainingLen-1-i]
		}

		bestIdx := -1
		bestScore := int64(0)
		bestPos := -1

		for i := 0; i < numCandidates; i++ {
			pos := remainingLen - 1 - i
			blockIdx := remaining[pos]
			blockSize := blocks[blockIdx]

			score := scorer(offset, blockSize)
			if score >= 0 && (bestIdx == -1 || score < bestScore) {
				bestIdx = blockIdx
				bestScore = score
				bestPos = pos
			}
		}

		if bestIdx == -1 {
			remainingLen -= numCandidates
			continue
		}

		blockSize := blocks[bestIdx]
		paddedSize := uint64(PaddedSize(uint64(blockSize)))
		alignedOff := alignedOffset(offset, paddedSize)

		total += uint64(blockSize)
		if padded {
			offset = alignedOff + paddedSize
		} else {
			offset = alignedOff + uint64(blockSize)
		}
		packed++

		remaining[bestPos] = remaining[remainingLen-1]
		remainingLen--
	}

	return total, packed
}

func scoreMinAlignmentWaste(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+uint64(blockSize) > containerSize {
		return -1
	}
	return int64(alignedOff - offset)
}

func scoreLargestFirst(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+uint64(blockSize) > containerSize {
		return -1
	}
	return int64(1<<62) - blockSize
}

func scoreMinTotalWaste(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+uint64(blockSize) > containerSize {
		return -1
	}
	return int64(alignedOff-offset) + int64(paddedSize) - blockSize
}

func scoreSmallestPaddedFirst(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+uint64(blockSize) > containerSize {
		return -1
	}
	return int64(paddedSize)
}

func scoreMinAlignmentWastePadded(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+paddedSize > containerSize {
		return -1
	}
	return int64(alignedOff - offset)
}

func scoreLargestFirstPadded(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+paddedSize > containerSize {
		return -1
	}
	return int64(1<<62) - blockSize
}

func scoreMinTotalWastePadded(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+paddedSize > containerSize {
		return -1
	}
	return int64(alignedOff-offset) + int64(paddedSize) - blockSize
}

func scoreSmallestPaddedFirstPadded(offset uint64, blockSize int64) int64 {
	paddedSize := uint64(PaddedSize(uint64(blockSize)))
	alignedOff := alignedOffset(offset, paddedSize)
	if alignedOff+paddedSize > containerSize {
		return -1
	}
	return int64(paddedSize)
}

func meanStd(data []float64) (mean, std float64) {
	n := float64(len(data))
	for _, v := range data {
		mean += v
	}
	mean /= n
	for _, v := range data {
		std += (v - mean) * (v - mean)
	}
	std = math.Sqrt(std / (n - 1))
	return
}

func tTestPValue(a, b []float64) float64 {
	meanA, stdA := meanStd(a)
	meanB, stdB := meanStd(b)
	nA, nB := float64(len(a)), float64(len(b))

	varA := stdA * stdA
	varB := stdB * stdB

	se := math.Sqrt(varA/nA + varB/nB)
	if se == 0 {
		return 1.0
	}
	t := (meanA - meanB) / se

	num := math.Pow(varA/nA+varB/nB, 2)
	denom := math.Pow(varA/nA, 2)/(nA-1) + math.Pow(varB/nB, 2)/(nB-1)
	df := num / denom

	return 2 * tDistCDF(-math.Abs(t), df)
}

func tDistCDF(t, df float64) float64 {
	x := df / (df + t*t)
	return 0.5 * betaIncomplete(df/2, 0.5, x)
}

func betaIncomplete(a, b, x float64) float64 {
	if x == 0 {
		return 0
	}
	if x == 1 {
		return 1
	}
	if x > (a+1)/(a+b+2) {
		return 1 - betaIncomplete(b, a, 1-x)
	}

	const maxIter = 200
	const epsilon = 1e-14

	qab := a + b
	qap := a + 1
	qam := a - 1
	c := 1.0
	d := 1 - qab*x/qap
	if math.Abs(d) < 1e-30 {
		d = 1e-30
	}
	d = 1 / d
	h := d

	for m := 1; m <= maxIter; m++ {
		m2 := 2 * m
		aa := float64(m) * (b - float64(m)) * x / ((qam + float64(m2)) * (a + float64(m2)))
		d = 1 + aa*d
		if math.Abs(d) < 1e-30 {
			d = 1e-30
		}
		c = 1 + aa/c
		if math.Abs(c) < 1e-30 {
			c = 1e-30
		}
		d = 1 / d
		h *= d * c
		aa = -(a + float64(m)) * (qab + float64(m)) * x / ((a + float64(m2)) * (qap + float64(m2)))
		d = 1 + aa*d
		if math.Abs(d) < 1e-30 {
			d = 1e-30
		}
		c = 1 + aa/c
		if math.Abs(c) < 1e-30 {
			c = 1e-30
		}
		d = 1 / d
		del := d * c
		h *= del
		if math.Abs(del-1) < epsilon {
			break
		}
	}

	logBeta, _ := math.Lgamma(a)
	logBetaB, _ := math.Lgamma(b)
	logBetaAB, _ := math.Lgamma(a + b)
	logBeta = logBeta + logBetaB - logBetaAB
	return math.Exp(a*math.Log(x)+b*math.Log(1-x)-logBeta) * h / a
}