// bench_soa_structs.v - V2 benchmark for @[soa] structs
// Build and run with:
//   ./cmd/v2/v2 -prod -backend cleanc bench/bench_soa_structs.v -o bench/bench_soa_structs
//   ./bench/bench_soa_structs
module main

import time

const particle_count = 500_000
const build_rounds = 5
const sum_rounds = 48
const hot_fields_rounds = 40
const full_struct_rounds = 20
const integrate_rounds = 20

@[soa]
struct Particle {
mut:
	x           f32
	y           f32
	z           f32
	vx          f32
	vy          f32
	vz          f32
	mass        f32
	temperature f32
	ax          f32
	ay          f32
	az          f32
	life        f32
	drag        f32
	color_r     f32
	color_g     f32
	color_b     f32
}

@[typedef]
struct C.Particle_SOA {
	len         int
	cap         int
	x           &f32
	y           &f32
	z           &f32
	vx          &f32
	vy          &f32
	vz          &f32
	mass        &f32
	temperature &f32
	ax          &f32
	ay          &f32
	az          &f32
	life        &f32
	drag        &f32
	color_r     &f32
	color_g     &f32
	color_b     &f32
}

fn C.Particle_SOA_new(int, int) C.Particle_SOA
fn C.Particle_SOA_push(&C.Particle_SOA, Particle)
fn C.Particle_SOA_get(C.Particle_SOA, int) Particle
fn C.Particle_SOA_free(&C.Particle_SOA)

struct BenchResult {
	elapsed_ms i64
	checksum   f64
}

fn make_particle(i int) Particle {
	base := f32(i % 1024)
	return Particle{
		x:           base * 0.25
		y:           base * 0.5
		z:           base * 0.75
		vx:          f32(0.001) * f32((i % 13) + 1)
		vy:          f32(0.0015) * f32((i % 17) + 1)
		vz:          f32(0.0005) * f32((i % 19) + 1)
		mass:        f32(1.0) + f32(i % 5) * f32(0.1)
		temperature: f32(20.0) + f32(i % 7)
		ax:          f32(0.0001) * f32((i % 11) + 1)
		ay:          f32(0.0002) * f32((i % 9) + 1)
		az:          f32(0.0003) * f32((i % 7) + 1)
		life:        f32(100.0) - f32(i % 80)
		drag:        f32(0.98)
		color_r:     f32(i % 255) / f32(255.0)
		color_g:     f32((i * 3) % 255) / f32(255.0)
		color_b:     f32((i * 7) % 255) / f32(255.0)
	}
}

fn build_aos() []Particle {
	mut particles := []Particle{cap: particle_count}
	for i in 0 .. particle_count {
		particles << make_particle(i)
	}
	return particles
}

fn build_soa() C.Particle_SOA {
	mut soa := C.Particle_SOA_new(0, particle_count)
	for i in 0 .. particle_count {
		C.Particle_SOA_push(&soa, make_particle(i))
	}
	return soa
}

fn build_soa_indexed() C.Particle_SOA {
	mut soa := C.Particle_SOA_new(particle_count, particle_count)
	unsafe {
		for i in 0 .. particle_count {
			p := make_particle(i)
			soa.x[i] = p.x
			soa.y[i] = p.y
			soa.z[i] = p.z
			soa.vx[i] = p.vx
			soa.vy[i] = p.vy
			soa.vz[i] = p.vz
			soa.mass[i] = p.mass
			soa.temperature[i] = p.temperature
			soa.ax[i] = p.ax
			soa.ay[i] = p.ay
			soa.az[i] = p.az
			soa.life[i] = p.life
			soa.drag[i] = p.drag
			soa.color_r[i] = p.color_r
			soa.color_g[i] = p.color_g
			soa.color_b[i] = p.color_b
		}
	}
	return soa
}

fn bench_build_aos() BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. build_rounds {
		particles := build_aos()
		sum += particles.len
		sum += particles[0].x
		sum += particles[particles.len - 1].life
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_build_soa() BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. build_rounds {
		mut soa := build_soa()
		sum += soa.len
		unsafe {
			sum += soa.x[0]
			sum += soa.life[soa.len - 1]
		}
		C.Particle_SOA_free(&soa)
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_build_soa_indexed() BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. build_rounds {
		mut soa := build_soa_indexed()
		sum += soa.len
		unsafe {
			sum += soa.x[0]
			sum += soa.life[soa.len - 1]
		}
		C.Particle_SOA_free(&soa)
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_sum_x_aos(particles []Particle) BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. sum_rounds {
		for i in 0 .. particles.len {
			sum += particles[i].x
		}
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_sum_x_soa(soa C.Particle_SOA) BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. sum_rounds {
		unsafe {
			for i in 0 .. soa.len {
				sum += soa.x[i]
			}
		}
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_sum_hot_fields_aos(particles []Particle) BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. hot_fields_rounds {
		for i in 0 .. particles.len {
			sum += particles[i].x + particles[i].y + particles[i].z + particles[i].life
		}
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_sum_hot_fields_soa(soa C.Particle_SOA) BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. hot_fields_rounds {
		unsafe {
			for i in 0 .. soa.len {
				sum += soa.x[i] + soa.y[i] + soa.z[i] + soa.life[i]
			}
		}
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_sum_all_fields_aos(particles []Particle) BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. full_struct_rounds {
		for i in 0 .. particles.len {
			sum += particles[i].x + particles[i].y + particles[i].z + particles[i].vx +
				particles[i].vy + particles[i].vz + particles[i].mass + particles[i].temperature +
				particles[i].ax + particles[i].ay + particles[i].az + particles[i].life +
				particles[i].drag + particles[i].color_r + particles[i].color_g +
				particles[i].color_b
		}
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_sum_all_fields_soa(soa C.Particle_SOA) BenchResult {
	sw := time.new_stopwatch()
	mut sum := f64(0.0)
	for _ in 0 .. full_struct_rounds {
		for i in 0 .. soa.len {
			p := C.Particle_SOA_get(soa, i)
			sum += p.x + p.y + p.z + p.vx + p.vy + p.vz + p.mass + p.temperature + p.ax + p.ay +
				p.az + p.life + p.drag + p.color_r + p.color_g + p.color_b
		}
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_integrate_aos(mut particles []Particle) BenchResult {
	sw := time.new_stopwatch()
	for _ in 0 .. integrate_rounds {
		for i in 0 .. particles.len {
			particles[i].vx += particles[i].ax
			particles[i].vy += particles[i].ay
			particles[i].vz += particles[i].az
			particles[i].x += particles[i].vx * particles[i].mass
			particles[i].y += particles[i].vy * particles[i].mass
			particles[i].z += particles[i].vz * particles[i].mass
			particles[i].life -= particles[i].drag
		}
	}
	mut sum := f64(0.0)
	for i in 0 .. particles.len {
		sum += particles[i].x + particles[i].y + particles[i].z + particles[i].life
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn bench_integrate_soa(mut soa C.Particle_SOA) BenchResult {
	sw := time.new_stopwatch()
	unsafe {
		for _ in 0 .. integrate_rounds {
			for i in 0 .. soa.len {
				soa.vx[i] += soa.ax[i]
				soa.vy[i] += soa.ay[i]
				soa.vz[i] += soa.az[i]
				soa.x[i] += soa.vx[i] * soa.mass[i]
				soa.y[i] += soa.vy[i] * soa.mass[i]
				soa.z[i] += soa.vz[i] * soa.mass[i]
				soa.life[i] -= soa.drag[i]
			}
		}
	}
	mut sum := f64(0.0)
	unsafe {
		for i in 0 .. soa.len {
			sum += soa.x[i] + soa.y[i] + soa.z[i] + soa.life[i]
		}
	}
	return BenchResult{
		elapsed_ms: sw.elapsed().milliseconds()
		checksum:   sum
	}
}

fn print_result(name string, aos BenchResult, soa BenchResult) {
	println(name)
	println('  aos: ${aos.elapsed_ms} ms, checksum=${aos.checksum}')
	println('  soa: ${soa.elapsed_ms} ms, checksum=${soa.checksum}')
	if soa.elapsed_ms > 0 {
		println('  speedup: ${f64(aos.elapsed_ms) / f64(soa.elapsed_ms)}x')
	}
}

fn print_build_results(aos BenchResult, soa_push BenchResult, soa_indexed BenchResult) {
	println('build particles')
	println('  aos: ${aos.elapsed_ms} ms, checksum=${aos.checksum}')
	println('  soa push: ${soa_push.elapsed_ms} ms, checksum=${soa_push.checksum}')
	println('  soa indexed: ${soa_indexed.elapsed_ms} ms, checksum=${soa_indexed.checksum}')
	if soa_indexed.elapsed_ms > 0 {
		println('  push vs indexed: ${f64(soa_push.elapsed_ms) / f64(soa_indexed.elapsed_ms)}x')
		println('  indexed speedup over aos: ${f64(aos.elapsed_ms) / f64(soa_indexed.elapsed_ms)}x')
	}
}

fn main() {
	println('soa struct bench (v2 cleanc)')
	println('particle_count=${particle_count}, build_rounds=${build_rounds}, sum_rounds=${sum_rounds}, hot_fields_rounds=${hot_fields_rounds}, full_struct_rounds=${full_struct_rounds}, integrate_rounds=${integrate_rounds}')

	result_build_aos := bench_build_aos()
	result_build_soa_push := bench_build_soa()
	result_build_soa_indexed := bench_build_soa_indexed()
	print_build_results(result_build_aos, result_build_soa_push, result_build_soa_indexed)

	aos_scan := build_aos()
	soa_scan := build_soa()
	result_aos_scan := bench_sum_x_aos(aos_scan)
	result_soa_scan := bench_sum_x_soa(soa_scan)
	print_result('sum x only', result_aos_scan, result_soa_scan)

	result_aos_hot_fields := bench_sum_hot_fields_aos(aos_scan)
	result_soa_hot_fields := bench_sum_hot_fields_soa(soa_scan)
	print_result('sum x/y/z/life', result_aos_hot_fields, result_soa_hot_fields)

	result_aos_full_struct := bench_sum_all_fields_aos(aos_scan)
	result_soa_full_struct := bench_sum_all_fields_soa(soa_scan)
	print_result('materialize full struct and sum all fields', result_aos_full_struct,
		result_soa_full_struct)
	C.Particle_SOA_free(&soa_scan)

	mut aos_integrate := build_aos()
	mut soa_integrate := build_soa()
	result_aos_integrate := bench_integrate_aos(mut aos_integrate)
	result_soa_integrate := bench_integrate_soa(mut soa_integrate)
	print_result('integrate position, velocity, and life', result_aos_integrate,
		result_soa_integrate)
	C.Particle_SOA_free(&soa_integrate)
}