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

package stats

import (
	"math"
	"sync"

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

// Histogram2D bins the paired samples (x[i], y[i]) into a
// xBins × yBins count matrix over equal-width bins spanning each
// axis's own data range, mirroring Histogram's conventions: bins are
// closed on the left and the last bin absorbs the maximum. The count
// matrix comes back as a (xBins × yBins) array, the edges as the
// bin boundaries on each axis. Real inputs only, and the bin counts
// are bounded by maxHistBins in total, exactly as Histogram's are.
func Histogram2D(x, y *core.Array, xBins, yBins int) (*core.Array, []float64, []float64, error) {
	const name = "Histogram2D"
	if x.Len() != y.Len() {
		return nil, nil, nil, base.Errf("%s: x and y must share their length, got %d and %d", name, x.Len(), y.Len())
	}
	if x.Len() == 0 {
		return nil, nil, nil, base.Errf("%s: empty samples have no histogram", name)
	}
	if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
		return nil, nil, nil, base.Errf("%s: complex inputs are not supported", name)
	}
	if xBins < 1 || yBins < 1 {
		return nil, nil, nil, base.Errf("%s: needs at least one bin per axis, got %d and %d", name, xBins, yBins)
	}
	// The feasibility check runs before any allocation: the count matrix
	// costs xBins*yBins cells, and refused up front it can neither
	// overflow the product nor ask the allocator for an unbounded
	// buffer.
	if xBins > maxHistBins || yBins > maxHistBins || xBins > maxHistBins/yBins {
		return nil, nil, nil, base.Errf("%s: the %d x %d bin request exceeds the %d-bin limit",
			name, xBins, yBins, maxHistBins)
	}
	xEdges, err := edgesFor(x, xBins, name, "x")
	if err != nil {
		return nil, nil, nil, err
	}
	yEdges, err := edgesFor(y, yBins, name, "y")
	if err != nil {
		return nil, nil, nil, err
	}
	counts := core.New(core.Int, xBins, yBins)
	ints := counts.RawInts()
	// The paired samples are read once: a dense float64 payload is walked
	// in place and any other layout is widened into a private slice, so
	// the sweep below reads no accessor and every value is the one
	// FloatAt returns.
	xVals, yVals := histogramAxis(x), histogramAxis(y)
	// Counting a cell is exact integer arithmetic and every sample
	// carries exactly one increment, so the sweep splits over disjoint
	// slices of the sample and the private counters merge in any order
	// into the totals a single pass produces.
	bins := xBins * yBins
	minPerWorker := max(hist2DSerialSamples, hist2DBinsPerWorker*bins)
	// The first non-finite pair is the one the serial sweep would name:
	// the chunks report their own first, and the lowest index wins.
	bad := -1
	var mu sync.Mutex
	note := func(i int) {
		if i < 0 {
			return
		}
		mu.Lock()
		if bad < 0 || i < bad {
			bad = i
		}
		mu.Unlock()
	}
	engine.ParallelMin(x.Len(), minPerWorker, func(start, end int) {
		if start == 0 && end == x.Len() {
			// The whole sample runs inline, the worker policy having
			// found the split uneconomical: count straight into the
			// result.
			note(histogram2DCount(xVals, yVals, 0, xEdges, yEdges, xBins, yBins, ints))
			return
		}
		local := make([]int64, bins)
		first := histogram2DCount(xVals[start:end], yVals[start:end], start, xEdges, yEdges, xBins, yBins, local)
		mu.Lock()
		for i, c := range local {
			ints[i] += c
		}
		mu.Unlock()
		note(first)
	})
	if bad >= 0 {
		return nil, nil, nil, base.Errf("%s: sample %d (%g, %g) is not finite", name, bad, xVals[bad], yVals[bad])
	}
	return counts, xEdges, yEdges, nil
}

// hist2DBinsPerWorker is the paired-sample count one counting worker
// must carry per cell before the binning splits across goroutines, and
// hist2DSerialSamples the count it must carry whatever the cell count:
// a worker counts into a private matrix of xBins·yBins cells, so the
// split only pays where the samples drained into it outweigh the cells
// it costs.
const (
	hist2DSerialSamples = 1 << 12
	hist2DBinsPerWorker = 8
)

// histogramAxis returns the axis's values as a plain slice, walking a
// dense float64 payload in place. Element i sits at payload index i,
// and a rebased view's payload may run past its own count, so the view
// is cut to the visible elements.
func histogramAxis(a *core.Array) []float64 {
	if fs := rawFloats(a); fs != nil {
		return fs[:a.Len()]
	}
	vals := make([]float64, a.Len())
	for i := range vals {
		vals[i] = a.FloatAt(i)
	}
	return vals
}

// histogram2DCount adds one slice of the paired sample to counts and
// returns the index of the first non-finite pair it met, or -1. The bin
// arithmetic is binOf's own, so a split of the sample changes no count.
func histogram2DCount(xv, yv []float64, base int, xEdges, yEdges []float64, xBins, yBins int, counts []int64) int {
	xBase, xStep := xEdges[0], xEdges[1]-xEdges[0]
	yBase, yStep := yEdges[0], yEdges[1]-yEdges[0]
	for i := range xv {
		xi := binOf(xv[i], xBase, xStep, xBins)
		yi := binOf(yv[i], yBase, yStep, yBins)
		if xi < 0 || yi < 0 {
			return base + i
		}
		counts[xi*yBins+yi]++
	}
	return -1
}

// edgesFor builds one axis's bin edges over the data's own range.
func edgesFor(a *core.Array, bins int, name, axis string) ([]float64, error) {
	lo, hi := math.Inf(1), math.Inf(-1)
	// The finiteness scan and the range share one pass over a dense
	// float64 payload, the values being the ones FloatAt returns; every
	// other layout keeps the accessor scan and the range pass after it.
	// The payload walk is bounded by Len: a rebased view's payload may
	// run past its own count.
	if fs := rawFloats(a); fs != nil {
		for i, v := range fs[:a.Len()] {
			if math.IsNaN(v) || math.IsInf(v, 0) {
				return nil, base.Errf("%s: %s sample %d is not finite (%g)", name, axis, i, v)
			}
			if v < lo {
				lo = v
			}
			if v > hi {
				hi = v
			}
		}
	} else {
		for i := range a.Len() {
			if v := a.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
				return nil, base.Errf("%s: %s sample %d is not finite (%g)", name, axis, i, v)
			}
		}
		for i := range a.Len() {
			v := a.FloatAt(i)
			if v < lo {
				lo = v
			}
			if v > hi {
				hi = v
			}
		}
	}
	if lo == hi {
		lo -= 0.5
		hi += 0.5
	}
	width := (hi - lo) / float64(bins)
	if math.IsInf(width, 0) || math.IsNaN(width) {
		// Both float extremes in one axis span more than the float64
		// range: no finite edges exist, and the bin arithmetic would
		// clamp every sample into bin 0 over ±Inf edges.
		return nil, base.Errf("%s: the %s samples span more than the float64 range (%g to %g)", name, axis, lo, hi)
	}
	edges := make([]float64, bins+1)
	for i := range bins + 1 {
		edges[i] = lo + float64(i)*width
	}
	return edges, nil
}

// binOf locates one sample's bin from the axis's base edge and the step
// between its first two edges, closing the last bin on the right; a
// non-finite sample returns −1.
func binOf(v, base, step float64, bins int) int {
	if math.IsNaN(v) || math.IsInf(v, 0) {
		return -1
	}
	i := int((v - base) / step)
	if i < 0 {
		return 0 // cannot happen over the data's own range, kept guarded
	}
	if i >= bins {
		i = bins - 1
	}
	return i
}