// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package base

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"sourcedock.dev/petrbalvin/tensor/internal/engine"
)

// The dispatch of Factor is the one place where a crew size, a chunk
// boundary and a spawn style can silently change the arithmetic, so the
// candidates live here beside a bit-identity test and the measurement
// that picks the constants. Every candidate must reproduce factorSerial
// byte for byte, and the production Factor must reproduce the
// parameterised mirror at the shipped constants.

func benchFactorRows(n int) (rows [][]float64, pristine []float64) {
	flat := make([]float64, n*n)
	for i := range n {
		for j := range n {
			flat[i*n+j] = float64((i*7+j*13)%11) - 5
		}
		flat[i*n+i] += float64(n)
	}
	rows = make([][]float64, n)
	for i := range n {
		rows[i] = flat[i*n : (i+1)*n]
	}
	pristine = make([]float64, len(flat))
	copy(pristine, flat)
	return rows, pristine
}

func resetFactorRows(rows [][]float64, pristine []float64) {
	off := 0
	for _, row := range rows {
		copy(row, pristine[off:off+len(row)])
		off += len(row)
	}
}

// factorSerial is Factor with the dispatch removed: the reference every
// candidate must match byte for byte.
func factorSerial(m [][]float64) ([]int, int) {
	n := len(m)
	perm := make([]int, n)
	for i := range perm {
		perm[i] = i
	}
	parity := 1
	for k := range n {
		pivot := k
		for i := k + 1; i < n; i++ {
			if absOf(m[i][k]) > absOf(m[pivot][k]) {
				pivot = i
			}
		}
		if m[pivot][k] == 0 {
			continue
		}
		if pivot != k {
			m[pivot], m[k] = m[k], m[pivot]
			perm[pivot], perm[k] = perm[k], perm[pivot]
			parity = -parity
		}
		factorRows(m[k+1:], k, n, m[k])
	}
	return perm, parity
}

// factorRetired is the dispatch Factor used before the crew moved onto
// preallocated slots: one goroutine per chunk through
// sync.WaitGroup.Go, so every chunk allocates a closure. It is the
// reference the current dispatch is measured against.
func factorRetired(m [][]float64, quantum int) ([]int, int) {
	n := len(m)
	perm := make([]int, n)
	for i := range perm {
		perm[i] = i
	}
	parity := 1
	for k := range n {
		pivot := k
		for i := k + 1; i < n; i++ {
			if absOf(m[i][k]) > absOf(m[pivot][k]) {
				pivot = i
			}
		}
		if m[pivot][k] == 0 {
			continue
		}
		if pivot != k {
			m[pivot], m[k] = m[k], m[pivot]
			perm[pivot], perm[k] = perm[k], perm[pivot]
			parity = -parity
		}
		pivotRow := m[k]
		rows := m[k+1:]
		work := len(rows) * (n - k)
		w := min(work/quantum+1, engine.WorkersFor(len(rows)))
		if w < 2 {
			factorRows(rows, k, n, pivotRow)
			continue
		}
		chunk := (len(rows) + w - 1) / w
		if chunk < factorMinRows {
			w = max(len(rows)/factorMinRows, 1)
			chunk = (len(rows) + w - 1) / w
		}
		var wg sync.WaitGroup
		for start := 0; start < len(rows); start += chunk {
			end := min(start+chunk, len(rows))
			wg.Go(func() {
				factorRows(rows[start:end], k, n, pivotRow)
			})
		}
		wg.Wait()
	}
	return perm, parity
}

// factorJobSlot is one crew member's preallocated work slot, mirroring
// factorJob for the float64 candidates.
type factorJobSlot struct {
	rows     [][]float64
	pivotRow []float64
	k, n     int
	start    int
	end      int
	wg       *sync.WaitGroup
}

func factorSlotWorker(job *factorJobSlot) {
	factorRows(job.rows[job.start:job.end], job.k, job.n, job.pivotRow)
	job.wg.Done()
}

// factorCrew is the shipped dispatch with both tuning constants
// exposed.
func factorCrew(m [][]float64, quantum, maxCrew int) ([]int, int) {
	n := len(m)
	perm := make([]int, n)
	for i := range perm {
		perm[i] = i
	}
	parity := 1
	var wg sync.WaitGroup
	var jobs []factorJobSlot
	for k := range n {
		pivot := k
		for i := k + 1; i < n; i++ {
			if absOf(m[i][k]) > absOf(m[pivot][k]) {
				pivot = i
			}
		}
		if m[pivot][k] == 0 {
			continue
		}
		if pivot != k {
			m[pivot], m[k] = m[k], m[pivot]
			perm[pivot], perm[k] = perm[k], perm[pivot]
			parity = -parity
		}
		pivotRow := m[k]
		rows := m[k+1:]
		work := len(rows) * (n - k)
		w := min(work/quantum+1, engine.WorkersFor(len(rows)))
		w = min(w, maxCrew)
		if w < 2 {
			factorRows(rows, k, n, pivotRow)
			continue
		}
		chunk := (len(rows) + w - 1) / w
		if chunk < factorMinRows {
			w = max(len(rows)/factorMinRows, 1)
			chunk = (len(rows) + w - 1) / w
			if w < 2 {
				factorRows(rows, k, n, pivotRow)
				continue
			}
		}
		if jobs == nil {
			jobs = make([]factorJobSlot, maxCrew)
		}
		spawned := 0
		for start := 0; start < len(rows); start += chunk {
			job := &jobs[spawned]
			job.rows, job.pivotRow, job.k, job.n = rows, pivotRow, k, n
			job.start, job.end, job.wg = start, min(start+chunk, len(rows)), &wg
			spawned++
		}
		wg.Add(spawned)
		for i := range spawned {
			go factorSlotWorker(&jobs[i])
		}
		wg.Wait()
	}
	return perm, parity
}

// TestFactorDispatchBitIdentical proves every crew size, chunk boundary
// and spawn style leaves the factor and the permutation byte for byte
// as the serial reference, and that the shipped Factor matches the
// parameterised mirror the tuning benchmark measures.
func TestFactorDispatchBitIdentical(t *testing.T) {
	for _, n := range []int{1, 2, 5, 16, 64, 129, 260} {
		base, _ := benchFactorRows(n)
		want := make([][]float64, n)
		for i := range n {
			want[i] = append([]float64(nil), base[i]...)
		}
		wp, wpar := factorSerial(want)
		check := func(name string, got [][]float64, gp []int, gpar int) {
			t.Helper()
			if len(gp) != len(wp) {
				t.Fatalf("%s: permutation length %d, want %d", name, len(gp), len(wp))
			}
			for i := range gp {
				if gp[i] != wp[i] {
					t.Fatalf("%s: permutation[%d] = %d, want %d", name, i, gp[i], wp[i])
				}
			}
			for i := range got {
				for j := range got[i] {
					if got[i][j] != want[i][j] {
						t.Fatalf("%s: factor[%d][%d] = %v, want %v", name, i, j, got[i][j], want[i][j])
					}
				}
			}
			if gpar != wpar {
				t.Fatalf("%s: parity %d, want %d", name, gpar, wpar)
			}
		}
		for _, q := range []int{1024, 8192, 65536} {
			for _, w0 := range []int{1, 2, 4, 8, 16, 32} {
				rows, _ := benchFactorRows(n)
				gp, gpar := factorCrew(rows, q, w0)
				check(fmt.Sprintf("crew q=%d w=%d n=%d", q, w0, n), rows, gp, gpar)
			}
			rows, _ := benchFactorRows(n)
			gp, gpar := factorRetired(rows, q)
			check(fmt.Sprintf("retired q=%d n=%d", q, n), rows, gp, gpar)
		}
		rows, _ := benchFactorRows(n)
		gp, gpar := Factor(rows)
		check(fmt.Sprintf("Factor n=%d", n), rows, gp, gpar)
		mirror, _ := benchFactorRows(n)
		mirrorGP, mirrorPar := factorCrew(mirror, factorWorkQuantum, factorMaxCrew)
		check(fmt.Sprintf("mirror n=%d", n), mirror, mirrorGP, mirrorPar)
	}
}

// BenchmarkForkJoinFloor measures an empty fork-join: the spawn, wake
// and join of w goroutines that touch nothing. It is the lower bound on
// what one dispatch costs, and the reason a pivot whose update is
// smaller than this cannot be worth splitting.
func BenchmarkForkJoinFloor(b *testing.B) {
	for _, w := range []int{1, 2, 4, 8, 16, 32} {
		b.Run(fmt.Sprintf("w=%d", w), func(b *testing.B) {
			for b.Loop() {
				var wg sync.WaitGroup
				for range w {
					wg.Go(func() {})
				}
				wg.Wait()
			}
		})
	}
}

// BenchmarkFactorStylesAB interleaves the dispatch styles iteration by
// iteration on separate matrices, so a drift in the machine's speed
// lands on every style equally: the reported ns/op per style is the
// paired comparison, and the ns/op of the group as a whole is not
// comparable across groups.
func BenchmarkFactorStylesAB(b *testing.B) {
	styles := []struct {
		name string
		run  func(m [][]float64)
	}{
		{"serial", func(m [][]float64) { factorSerial(m) }},
		{"retired/q=16384", func(m [][]float64) { factorRetired(m, 16384) }},
		{"slots/q=8192/cap=16", func(m [][]float64) { factorCrew(m, 8192, 16) }},
		{"slots/q=16384/cap=16", func(m [][]float64) { factorCrew(m, 16384, 16) }},
		{"slots/q=24576/cap=16", func(m [][]float64) { factorCrew(m, 24576, 16) }},
		{"slots/q=32768/cap=16", func(m [][]float64) { factorCrew(m, 32768, 16) }},
	}
	for _, n := range []int{256, 512} {
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			rows := make([][][]float64, len(styles))
			pristine := make([][]float64, len(styles))
			for i := range styles {
				rows[i], pristine[i] = benchFactorRows(n)
			}
			elapsed := make([]time.Duration, len(styles))
			for b.Loop() {
				for i, s := range styles {
					resetFactorRows(rows[i], pristine[i])
					start := time.Now()
					s.run(rows[i])
					elapsed[i] += time.Since(start)
				}
			}
			for i, s := range styles {
				b.ReportMetric(float64(elapsed[i].Nanoseconds())/float64(b.N), "ns/op-"+s.name)
			}
		})
	}
}

// transposeFlatPlain is the untiled copy the tiled TransposeFlat
// replaces: the read runs along the rows, the write strides by the row
// length.
func transposeFlatPlain(a []float64, m, n int) []float64 {
	out := make([]float64, m*n)
	for i := range m {
		for j := range n {
			out[j*m+i] = a[i*n+j]
		}
	}
	return out
}

// TestTransposeFlatTiledBitIdentical pins the tiled copy against the
// untiled one: a transpose is a permutation, so it must match value for
// value, and the tile boundary must not lose a corner of a ragged
// matrix.
func TestTransposeFlatTiledBitIdentical(t *testing.T) {
	for _, dim := range [][2]int{{1, 1}, {1, 7}, {7, 1}, {16, 16}, {17, 16}, {16, 17}, {33, 5}, {5, 33}, {64, 64}} {
		m, n := dim[0], dim[1]
		a := make([]float64, m*n)
		for i := range a {
			a[i] = float64(i%13) - 6
		}
		got := TransposeFlat(a, m, n)
		want := transposeFlatPlain(a, m, n)
		if len(got) != len(want) {
			t.Fatalf("%dx%d: length %d, want %d", m, n, len(got), len(want))
		}
		for i := range got {
			if got[i] != want[i] {
				t.Fatalf("%dx%d: entry %d = %v, want %v", m, n, i, got[i], want[i])
			}
		}
	}
}

// BenchmarkTransposeFlatAB interleaves the tiled copy with the untiled
// one on separate buffers, so a drift in the machine's speed lands on
// both: the reported ns/op per style is the paired comparison.
func BenchmarkTransposeFlatAB(b *testing.B) {
	for _, n := range []int{64, 256} {
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			flat := make([]float64, n*n)
			for i := range flat {
				flat[i] = float64(i%17) - 8
			}
			var tiled, plain time.Duration
			for b.Loop() {
				start := time.Now()
				out := TransposeFlat(flat, n, n)
				tiled += time.Since(start)
				if len(out) != len(flat) {
					b.Fatal("transpose lost entries")
				}
				start = time.Now()
				out = transposeFlatPlain(flat, n, n)
				plain += time.Since(start)
				if len(out) != len(flat) {
					b.Fatal("transpose lost entries")
				}
			}
			b.ReportMetric(float64(tiled.Nanoseconds())/float64(b.N), "ns/op-tiled")
			b.ReportMetric(float64(plain.Nanoseconds())/float64(b.N), "ns/op-plain")
		})
	}
}

// BenchmarkTransposeFlat measures the tiled transpose alone.
func BenchmarkTransposeFlat(b *testing.B) {
	for _, n := range []int{64, 256} {
		flat := make([]float64, n*n)
		for i := range flat {
			flat[i] = float64(i%17) - 8
		}
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				out := TransposeFlat(flat, n, n)
				if len(out) != len(flat) {
					b.Fatal("transpose lost entries")
				}
			}
		})
	}
}

// BenchmarkSolveSystemColumns measures the column loop of SolveSystem
// with one column and with several: the several-column path permutes
// each column through a bitmap the crew reuses, which is what the
// one-column path does too.
func BenchmarkSolveSystemColumns(b *testing.B) {
	const n = 256
	for _, cols := range []int{1, 8} {
		b.Run(fmt.Sprintf("cols=%d", cols), func(b *testing.B) {
			rhs := make([][]float64, cols)
			for c := range rhs {
				rhs[c] = make([]float64, n)
				for i := range rhs[c] {
					rhs[c][i] = float64((i+c)%9) - 4
				}
			}
			b.ReportAllocs()
			for b.Loop() {
				m, _ := benchFactorRows(n)
				work := make([][]float64, cols)
				for c := range work {
					work[c] = append([]float64(nil), rhs[c]...)
				}
				if _, err := SolveSystem("Solve", m, work); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}