package dsp

import (
	"errors"
	"fmt"
	"math"
)

// This file holds ArrayNode ops: those that operate natively on waveform
// (float64 array) Samples — reductions (array→scalar), producers (array→array),
// and element access. Each also implements the legacy scalar Node interface
// (treating a scalar as a single-element array) so it remains usable from the
// scalar eval path.

// reductionProcess adapts a scalar Process call to a reduction ArrayNode.
func reductionProcess(n ArrayNode, in []float64, st map[string]any) (float64, error) {
	s, err := n.ProcessSample(scalarInputs(in), st)
	if err != nil {
		return 0, err
	}
	return s.F, nil
}

// ── IndexNode ───────────────────────────────────────────────────────────────

// IndexNode extracts element I of an array input (array→scalar).
type IndexNode struct{ I int }

func (n *IndexNode) Type() string { return "index" }
func (n *IndexNode) Process(in []float64, st map[string]any) (float64, error) {
	return reductionProcess(n, in, st)
}
func (n *IndexNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("index: no inputs")
	}
	arr := in[0].AsArray()
	if n.I < 0 || n.I >= len(arr) {
		return Sample{}, fmt.Errorf("index: %d out of range [0,%d)", n.I, len(arr))
	}
	return Scalar(arr[n.I]), nil
}

// ── SliceNode ───────────────────────────────────────────────────────────────

// SliceNode returns a sub-range [Start,End) of an array input (array→array),
// clamped to the array bounds. End <= 0 means "to the end".
type SliceNode struct{ Start, End int }

func (n *SliceNode) Type() string { return "slice" }
func (n *SliceNode) Process(in []float64, st map[string]any) (float64, error) {
	s, err := n.ProcessSample(scalarInputs(in), st)
	if err != nil {
		return 0, err
	}
	if len(s.Arr) == 0 {
		return 0, nil
	}
	return s.Arr[0], nil
}
func (n *SliceNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("slice: no inputs")
	}
	arr := in[0].AsArray()
	start := n.Start
	if start < 0 {
		start = 0
	}
	if start > len(arr) {
		start = len(arr)
	}
	end := n.End
	if end <= 0 || end > len(arr) {
		end = len(arr)
	}
	if end < start {
		end = start
	}
	out := make([]float64, end-start)
	copy(out, arr[start:end])
	return Array(out), nil
}

// ── reductions ────────────────────────────────────────────────────────────────

// SumNode sums an array input (array→scalar).
type SumNode struct{}

func (n *SumNode) Type() string { return "sum" }
func (n *SumNode) Process(in []float64, st map[string]any) (float64, error) {
	return reductionProcess(n, in, st)
}
func (n *SumNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("sum: no inputs")
	}
	var s float64
	for _, v := range in[0].AsArray() {
		s += v
	}
	return Scalar(s), nil
}

// MeanNode averages an array input (array→scalar).
type MeanNode struct{}

func (n *MeanNode) Type() string { return "mean" }
func (n *MeanNode) Process(in []float64, st map[string]any) (float64, error) {
	return reductionProcess(n, in, st)
}
func (n *MeanNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("mean: no inputs")
	}
	arr := in[0].AsArray()
	if len(arr) == 0 {
		return Scalar(0), nil
	}
	var s float64
	for _, v := range arr {
		s += v
	}
	return Scalar(s / float64(len(arr))), nil
}

// MinNode returns the minimum element of an array input (array→scalar).
type MinNode struct{}

func (n *MinNode) Type() string { return "min" }
func (n *MinNode) Process(in []float64, st map[string]any) (float64, error) {
	return reductionProcess(n, in, st)
}
func (n *MinNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("min: no inputs")
	}
	arr := in[0].AsArray()
	if len(arr) == 0 {
		return Scalar(0), nil
	}
	m := arr[0]
	for _, v := range arr[1:] {
		if v < m {
			m = v
		}
	}
	return Scalar(m), nil
}

// MaxNode returns the maximum element of an array input (array→scalar).
type MaxNode struct{}

func (n *MaxNode) Type() string { return "max" }
func (n *MaxNode) Process(in []float64, st map[string]any) (float64, error) {
	return reductionProcess(n, in, st)
}
func (n *MaxNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("max: no inputs")
	}
	arr := in[0].AsArray()
	if len(arr) == 0 {
		return Scalar(0), nil
	}
	m := arr[0]
	for _, v := range arr[1:] {
		if v > m {
			m = v
		}
	}
	return Scalar(m), nil
}

// LengthNode returns the element count of an array input (array→scalar).
type LengthNode struct{}

func (n *LengthNode) Type() string { return "length" }
func (n *LengthNode) Process(in []float64, st map[string]any) (float64, error) {
	return reductionProcess(n, in, st)
}
func (n *LengthNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("length: no inputs")
	}
	return Scalar(float64(len(in[0].AsArray()))), nil
}

// ── FFTNode ────────────────────────────────────────────────────────────────

// FFTNode computes the magnitude spectrum of an array input (array→array). The
// input is zero-padded to the next power of two; the output has that length and
// holds |X[k]| for each frequency bin.
type FFTNode struct{}

func (n *FFTNode) Type() string { return "fft" }
func (n *FFTNode) Process(in []float64, st map[string]any) (float64, error) {
	s, err := n.ProcessSample(scalarInputs(in), st)
	if err != nil {
		return 0, err
	}
	if len(s.Arr) == 0 {
		return 0, nil
	}
	return s.Arr[0], nil
}
func (n *FFTNode) ProcessSample(in []Sample, _ map[string]any) (Sample, error) {
	if len(in) == 0 {
		return Sample{}, errors.New("fft: no inputs")
	}
	return Array(fftMagnitude(in[0].AsArray())), nil
}

// fftMagnitude returns the magnitude spectrum of x, zero-padded to the next
// power of two. Returns an empty slice for empty input.
func fftMagnitude(x []float64) []float64 {
	if len(x) == 0 {
		return nil
	}
	n := nextPow2(len(x))
	re := make([]float64, n)
	im := make([]float64, n)
	copy(re, x)
	fftRadix2(re, im)
	mag := make([]float64, n)
	for i := range mag {
		mag[i] = math.Hypot(re[i], im[i])
	}
	return mag
}