package ops

import (
	"makarna/pkg/backend/cpu"
	"makarna/pkg/tensor"
)

// Slice extracts a portion of tensor along specified dimension
// dim: dimension to slice
// start, end: range [start, end)
func Slice(t *cpu.Tensor, dim, start, end int) *cpu.Tensor {
	shape := t.Shape()
	data := t.DataFloat32()
	
	newShape := make(tensor.Shape, len(shape))
	copy(newShape, shape)
	newShape[dim] = end - start
	
	// Calculate strides
	strides := make([]int, len(shape))
	strides[len(shape)-1] = 1
	for i := len(shape) - 2; i >= 0; i-- {
		strides[i] = strides[i+1] * shape[i+1]
	}
	
	newData := make([]float32, newShape.NumElements())
	
	// Iterate and copy
	newIdx := 0
	indices := make([]int, len(shape))
	for i := 0; i < len(data); i++ {
		// Check if this index is within slice range for target dim
		if indices[dim] >= start && indices[dim] < end {
			newData[newIdx] = data[i]
			newIdx++
		}
		
		// Increment indices
		for d := len(indices) - 1; d >= 0; d-- {
			indices[d]++
			if indices[d] < shape[d] {
				break
			}
			indices[d] = 0
		}
	}
	
	return cpu.NewTensor(newShape, newData)
}

// Concat concatenates tensors along specified dimension
func Concat(tensors []*cpu.Tensor, dim int) *cpu.Tensor {
	if len(tensors) == 0 {
		return nil
	}
	if len(tensors) == 1 {
		return tensors[0]
	}
	
	// Calculate new shape
	refShape := tensors[0].Shape()
	newShape := make(tensor.Shape, len(refShape))
	copy(newShape, refShape)
	
	totalDim := 0
	for _, t := range tensors {
		totalDim += t.Shape()[dim]
	}
	newShape[dim] = totalDim
	
	// Simple case: concat along last dimension
	if dim == len(refShape)-1 {
		newData := make([]float32, 0, newShape.NumElements())
		// For each row, append all tensors' data
		numRows := refShape.NumElements() / refShape[dim]
		for row := 0; row < numRows; row++ {
			for _, t := range tensors {
				tData := t.DataFloat32()
				rowSize := t.Shape()[dim]
				start := row * rowSize
				newData = append(newData, tData[start:start+rowSize]...)
			}
		}
		return cpu.NewTensor(newShape, newData)
	}
	
	// General case: just concatenate flat data (works for dim=0)
	if dim == 0 {
		newData := make([]float32, 0, newShape.NumElements())
		for _, t := range tensors {
			newData = append(newData, t.DataFloat32()...)
		}
		return cpu.NewTensor(newShape, newData)
	}
	
	// General case: concat along middle dimension
	// Calculate outer (before dim), inner (after dim) sizes
	outerSize := 1
	for i := 0; i < dim; i++ {
		outerSize *= refShape[i]
	}
	innerSize := 1
	for i := dim + 1; i < len(refShape); i++ {
		innerSize *= refShape[i]
	}
	
	newData := make([]float32, 0, newShape.NumElements())
	
	// For each outer index, copy all tensors' slices
	for outer := 0; outer < outerSize; outer++ {
		for _, t := range tensors {
			tData := t.DataFloat32()
			dimSize := t.Shape()[dim]
			sliceSize := dimSize * innerSize
			start := outer * sliceSize
			newData = append(newData, tData[start:start+sliceSize]...)
		}
	}
	
	return cpu.NewTensor(newShape, newData)
}