// 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"
)

// Paired crossover probe for Factor's dispatch: the serial elimination
// against the shipped crew constants on the same deterministic matrix,
// interleaved iteration by iteration so a drift in the machine's speed
// lands on both styles equally. The reported ns/op per style is the
// paired comparison; the ns/op of the group as a whole is not
// comparable across groups. The spawn floor factorSpawnFloor sits where
// the crew starts winning.

// probeFactorRows builds the deterministic n×n fixture: a diagonally
// dominant integer pattern every style factors identically.
func probeFactorRows(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 probeFactorReset(rows [][]float64, pristine []float64) {
	off := 0
	for _, row := range rows {
		copy(row, pristine[off:off+len(row)])
		off += len(row)
	}
}

// probeFactorSerial is Factor with the dispatch removed: the reference
// every candidate must match byte for byte.
func probeFactorSerial(m [][]float64) {
	n := len(m)
	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]
		}
		factorRows(m[k+1:], k, n, m[k])
	}
}

// probeFactorJob mirrors factorJob for the probe's own dispatch, so the
// probe keeps compiling whatever the production type later gains.
type probeFactorJob struct {
	rows     [][]float64
	pivotRow []float64
	k, n     int
	start    int
	end      int
	wg       *sync.WaitGroup
}

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

// probeFactorCrew mirrors Factor's dispatch at the production constants,
// which it reads directly: the quantum sizes the crew, the spawn floor
// decides whether any crew runs, and the per-row update is the same
// factorRows call the serial reference makes.
func probeFactorCrew(m [][]float64) {
	n := len(m)
	var wg sync.WaitGroup
	var jobs []probeFactorJob
	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]
		}
		pivotRow := m[k]
		rows := m[k+1:]
		work := len(rows) * (n - k)
		w := min(work/factorWorkQuantum+1, engine.WorkersFor(len(rows)))
		w = min(w, factorMaxCrew)
		if work < factorSpawnFloor {
			w = 1
		}
		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([]probeFactorJob, factorMaxCrew)
		}
		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 probeFactorWorker(&jobs[i])
		}
		wg.Wait()
	}
}

// TestFactorSpawnFloorBitIdentical pins the crew mirror against the
// serial reference at the sizes the crossover probe sweeps, whatever
// the dispatch decides: crew size never moves a bit.
func TestFactorSpawnFloorBitIdentical(t *testing.T) {
	for _, n := range []int{2, 3, 16, 64, 256, 320, 384} {
		want, _ := probeFactorRows(n)
		probeFactorSerial(want)
		got, _ := probeFactorRows(n)
		probeFactorCrew(got)
		for i := range n {
			for j := range n {
				if got[i][j] != want[i][j] {
					t.Fatalf("n=%d factor[%d][%d] = %v, want %v", n, i, j, got[i][j], want[i][j])
				}
			}
		}
	}
}

// BenchmarkFactorCrossoverAB interleaves the serial elimination with the
// shipped dispatch iteration by iteration across the sizes around the
// spawn floor.
func BenchmarkFactorCrossoverAB(b *testing.B) {
	for _, n := range []int{256, 320, 384, 448, 512} {
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			rowsS, prS := probeFactorRows(n)
			rowsC, prC := probeFactorRows(n)
			var tSerial, tCrew time.Duration
			for b.Loop() {
				probeFactorReset(rowsS, prS)
				probeFactorReset(rowsC, prC)
				start := time.Now()
				probeFactorSerial(rowsS)
				tSerial += time.Since(start)
				start = time.Now()
				probeFactorCrew(rowsC)
				tCrew += time.Since(start)
			}
			b.ReportMetric(float64(tSerial.Nanoseconds())/float64(b.N), "ns/op-serial")
			b.ReportMetric(float64(tCrew.Nanoseconds())/float64(b.N), "ns/op-crew")
		})
	}
}