package dsp

import (
	"math"
	"testing"
)

func TestSampleRoundTrip(t *testing.T) {
	s := Scalar(3.5)
	if s.IsArray {
		t.Error("Scalar should not be an array")
	}
	if s.Type() != ValScalar {
		t.Errorf("Scalar type: want ValScalar, got %v", s.Type())
	}
	if s.AsAny() != 3.5 {
		t.Errorf("Scalar AsAny: want 3.5, got %v", s.AsAny())
	}
	if got := s.AsArray(); len(got) != 1 || got[0] != 3.5 {
		t.Errorf("Scalar AsArray: want [3.5], got %v", got)
	}

	a := Array([]float64{1, 2, 3})
	if !a.IsArray {
		t.Error("Array should be an array")
	}
	if a.Type() != ValArray {
		t.Errorf("Array type: want ValArray, got %v", a.Type())
	}
	got, ok := a.AsAny().([]float64)
	if !ok || len(got) != 3 {
		t.Errorf("Array AsAny: want []float64 len 3, got %v", a.AsAny())
	}
}

func TestApplyElementwiseAllScalar(t *testing.T) {
	n := &AddNode{}
	out, err := ApplyElementwise(n, []Sample{Scalar(2), Scalar(3)}, map[string]any{})
	if err != nil {
		t.Fatal(err)
	}
	if out.IsArray || out.F != 5 {
		t.Errorf("all-scalar add: want scalar 5, got %v", out)
	}
}

func TestApplyElementwiseBroadcast(t *testing.T) {
	// array ⊕ scalar: scalar is a constant broadcast across the array.
	n := &AddNode{}
	out, err := ApplyElementwise(n, []Sample{Array([]float64{1, 2, 3}), Scalar(10)}, map[string]any{})
	if err != nil {
		t.Fatal(err)
	}
	if !out.IsArray {
		t.Fatalf("array+scalar: want array, got %v", out)
	}
	want := []float64{11, 12, 13}
	for i, v := range want {
		if out.Arr[i] != v {
			t.Errorf("array+scalar[%d]: want %v, got %v", i, v, out.Arr[i])
		}
	}
}

func TestApplyElementwiseArrayArray(t *testing.T) {
	n := &MultiplyNode{}
	out, err := ApplyElementwise(n, []Sample{Array([]float64{1, 2, 3}), Array([]float64{4, 5, 6})}, map[string]any{})
	if err != nil {
		t.Fatal(err)
	}
	want := []float64{4, 10, 18}
	for i, v := range want {
		if out.Arr[i] != v {
			t.Errorf("array*array[%d]: want %v, got %v", i, v, out.Arr[i])
		}
	}
}

func TestApplyElementwiseLengthMismatch(t *testing.T) {
	n := &AddNode{}
	_, err := ApplyElementwise(n, []Sample{Array([]float64{1, 2}), Array([]float64{1, 2, 3})}, map[string]any{})
	if err == nil {
		t.Error("expected length-mismatch error")
	}
}

func TestReductionNodes(t *testing.T) {
	arr := []Sample{Array([]float64{2, 4, 6, 8})}
	cases := []struct {
		name string
		node ArrayNode
		want float64
	}{
		{"sum", &SumNode{}, 20},
		{"mean", &MeanNode{}, 5},
		{"min", &MinNode{}, 2},
		{"max", &MaxNode{}, 8},
		{"length", &LengthNode{}, 4},
		{"index", &IndexNode{I: 2}, 6},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			out, err := tc.node.ProcessSample(arr, map[string]any{})
			if err != nil {
				t.Fatal(err)
			}
			if out.IsArray || out.F != tc.want {
				t.Errorf("%s: want scalar %v, got %v", tc.name, tc.want, out)
			}
		})
	}
}

func TestIndexNodeOutOfRange(t *testing.T) {
	n := &IndexNode{I: 9}
	_, err := n.ProcessSample([]Sample{Array([]float64{1, 2, 3})}, map[string]any{})
	if err == nil {
		t.Error("expected out-of-range error")
	}
}

func TestSliceNode(t *testing.T) {
	n := &SliceNode{Start: 1, End: 3}
	out, err := n.ProcessSample([]Sample{Array([]float64{10, 20, 30, 40})}, map[string]any{})
	if err != nil {
		t.Fatal(err)
	}
	want := []float64{20, 30}
	if len(out.Arr) != len(want) {
		t.Fatalf("slice: want len %d, got %d", len(want), len(out.Arr))
	}
	for i, v := range want {
		if out.Arr[i] != v {
			t.Errorf("slice[%d]: want %v, got %v", i, v, out.Arr[i])
		}
	}
}

func TestFFTMagnitude(t *testing.T) {
	// A constant signal has all energy in bin 0 (the DC term equals the sum).
	x := []float64{1, 1, 1, 1}
	mag := fftMagnitude(x)
	if len(mag) != 4 {
		t.Fatalf("fft len: want 4, got %d", len(mag))
	}
	if math.Abs(mag[0]-4) > 1e-9 {
		t.Errorf("fft DC bin: want 4, got %v", mag[0])
	}
	for k := 1; k < len(mag); k++ {
		if math.Abs(mag[k]) > 1e-9 {
			t.Errorf("fft bin %d: want ~0, got %v", k, mag[k])
		}
	}
}

func TestFFTSingleTone(t *testing.T) {
	// One full cycle of a cosine over 8 samples → energy in bins 1 and N-1.
	n := 8
	x := make([]float64, n)
	for i := range x {
		x[i] = math.Cos(2 * math.Pi * float64(i) / float64(n))
	}
	mag := fftMagnitude(x)
	if math.Abs(mag[1]-float64(n)/2) > 1e-6 {
		t.Errorf("fft tone bin 1: want %v, got %v", float64(n)/2, mag[1])
	}
	if math.Abs(mag[n-1]-float64(n)/2) > 1e-6 {
		t.Errorf("fft tone bin %d: want %v, got %v", n-1, float64(n)/2, mag[n-1])
	}
}

func TestOpOutputType(t *testing.T) {
	cases := []struct {
		op      string
		in      []ValType
		want    ValType
		wantErr bool
	}{
		// reductions → scalar regardless of input
		{"sum", []ValType{ValArray}, ValScalar, false},
		{"mean", []ValType{ValScalar}, ValScalar, false},
		{"index", []ValType{ValUnknown}, ValScalar, false},
		// array producers require array, yield array
		{"fft", []ValType{ValArray}, ValArray, false},
		{"slice", []ValType{ValUnknown}, ValArray, false},
		{"fft", []ValType{ValScalar}, ValUnknown, true},
		// scalar-only reject arrays
		{"moving_average", []ValType{ValScalar}, ValScalar, false},
		{"lua", []ValType{ValArray}, ValUnknown, true},
		{"rms", []ValType{ValUnknown}, ValScalar, false},
		// elementwise: array if any array, scalar if all scalar, else unknown
		{"add", []ValType{ValScalar, ValScalar}, ValScalar, false},
		{"add", []ValType{ValArray, ValScalar}, ValArray, false},
		{"gain", []ValType{ValUnknown}, ValUnknown, false},
		{"expr", []ValType{ValArray, ValScalar}, ValArray, false},
	}
	for _, tc := range cases {
		got, err := OpOutputType(tc.op, tc.in)
		if tc.wantErr {
			if err == nil {
				t.Errorf("OpOutputType(%q,%v): expected error", tc.op, tc.in)
			}
			continue
		}
		if err != nil {
			t.Errorf("OpOutputType(%q,%v): unexpected error %v", tc.op, tc.in, err)
			continue
		}
		if got != tc.want {
			t.Errorf("OpOutputType(%q,%v): want %v, got %v", tc.op, tc.in, tc.want, got)
		}
	}
}