package gotensor

import (
	"git.kingecg.top/kingecg/gomatrix"
	"math"
)

// Flatten 将多维张量展平为一维
func (t *Tensor) Flatten() *Tensor {
	totalSize := t.Size()
	flatShape := []int{totalSize}
	
	// 创建新数据切片
	flatData := make([]float64, totalSize)
	
	// 将原数据复制到扁平化数组
	shape := t.Data.Shape()
	if len(shape) == 2 {
		rows, cols := shape[0], shape[1]
		for i := 0; i < rows; i++ {
			for j := 0; j < cols; j++ {
				val, _ := t.Data.Get(i, j)
				flatData[i*cols+j] = val
			}
		}
	} else if len(shape) == 4 {
		// 假设为 [batch, channel, height, width] 格式
		b, c, h, w := shape[0], shape[1], shape[2], shape[3]
		for bi := 0; bi < b; bi++ {
			for ci := 0; ci < c; ci++ {
				for hi := 0; hi < h; hi++ {
					for wi := 0; wi < w; wi++ {
						srcIdx := ((bi*c+ci)*h+hi)*w + wi
						val, _ := t.Data.Get(bi, ci, hi, wi)
						flatData[srcIdx] = val
					}
				}
			}
		}
	} else {
		// 对于其他情况，直接复制数据
		for i := 0; i < totalSize; i++ {
			val, _ := t.Data.Get(i)
			flatData[i] = val
		}
	}

	resultMatrix, _ := gomatrix.NewMatrix(flatData, flatShape)
	
	output := &Tensor{
		Data: resultMatrix,
		Op:   "flatten",
	}
	
	output.Prevs[0] = t
	output.Num_Prevs = 1
	
	output.backwardFunc = func() {
		if t.Grad != nil {
			// 将梯度重塑回原始形状
			originalShape := t.Data.Shape()
			originalSize := t.Size()
			
			// 创建一个临时数组来保存重塑后的梯度
			reshapedGradData := make([]float64, originalSize)
			
			// 根据原始形状复制梯度数据
			if len(originalShape) == 2 {
				rows, cols := originalShape[0], originalShape[1]
				for i := 0; i < rows; i++ {
					for j := 0; j < cols; j++ {
						srcIdx := i*cols + j
						gradVal, _ := output.Grad.Get(srcIdx)
						reshapedGradData[srcIdx] = gradVal
					}
				}
			} else if len(originalShape) == 4 {
				b, c, h, w := originalShape[0], originalShape[1], originalShape[2], originalShape[3]
				for bi := 0; bi < b; bi++ {
					for ci := 0; ci < c; ci++ {
						for hi := 0; hi < h; hi++ {
							for wi := 0; wi < w; wi++ {
								srcIdx := ((bi*c+ci)*h+hi)*w + wi
								gradVal, _ := output.Grad.Get(srcIdx)
								reshapedGradData[srcIdx] = gradVal
							}
						}
					}
				}
			} else {
				for i := 0; i < originalSize; i++ {
					gradVal, _ := output.Grad.Get(i)
					reshapedGradData[i] = gradVal
				}
			}
			
			// 创建重塑后的梯度矩阵
			reshapedGradMatrix, _ := gomatrix.NewMatrix(reshapedGradData, originalShape)
			
			// 将重塑后的梯度加到原张量的梯度上
			newGrad, _ := t.Grad.Add(reshapedGradMatrix)
			t.Grad = newGrad
		}
	}
	
	return output
}

// CrossEntropy 交叉熵损失函数
func (t *Tensor) CrossEntropy(target *Tensor) *Tensor {
	// 计算预测值的概率分布
	predictions := t.Softmax()
	
	// 计算交叉熵损失
	// loss = -sum(target * log(predictions + epsilon))
	epsilon := 1e-15 // 防止log(0)
	
	predData := predictions.Data
	targetData := target.Data
	
	// 计算log(predictions + epsilon)
	logPredData := make([]float64, predData.Size())
	for i := 0; i < predData.Size(); i++ {
		var val float64
		var err error
		
		shape := predData.Shape()
		if len(shape) == 1 {
			val, err = predData.Get(i)
		} else if len(shape) == 2 {
			_, cols := shape[0], shape[1]
			val, err = predData.Get(i/cols, i%cols)
		} else {
			// 对于其他情况，简单地按索引获取
			val, err = predData.Get(i)
		}
		
		if err != nil {
			continue
		}
		
		logPredData[i] = math.Log(val + epsilon)
	}
	logPredMatrix, _ := gomatrix.NewMatrix(logPredData, predData.Shape())
	
	// 计算target * log(predictions + epsilon)
	multiplied, _ := targetData.Multiply(logPredMatrix)
	
	// 求和得到损失值
	var totalLoss float64
	for i := 0; i < multiplied.Size(); i++ {
		var val float64
		var err error
		
		shape := multiplied.Shape()
		if len(shape) == 1 {
			val, err = multiplied.Get(i)
		} else if len(shape) == 2 {
			_, cols := shape[0], shape[1]
			val, err = multiplied.Get(i/cols, i%cols)
		} else {
			// 对于其他情况，简单地按索引获取
			val, err = multiplied.Get(i)
		}
		
		if err != nil {
			continue
		}
		
		totalLoss += val
	}
	totalLoss = -totalLoss
	
	// 创建损失张量（标量）
	lossValue := []float64{totalLoss}
	lossMatrix, _ := gomatrix.NewMatrix(lossValue, []int{1})
	
	output := &Tensor{
		Data: lossMatrix,
		Op:   "cross_entropy",
	}
	
	output.Prevs[0] = t
	output.Num_Prevs = 1
	
	output.backwardFunc = func() {
		if t.Grad != nil {
			// 计算交叉熵损失相对于预测值的梯度
			predProbs := predictions.Data
			targetVals := target.Data
			
			gradData := make([]float64, predProbs.Size())
			for i := 0; i < predProbs.Size(); i++ {
				var predVal, targetVal float64
				var err error
				
				predShape := predProbs.Shape()
				if len(predShape) == 1 {
					predVal, err = predProbs.Get(i)
					if err != nil {
						continue
					}
					targetVal, err = targetVals.Get(i)
					if err != nil {
						continue
					}
				} else if len(predShape) == 2 {
					_, cols := predShape[0], predShape[1]
					predVal, err = predProbs.Get(i/cols, i%cols)
					if err != nil {
						continue
					}
					targetVal, err = targetVals.Get(i/cols, i%cols)
					if err != nil {
						continue
					}
				} else {
					predVal, err = predProbs.Get(i)
					if err != nil {
						continue
					}
					targetVal, err = targetVals.Get(i)
					if err != nil {
						continue
					}
				}
				
				// 交叉熵损失的梯度是 pred - target
				gradData[i] = (predVal - targetVal) / float64(predProbs.Size())
			}
			
			gradMatrix, _ := gomatrix.NewMatrix(gradData, predProbs.Shape())
			
			// 将梯度加到原张量的梯度上
			newGrad, _ := t.Grad.Add(gradMatrix)
			t.Grad = newGrad
		}
	}
	
	return output
}

// Softmax 激活函数
func (t *Tensor) Softmax() *Tensor {
	// 假设输入是二维的，第二维是类别维度
	shape := t.Data.Shape()
	
	if len(shape) != 2 {
		// 如果不是二维，则将其视为一维（单个样本）
		origSize := t.Size()
		shape = []int{1, origSize}
	}
	
	rows, cols := shape[0], shape[1]
	
	// 计算softmax，按行进行
	resultData := make([]float64, t.Size())
	
	for r := 0; r < rows; r++ {
		// 找到当前行的最大值，用于数值稳定性
		maxVal := -math.Inf(1)
		for c := 0; c < cols; c++ {
			var val float64
			var err error
			
			if len(t.Data.Shape()) == 2 {
				val, err = t.Data.Get(r, c)
			} else {
				// 如果原数据不是二维的，我们按顺序取值
				val, err = t.Data.Get(r*cols+c)
			}
			
			if err != nil {
				continue
			}
			
			if val > maxVal {
				maxVal = val
			}
		}
		
		// 计算指数和
		expSum := 0.0
		tempVals := make([]float64, cols)
		for c := 0; c < cols; c++ {
			var val float64
			var err error
			
			if len(t.Data.Shape()) == 2 {
				val, err = t.Data.Get(r, c)
			} else {
				// 如果原数据不是二维的，我们按顺序取值
				val, err = t.Data.Get(r*cols+c)
			}
			
			if err != nil {
				continue
			}
			
			expVal := math.Exp(val - maxVal) // 减去最大值提高数值稳定性
			tempVals[c] = expVal
			expSum += expVal
		}
		
		// 计算softmax值
		for c := 0; c < cols; c++ {
			resultData[r*cols+c] = tempVals[c] / expSum
		}
	}
	
	resultMatrix, _ := gomatrix.NewMatrix(resultData, shape)
	
	output := &Tensor{
		Data: resultMatrix,
		Op:   "softmax",
	}
	
	output.Prevs[0] = t
	output.Num_Prevs = 1
	
	output.backwardFunc = func() {
		if t.Grad != nil {
			// 计算softmax的梯度
			gradData := make([]float64, t.Size())
			
			for r := 0; r < rows; r++ {
				for c1 := 0; c1 < cols; c1++ {
					gradSum := 0.0
					for c2 := 0; c2 < cols; c2++ {
						var s1, s2 float64
						var err error
						
						if len(resultMatrix.Shape()) == 2 {
							s1, err = resultMatrix.Get(r, c1)
							if err != nil {
								continue
							}
							s2, err = resultMatrix.Get(r, c2)
							if err != nil {
								continue
							}
						} else {
							s1, err = resultMatrix.Get(r*cols+c1)
							if err != nil {
								continue
							}
							s2, err = resultMatrix.Get(r*cols+c2)
							if err != nil {
								continue
							}
						}
						
						var ds_din float64
						if c1 == c2 {
							ds_din = s1 * (1 - s2)  // softmax导数：si*(1-sj) if i=j
						} else {
							ds_din = -s1 * s2         // softmax导数：-si*sj if i!=j
						}
						
						var outGrad float64
						if len(output.Grad.Shape()) == 2 {
							outGrad, _ = output.Grad.Get(r, c2)
						} else {
							outGrad, _ = output.Grad.Get(r*cols+c2)
						}
						gradSum += ds_din * outGrad
					}
					gradData[r*cols+c1] = gradSum
				}
			}
			
			gradMatrix, _ := gomatrix.NewMatrix(gradData, shape)
			
			// 将梯度加到原张量的梯度上
			newGrad, _ := t.Grad.Add(gradMatrix)
			t.Grad = newGrad
		}
	}
	
	return output
}

// MeanSquaredError 均方误差损失函数
func (t *Tensor) MeanSquaredError(target *Tensor) *Tensor {
	// 计算预测值和真实值之间的差值
	diff, _ := t.Data.Subtract(target.Data)
	
	// 计算平方差
	squaredDiff, _ := diff.Multiply(diff)
	
	// 计算均值
	var total float64
	for i := 0; i < squaredDiff.Size(); i++ {
		var val float64
		var err error
		
		shape := squaredDiff.Shape()
		if len(shape) == 1 {
			val, err = squaredDiff.Get(i)
		} else if len(shape) == 2 {
			_, cols := shape[0], shape[1]
			val, err = squaredDiff.Get(i/cols, i%cols)
		} else {
			// 对于其他情况，简单地按索引获取
			val, err = squaredDiff.Get(i)
		}
		
		if err != nil {
			continue
		}
		
		total += val
	}
	
	mseVal := total / float64(squaredDiff.Size())
	mseData := []float64{mseVal}
	mseMatrix, _ := gomatrix.NewMatrix(mseData, []int{1})
	
	output := &Tensor{
		Data: mseMatrix,
		Op:   "mse",
	}
	
	output.Prevs[0] = t
	output.Num_Prevs = 1
	
	output.backwardFunc = func() {
		if t.Grad != nil {
			// MSE损失函数的梯度是 2*(prediction - target)/N
			gradData := make([]float64, t.Size())
			n := float64(t.Size())
			
			for i := 0; i < t.Size(); i++ {
				var predVal, targetVal float64
				var err error
				
				shape := t.Data.Shape()
				if len(shape) == 1 {
					predVal, err = t.Data.Get(i)
					if err != nil {
						continue
					}
					targetVal, err = target.Data.Get(i)
					if err != nil {
						continue
					}
				} else if len(shape) == 2 {
					_, cols := shape[0], shape[1]
					predVal, err = t.Data.Get(i/cols, i%cols)
					if err != nil {
						continue
					}
					targetVal, err = target.Data.Get(i/cols, i%cols)
					if err != nil {
						continue
					}
				} else {
					// 对于其他情况
					predVal, err = t.Data.Get(i)
					if err != nil {
						continue
					}
					targetVal, err = target.Data.Get(i)
					if err != nil {
						continue
					}
				}
				
				gradData[i] = 2 * (predVal - targetVal) / n
			}
			
			gradMatrix, _ := gomatrix.NewMatrix(gradData, t.Data.Shape())
			
			// 将梯度加到原张量的梯度上
			newGrad, _ := t.Grad.Add(gradMatrix)
			t.Grad = newGrad
		}
	}
	
	return output
}