package nn

import (
	"math"
	"testing"

	"github.com/ollama/ollama/x/mlxrunner/batch"
	"github.com/ollama/ollama/x/mlxrunner/mlx"
)

func ones(dtype mlx.DType, shape ...int) *mlx.Array {
	return mlx.AddScalar(mlx.Zeros(dtype, shape...), 1)
}

// lastState returns the forward-end state — the last boundary the recurrent
// wrappers return.
func lastState(states []*mlx.Array) *mlx.Array { return states[len(states)-1] }

// fromValues builds a tensor with sequentially-numbered float32
// values so element-by-element parity actually exercises the kernel.
func fromValues(seed float32, shape ...int) *mlx.Array {
	n := 1
	for _, d := range shape {
		n *= d
	}
	vals := make([]float32, n)
	for i := range vals {
		vals[i] = seed + 0.1*float32(i)
	}
	return mlx.FromValues(vals, shape...)
}

// convFromKernel builds the depthwise causal Conv1d the model constructs at
// load time from a bare [C, K] kernel, so the wrapper tests drive the same
// mlx.Conv1d path production runs.
func convFromKernel(w *mlx.Array) *Conv1d {
	return NewConv1d(mlx.ExpandDims(w, 2), nil, 1, 0, 1, int32(w.Dim(0)))
}

// TestCausalConv1DPaddedRowParity drives a B=2 batch with one short
// row (qLen<L). For the short row, (a) `out` positions [0..qLen)
// must equal a B=1 reference at length qLen, (b) `nextConv` for the
// short row must be the row's last convTail real positions (not the
// padded tail), (c) the full row must be unaffected.
func TestCausalConv1DPaddedRowParity(t *testing.T) {
	skipIfNoMLX(t)
	L, D, convTail := 4, 3, 2
	qLenShort := 2
	K := convTail + 1

	weight := fromValues(0.2, D, K)
	conv := convFromKernel(weight)
	priorFull := fromValues(0.5, 2, convTail, D)
	priorShort := mlx.SliceStartStop(priorFull,
		[]int32{1, 0, 0},
		[]int32{2, int32(convTail), int32(D)})

	// Pad row 1 with arbitrary values past qLenShort — the wrapper
	// must zero them before convolving. Distinct values let us catch
	// any leak.
	inputFull := fromValues(1.0, 1, L, D)
	inputShortReal := mlx.FromValues([]float32{
		2.0, 2.1, 2.2,
		2.3, 2.4, 2.5,
	}, 1, qLenShort, D)
	inputShortPad := mlx.FromValues([]float32{
		99, 99, 99,
		99, 99, 99,
	}, 1, L-qLenShort, D)
	inputShortFull := mlx.Concatenate([]*mlx.Array{inputShortReal, inputShortPad}, 1)
	input := mlx.Concatenate([]*mlx.Array{inputFull, inputShortFull}, 0)

	b := &batch.Batch{
		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 2, L),
		SeqOffsets:   []int32{0, 0},
		SeqQueryLens: []int32{int32(L), int32(qLenShort)},
	}

	out, convStates := CausalConv1D(b, input, conv, convTail, WithRecurrentState(priorFull, nil))
	nextConv := lastState(convStates)
	mlx.Eval(out, nextConv)

	// Reference for row 0: B=1 unpadded length-L call.
	refOut0, refConvStates0 := CausalConv1D(&batch.Batch{},
		inputFull, conv, convTail,
		WithRecurrentState(mlx.SliceStartStop(priorFull,
			[]int32{0, 0, 0},
			[]int32{1, int32(convTail), int32(D)}), nil))
	refNextConv0 := lastState(refConvStates0)
	// Reference for row 1: B=1 unpadded length-qLenShort call.
	refOut1, refConvStates1 := CausalConv1D(&batch.Batch{},
		inputShortReal, conv, convTail,
		WithRecurrentState(priorShort, nil))
	refNextConv1 := lastState(refConvStates1)
	mlx.Eval(refOut0, refNextConv0, refOut1, refNextConv1)

	gotOut := out.Floats()
	wantOut0 := refOut0.Floats()
	wantOut1 := refOut1.Floats()

	for q := range L {
		for d := range D {
			i := q*D + d
			if gotOut[i] != wantOut0[i] {
				t.Fatalf("row 0 out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[i], wantOut0[i])
			}
		}
	}
	for q := range qLenShort {
		for d := range D {
			gotI := L*D + q*D + d
			refI := q*D + d
			if math.Abs(float64(gotOut[gotI]-wantOut1[refI])) > 1e-5 {
				t.Fatalf("row 1 real out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[gotI], wantOut1[refI])
			}
		}
	}

	// nextConv: row 0 unaffected, row 1 must be the row's real tail
	// (positions [qLenShort - convTail, qLenShort) of the per-row
	// concat, i.e. the last two real input rows in this setup).
	gotTail := nextConv.Floats()
	wantTail0 := refNextConv0.Floats()
	wantTail1 := refNextConv1.Floats()
	for k := range convTail {
		for d := range D {
			i := k*D + d
			if gotTail[i] != wantTail0[i] {
				t.Fatalf("row 0 nextConv[k=%d,d=%d]: got %v, want %v", k, d, gotTail[i], wantTail0[i])
			}
		}
	}
	for k := range convTail {
		for d := range D {
			gotI := convTail*D + k*D + d
			refI := k*D + d
			if gotTail[gotI] != wantTail1[refI] {
				t.Fatalf("row 1 nextConv[k=%d,d=%d]: got %v, want %v (must come from real positions, not the padded tail)",
					k, d, gotTail[gotI], wantTail1[refI])
			}
		}
	}
}

// TestGatedDeltaDelegatesToKernel checks the wrapper produces the same output
// and final state as a direct mlx.FastGatedDelta call, for both a zero prior
// state and a non-zero one (so the wrapper is shown to thread the prior through,
// not just handle the zero path).
func TestGatedDeltaDelegatesToKernel(t *testing.T) {
	skipIfNoMLX(t)
	B, L, nK, nV, dK, dV := 1, 2, 1, 1, 4, 4
	q := ones(mlx.DTypeFloat32, B, L, nK, dK)
	k := ones(mlx.DTypeFloat32, B, L, nK, dK)
	v := ones(mlx.DTypeFloat32, B, L, nV, dV)
	gDecay := ones(mlx.DTypeFloat32, B, L, nV)
	beta := ones(mlx.DTypeFloat32, B, L, nV)

	cases := []struct {
		name  string
		prior *mlx.Array
	}{
		{"zero", mlx.Zeros(mlx.DTypeFloat32, B, nV, dV, dK)},
		{"non-zero", mlx.MulScalar(ones(mlx.DTypeFloat32, B, nV, dV, dK), 3)},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			outA, statesA := GatedDelta(&batch.Batch{}, q, k, v, gDecay, beta, WithRecurrentState(nil, tc.prior))
			stateA := lastState(statesA)
			outB, stateB := mlx.FastGatedDelta(q, k, v, gDecay, beta, tc.prior, nil)
			mlx.Eval(outA, stateA, outB, stateB)

			gotOut, wantOut := outA.Floats(), outB.Floats()
			for i := range wantOut {
				if gotOut[i] != wantOut[i] {
					t.Fatalf("output[%d]: wrapper=%v direct=%v", i, gotOut[i], wantOut[i])
				}
			}
			gotState, wantState := stateA.Floats(), stateB.Floats()
			for i := range wantState {
				if gotState[i] != wantState[i] {
					t.Fatalf("state[%d]: wrapper=%v direct=%v", i, gotState[i], wantState[i])
				}
			}
		})
	}
}

// TestGatedDeltaPaddedRowParity drives a B=2 batch where row 1 is
// short (qLen < L). The wrapper must substitute neutral values
// (q=k=v=beta=0, g=1) at row 1's padded positions so the recurrence
// is a no-op there — and row 1's final state must equal the state
// after its last real token. Pinned via parity against a B=1 length-
// qLen call on the same row.
func TestGatedDeltaPaddedRowParity(t *testing.T) {
	skipIfNoMLX(t)
	L, nK, nV, dK, dV := 4, 1, 1, 4, 4
	qLenShort := 2

	makeRows := func(seedA, seedB float32, shape ...int) *mlx.Array {
		// Build a rank-(len(shape)+1) tensor with B=2 rows from two
		// distinct seeds so the rows are not accidentally identical.
		n := 1
		for _, d := range shape {
			n *= d
		}
		vals := make([]float32, 2*n)
		for i := range n {
			vals[i] = seedA + 0.1*float32(i)
		}
		for i := range n {
			vals[n+i] = seedB + 0.1*float32(i)
		}
		full := append([]int{2}, shape...)
		return mlx.FromValues(vals, full...)
	}

	q := makeRows(0.5, -0.5, L, nK, dK)
	k := makeRows(0.7, -0.7, L, nK, dK)
	v := makeRows(0.3, -0.3, L, nV, dV)
	gDecay := makeRows(0.1, -0.1, L, nV)
	beta := makeRows(0.4, -0.4, L, nV)
	priorState := makeRows(0.2, -0.2, nV, dV, dK)

	b := &batch.Batch{
		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 2, L),
		SeqOffsets:   []int32{0, 0},
		SeqQueryLens: []int32{int32(L), int32(qLenShort)},
	}
	_, states := GatedDelta(b, q, k, v, gDecay, beta, WithRecurrentState(nil, priorState))
	state := lastState(states)
	mlx.Eval(state)

	// Reference for row 1: B=1 length-qLenShort call against the
	// row's real prefix and its prior state slice.
	row1Slice := func(a *mlx.Array, axisLens ...int32) *mlx.Array {
		dims := a.Dims()
		start := make([]int32, len(dims))
		stop := make([]int32, len(dims))
		start[0], stop[0] = 1, 2
		for i := 1; i < len(dims); i++ {
			stop[i] = int32(dims[i])
		}
		// Optionally truncate axis 1 (sequence axis) to qLenShort.
		if len(axisLens) >= 1 && len(dims) >= 2 {
			stop[1] = axisLens[0]
		}
		return mlx.SliceStartStop(a, start, stop)
	}
	q1 := row1Slice(q, int32(qLenShort))
	k1 := row1Slice(k, int32(qLenShort))
	v1 := row1Slice(v, int32(qLenShort))
	gDecay1 := row1Slice(gDecay, int32(qLenShort))
	beta1 := row1Slice(beta, int32(qLenShort))
	priorRow1 := row1Slice(priorState)

	_, refState := mlx.FastGatedDelta(q1, k1, v1, gDecay1, beta1, priorRow1, nil)
	mlx.Eval(refState)

	gotState := state.Floats()
	wantState := refState.Floats()
	row1Stride := nV * dV * dK
	for i := range row1Stride {
		gotV := gotState[row1Stride+i]
		wantV := wantState[i]
		if math.Abs(float64(gotV-wantV)) > 1e-4 {
			t.Fatalf("row 1 final state[%d]: got %v, want %v", i, gotV, wantV)
		}
	}
}

// floatsClose compares two flat float slices within tolerance.
func floatsClose(t *testing.T, label string, got, want []float32, tol float64) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("%s: len %d, want %d", label, len(got), len(want))
	}
	for i := range want {
		if math.Abs(float64(got[i]-want[i])) > tol {
			t.Fatalf("%s[%d]: got %v, want %v", label, i, got[i], want[i])
		}
	}
}

// TestGatedDeltaSegmentEquivalence checks that running the scan in segments
// cut at WithSnapshotSplits offsets yields identical output and final state to
// the single-shot scan, and that each boundary state equals the single-shot
// state over the corresponding prefix.
func TestGatedDeltaSegmentEquivalence(t *testing.T) {
	skipIfNoMLX(t)
	B, T, Hk, Dk, Hv, Dv := 1, 4, 1, 32, 1, 32

	q := fromValues(0.1, B, T, Hk, Dk)
	k := fromValues(-0.2, B, T, Hk, Dk)
	v := fromValues(0.3, B, T, Hv, Dv)
	gDecay := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.9)
	beta := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.5)
	prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)

	full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(T)}}

	refOut, refStates := GatedDelta(full, q, k, v, gDecay, beta, WithRecurrentState(nil, prior))
	if len(refStates) != 1 {
		t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
	}

	segOut, segStates := GatedDelta(full, q, k, v, gDecay, beta,
		WithRecurrentState(nil, prior), WithSnapshotSplits([]int{1, 2, 3}))
	mlx.Eval(refOut, segOut)

	floatsClose(t, "out", segOut.Floats(), refOut.Floats(), 1e-4)
	// 3 splits + end = 4 boundary states.
	if len(segStates) != 4 {
		t.Fatalf("got %d boundary states, want 4", len(segStates))
	}
	mlx.Eval(lastState(segStates), lastState(refStates))
	floatsClose(t, "final state", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)

	// boundary i (offset i+1) must equal a single-shot scan over prefix [0, i+1).
	for i := range segStates {
		n := int32(i + 1)
		pb := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{n}}
		_, want := GatedDelta(pb,
			mlx.SliceStartStop(q, []int32{0, 0, 0, 0}, []int32{int32(B), n, int32(Hk), int32(Dk)}),
			mlx.SliceStartStop(k, []int32{0, 0, 0, 0}, []int32{int32(B), n, int32(Hk), int32(Dk)}),
			mlx.SliceStartStop(v, []int32{0, 0, 0, 0}, []int32{int32(B), n, int32(Hv), int32(Dv)}),
			mlx.SliceStartStop(gDecay, []int32{0, 0, 0}, []int32{int32(B), n, int32(Hv)}),
			mlx.SliceStartStop(beta, []int32{0, 0, 0}, []int32{int32(B), n, int32(Hv)}),
			WithRecurrentState(nil, prior))
		mlx.Eval(segStates[i], lastState(want))
		floatsClose(t, "boundary delta", segStates[i].Floats(), lastState(want).Floats(), 1e-4)
	}
}

// TestCausalConv1DSegmentEquivalence checks the conv segmented path matches the
// single-shot conv for output, final conv tail, and each boundary conv state.
func TestCausalConv1DSegmentEquivalence(t *testing.T) {
	skipIfNoMLX(t)
	B, L, D, convTail := 1, 4, 3, 2
	K := convTail + 1

	input := fromValues(0.5, B, L, D)
	prior := fromValues(-0.3, B, convTail, D)
	weight := fromValues(0.2, D, K)
	conv := convFromKernel(weight)

	full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(L)}}

	refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
	if len(refStates) != 1 {
		t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
	}

	segOut, segStates := CausalConv1D(full, input, conv, convTail,
		WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
	mlx.Eval(refOut, segOut)

	floatsClose(t, "conv out", segOut.Floats(), refOut.Floats(), 1e-4)
	if len(segStates) != 4 {
		t.Fatalf("got %d boundary conv states, want 4", len(segStates))
	}
	mlx.Eval(lastState(segStates), lastState(refStates))
	floatsClose(t, "conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)

	for i := range segStates {
		n := int32(i + 1)
		pb := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{n}}
		_, want := CausalConv1D(pb,
			mlx.SliceStartStop(input, []int32{0, 0, 0}, []int32{int32(B), n, int32(D)}),
			conv, convTail, WithRecurrentState(prior, nil))
		mlx.Eval(segStates[i], lastState(want))
		floatsClose(t, "boundary conv", segStates[i].Floats(), lastState(want).Floats(), 1e-4)
	}
}

// TestGatedDeltaSegmentEquivalenceBatched checks the segmented scan matches the
// single-shot scan for B>1, including a ragged batch where rows have different
// real lengths — the per-segment sliced mask must zero each row's padded
// positions so a short row's boundary state freezes at its real end.
func TestGatedDeltaSegmentEquivalenceBatched(t *testing.T) {
	skipIfNoMLX(t)
	B, T, Hk, Dk, Hv, Dv := 2, 4, 1, 32, 1, 32

	q := fromValues(0.1, B, T, Hk, Dk)
	k := fromValues(-0.2, B, T, Hk, Dk)
	v := fromValues(0.3, B, T, Hv, Dv)
	gDecay := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.9)
	beta := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.5)
	prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)

	// Row 0 full length T; row 1 ends at 3 (so segment [3,4) is all padding
	// for row 1).
	full := &batch.Batch{SeqOffsets: []int32{0, 0}, SeqQueryLens: []int32{int32(T), 3}}

	refOut, refStates := GatedDelta(full, q, k, v, gDecay, beta, WithRecurrentState(nil, prior))
	segOut, segStates := GatedDelta(full, q, k, v, gDecay, beta,
		WithRecurrentState(nil, prior), WithSnapshotSplits([]int{1, 2, 3}))
	mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))

	floatsClose(t, "batched out", segOut.Floats(), refOut.Floats(), 1e-4)
	floatsClose(t, "batched final state", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
	if len(segStates) != 4 {
		t.Fatalf("got %d boundary states, want 4", len(segStates))
	}

	// Each row's boundary i (offset i+1) must equal a B=1 single-shot scan over
	// that row's real prefix: row 0 advances the full length, row 1 freezes once
	// it reaches its real length 3. Per-row B=1 references avoid the ambiguity of
	// re-declaring a ragged length over a uniform input slice.
	rowReal := []int32{int32(T), 3}
	for i := range segStates {
		for r := range B {
			n := min(int32(i+1), rowReal[r])
			lo, hi := int32(r), int32(r)+1
			rowPrior := mlx.SliceStartStop(prior, []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
			_, want := GatedDelta(&batch.Batch{},
				mlx.SliceStartStop(q, []int32{lo, 0, 0, 0}, []int32{hi, n, int32(Hk), int32(Dk)}),
				mlx.SliceStartStop(k, []int32{lo, 0, 0, 0}, []int32{hi, n, int32(Hk), int32(Dk)}),
				mlx.SliceStartStop(v, []int32{lo, 0, 0, 0}, []int32{hi, n, int32(Hv), int32(Dv)}),
				mlx.SliceStartStop(gDecay, []int32{lo, 0, 0}, []int32{hi, n, int32(Hv)}),
				mlx.SliceStartStop(beta, []int32{lo, 0, 0}, []int32{hi, n, int32(Hv)}),
				WithRecurrentState(nil, rowPrior))
			gotRow := mlx.SliceStartStop(segStates[i], []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
			mlx.Eval(gotRow, lastState(want))
			floatsClose(t, "batched boundary delta", gotRow.Floats(), lastState(want).Floats(), 1e-4)
		}
	}
}

// TestCausalConv1DSegmentEquivalenceBatched is the conv analog of the gated-delta
// batched test: boundary tails from the single conv pass vs per-row single-shot
// references for a ragged B>1 batch, where a short row must freeze its tail at
// its real end rather than reach into padding.
func TestCausalConv1DSegmentEquivalenceBatched(t *testing.T) {
	skipIfNoMLX(t)
	B, L, D, convTail := 2, 4, 3, 2
	K := convTail + 1

	input := fromValues(0.5, B, L, D)
	prior := fromValues(-0.3, B, convTail, D)
	weight := fromValues(0.2, D, K)
	conv := convFromKernel(weight)

	full := &batch.Batch{SeqOffsets: []int32{0, 0}, SeqQueryLens: []int32{int32(L), 3}}

	refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
	segOut, segStates := CausalConv1D(full, input, conv, convTail,
		WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
	mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))

	floatsClose(t, "batched conv out", segOut.Floats(), refOut.Floats(), 1e-4)
	floatsClose(t, "batched conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
	if len(segStates) != 4 {
		t.Fatalf("got %d boundary conv states, want 4", len(segStates))
	}

	// Each row's boundary i (offset i+1) must equal a B=1 single-shot conv over
	// that row's real prefix: row 0 advances the full length, row 1 freezes once
	// it reaches its real length 3. Per-row B=1 references avoid the ambiguity of
	// re-declaring a ragged length over a uniform input slice.
	rowReal := []int32{int32(L), 3}
	for i := range segStates {
		for r := range B {
			n := min(int32(i+1), rowReal[r])
			rowPrior := mlx.SliceStartStop(prior,
				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
			rowInput := mlx.SliceStartStop(input,
				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, n, int32(D)})
			_, want := CausalConv1D(&batch.Batch{}, rowInput, conv, convTail,
				WithRecurrentState(rowPrior, nil))
			gotRow := mlx.SliceStartStop(segStates[i],
				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
			mlx.Eval(gotRow, lastState(want))
			floatsClose(t, "batched boundary conv", gotRow.Floats(), lastState(want).Floats(), 1e-4)
		}
	}
}
