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

package signal

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

import (
	"math"
	"sync"
)

// Periodograms. The Lomb-Scargle periodogram answers "at which
// frequency does unevenly sampled data oscillate" without the
// interpolation a resampled FFT would need: each trial frequency gets
// its own least-squares fit of a sine and cosine through the actual
// observation times, with the phase reference τ chosen so the two
// fitted components are exactly orthogonal at that frequency.

// lombParallelMinN is the observation count above which a single
// frequency's two sine/cosine passes are worth a worker's spawn cost.
// Below it the frequency grid walk stays on the calling goroutine.
const lombParallelMinN = 1 << 9

// LombScargle computes the normalised Lomb-Scargle periodogram of the
// observations values taken at times, over the nFreq frequencies
// evenly spaced from minFreq to maxFreq inclusive (a single frequency
// when nFreq is 1) and returns the frequency grid and the power at
// each frequency. The power carries the classical
// normalisation: a pure sinusoid of amplitude A at a frequency on the
// grid peaks near A²·n/(4·var(values)), so the scale is comparable
// across data sets. An empty or two-point time base, a length
// mismatch, an all-equal time base, a non-positive variance, or a
// frequency range that does not satisfy 0 < minFreq ≤ maxFreq is an
// error.
func LombScargle(times, values *core.Array, minFreq, maxFreq float64, nFreq int) (freqs, power *core.Array, err error) {
	const name = "LombScargle"
	if times.NDim() != 1 || values.NDim() != 1 {
		return nil, nil, base.Errf("%s: times and values must be vectors", name)
	}
	if times.Dtype() == core.Complex || values.Dtype() == core.Complex {
		return nil, nil, base.Errf("%s: complex arrays are not supported", name)
	}
	n := values.Len()
	if n < 3 {
		return nil, nil, base.Errf("%s: at least three observations are needed, got %d", name, n)
	}
	if times.Len() != n {
		return nil, nil, base.Errf("%s: times has %d entries for %d values", name, times.Len(), n)
	}
	if nFreq < 1 {
		return nil, nil, base.Errf("%s: nFreq must be at least 1, got %d", name, nFreq)
	}
	// The gate is NaN-rejecting and Inf-rejecting at once: +Inf passes
	// a bare > 0, and Inf endpoints turn every interpolated frequency
	// into NaN with no error.
	if !(minFreq > 0) || math.IsInf(minFreq, 0) || math.IsInf(maxFreq, 0) || maxFreq < minFreq {
		return nil, nil, base.Errf("%s: the frequency range must satisfy 0 < minFreq ≤ maxFreq over finite frequencies, got [%g, %g]",
			name, minFreq, maxFreq)
	}
	t := make([]float64, n)
	x := make([]float64, n)
	mean := 0.0
	// A non-finite time or value would drive the variance NaN, slip
	// past its gate and publish NaN powers with no error, so both
	// arrays are refused up front (the guard SolvePoissonPeriodic
	// applies to its source).
	for i := range n {
		t[i] = times.FloatAt(i)
		if math.IsNaN(t[i]) || math.IsInf(t[i], 0) {
			return nil, nil, base.Errf("%s: times holds the non-finite value %g at %d", name, t[i], i)
		}
		x[i] = values.FloatAt(i)
		if math.IsNaN(x[i]) || math.IsInf(x[i], 0) {
			return nil, nil, base.Errf("%s: values holds the non-finite value %g at %d", name, x[i], i)
		}
		mean += x[i]
	}
	mean /= float64(n)
	// A constant time base carries no phase information: every trial
	// frequency drives the sine fit to 0/0.
	if allEqual(t) {
		return nil, nil, base.Errf("%s: the times must not all be equal", name)
	}
	variance := 0.0
	for i := range n {
		x[i] -= mean
		variance += x[i] * x[i]
	}
	variance /= float64(n - 1)
	if variance <= 0 {
		return nil, nil, base.Errf("%s: the values have zero variance", name)
	}
	tCenter := t[n/2]
	// The offsets from the phase centre feed every trig argument of
	// every frequency: (t[i]−tCenter) is recomputed twice per
	// observation per frequency, so it is evaluated once here and
	// reused. The stored value is the subtraction result itself, so
	// every argument keeps the exact bits it had.
	dt := make([]float64, n)
	for i := range n {
		dt[i] = t[i] - tCenter
	}
	freqsArr := core.New(core.Float, nFreq)
	powerArr := core.New(core.Float, nFreq)
	freqRow := freqsArr.RawFloats()
	powerRow := powerArr.RawFloats()
	// A frequency whose sine or cosine sum vanishes cannot be fitted
	// on this time base. The serial walk reported the lowest such
	// frequency; the split keeps that contract by remembering the
	// smallest offending index and erroring after the join.
	var (
		badMu sync.Mutex
		bad   = -1
	)
	unresolvable := func(f int) {
		badMu.Lock()
		defer badMu.Unlock()
		if bad < 0 || f < bad {
			bad = f
		}
	}
	// fitAt runs the whole per-frequency pipeline: the grid frequency,
	// the orthogonalising phase reference τ and both least-squares
	// fits. Every read is from the shared time and value slices, every
	// write lands in this frequency's own slot of the two outputs, and
	// the per-frequency arithmetic sequence is the serial one
	// unchanged, so the split cannot move an addend.
	fitAt := func(f int) {
		freq := minFreq
		if nFreq > 1 {
			freq = minFreq + (maxFreq-minFreq)*float64(f)/float64(nFreq-1)
		}
		freqRow[f] = freq
		omega := 2 * math.Pi * freq
		// The phase reference τ keeps the sine and cosine fits
		// orthogonal at this frequency. Both passes need the sine and
		// the cosine of the same argument; math.Sincos shares the range
		// reduction between the two and returns exactly the pair
		// math.Sin and math.Cos produce (verified bit-for-bit), so the
		// sums are unchanged while the trig work halves.
		sumSin2, sumCos2 := 0.0, 0.0
		for i := range n {
			arg := omega * dt[i]
			s, c := math.Sincos(2 * arg)
			sumSin2 += s
			sumCos2 += c
		}
		tau := 0.5 * math.Atan2(sumSin2, sumCos2) / omega
		sumCos, sumSin, sumCosSq, sumSinSq := 0.0, 0.0, 0.0, 0.0
		for i := range n {
			arg := omega * (dt[i] - tau)
			c, s := math.Sincos(arg)
			sumCos += x[i] * c
			sumSin += x[i] * s
			sumCosSq += c * c
			sumSinSq += s * s
		}
		if sumSinSq == 0 || sumCosSq == 0 {
			unresolvable(f)
			return
		}
		power := (sumCos*sumCos)/sumCosSq + (sumSin*sumSin)/sumSinSq
		powerRow[f] = power / (2 * variance)
	}
	if n >= lombParallelMinN {
		// The frequencies split across workers: disjoint output slots,
		// per-frequency normalisations computed inside the worker that
		// owns the frequency.
		engine.Parallel(nFreq, func(fs, fe int) {
			for f := fs; f < fe; f++ {
				fitAt(f)
			}
		})
	} else {
		for f := range nFreq {
			fitAt(f)
		}
	}
	if bad >= 0 {
		freq := minFreq
		if nFreq > 1 {
			freq = minFreq + (maxFreq-minFreq)*float64(bad)/float64(nFreq-1)
		}
		return nil, nil, base.Errf("%s: the time base cannot resolve the frequency %g", name, freq)
	}
	return freqsArr, powerArr, nil
}

// allEqual reports whether every slice entry matches the first.
func allEqual(v []float64) bool {
	for _, x := range v[1:] {
		if x != v[0] {
			return false
		}
	}
	return true
}