package nn

import (
	"math"

	"makarna/pkg/backend/cpu"
)

// Sigmoid applies the sigmoid function: 1 / (1 + exp(-x))
// Uses numerically stable computation for both positive and negative inputs.
func Sigmoid(x float32) float32 {
	if x >= 0 {
		z := float32(math.Exp(float64(-x)))
		return 1 / (1 + z)
	}
	z := float32(math.Exp(float64(x)))
	return z / (1 + z)
}

// SigmoidInplace applies sigmoid to each element of the slice in-place.
func SigmoidInplace(data []float32) {
	for i := range data {
		data[i] = Sigmoid(data[i])
	}
}

// SigmoidTensor applies sigmoid in-place to a tensor.
func SigmoidTensor(x *cpu.Tensor) error {
	SigmoidInplace(x.DataFloat32())
	return nil
}

// Softplus computes log(1 + exp(x)), with numerical stability for large x.
func Softplus(x float32) float32 {
	if x > 20 {
		return x
	}
	return float32(math.Log1p(math.Exp(float64(x))))
}

// SoftplusInplace applies softplus to each element of the slice in-place.
func SoftplusInplace(data []float32) {
	for i := range data {
		data[i] = Softplus(data[i])
	}
}

// SoftplusTensor applies softplus in-place to a tensor.
func SoftplusTensor(x *cpu.Tensor) error {
	SoftplusInplace(x.DataFloat32())
	return nil
}

// L2NormInplace normalizes a vector to unit length in-place.
// eps is added to prevent division by zero.
func L2NormInplace(x []float32, eps float32) {
	if len(x) == 0 {
		return
	}
	ss := float32(0)
	for _, v := range x {
		ss += v * v
	}
	inv := float32(1.0 / math.Sqrt(float64(ss+eps)))
	for i := range x {
		x[i] *= inv
	}
}

// L2NormHeads normalizes Q and K vectors per-head for multi-head attention.
// Shape: [tokens, numHeads * headDim], normalizes each [headDim] segment.
func L2NormHeads(q, k []float32, tokens, numHeads, headDim int, eps float32) {
	stride := numHeads * headDim
	for t := 0; t < tokens; t++ {
		base := t * stride
		for h := 0; h < numHeads; h++ {
			off := base + h*headDim
			L2NormInplace(q[off:off+headDim], eps)
			L2NormInplace(k[off:off+headDim], eps)
		}
	}
}

// Exp computes e^x
func Exp(x float32) float32 {
	return float32(math.Exp(float64(x)))
}

// ExpInplace applies exp to each element of the slice in-place.
func ExpInplace(data []float32) {
	for i := range data {
		data[i] = Exp(data[i])
	}
}