package gotensor

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

const (
	Max_Prevs = 10
	Max_Args  = 10
)

type Arg interface {
	int | float64 | []int
}

type Tensor struct {
	Data      *gomatrix.Matrix
	Grad      *gomatrix.Matrix
	Prevs     [Max_Prevs]*Tensor
	Num_Prevs int
	Args      [Max_Args]any // must meet Arg
	Op        string       // 操作类型，用于反向传播
	backwardFunc  func()       // 反向传播函数
}

// NewTensor 创建一个新的Tensor
func NewTensor(data []float64, shape []int) (*Tensor, error) {
	matrix, err := gomatrix.NewMatrix(data, shape)
	if err != nil {
		return nil, err
	}

	// 创建梯度矩阵，初始为零
	grad, err := gomatrix.NewZeros(shape)
	if err != nil {
		return nil, err
	}

	return &Tensor{
		Data: matrix,
		Grad: grad,
		Op:   "",
	}, nil
}

// NewZeros 创建一个全零Tensor
func NewZeros(shape []int) (*Tensor, error) {
	matrix, err := gomatrix.NewZeros(shape)
	if err != nil {
		return nil, err
	}

	grad, err := gomatrix.NewZeros(shape)
	if err != nil {
		return nil, err
	}

	return &Tensor{
		Data: matrix,
		Grad: grad,
		Op:   "zeros",
	}, nil
}

// NewOnes 创建一个全一Tensor
func NewOnes(shape []int) (*Tensor, error) {
	matrix, err := gomatrix.NewOnes(shape)
	if err != nil {
		return nil, err
	}

	grad, err := gomatrix.NewZeros(shape)
	if err != nil {
		return nil, err
	}

	return &Tensor{
		Data: matrix,
		Grad: grad,
		Op:   "ones",
	}, nil
}

// NewIdentity 创建一个单位Tensor（矩阵）
func NewIdentity(size int) (*Tensor, error) {
	matrix, err := gomatrix.NewIdentity(size)
	if err != nil {
		return nil, err
	}

	grad, err := gomatrix.NewZeros([]int{size, size})
	if err != nil {
		return nil, err
	}

	return &Tensor{
		Data: matrix,
		Grad: grad,
		Op:   "identity",
	}, nil
}

// Add Tensor加法
func (t *Tensor) Add(other *Tensor) (*Tensor, error) {
	result, err := t.Data.Add(other.Data)
	if err != nil {
		return nil, err
	}

	// 创建结果Tensor
	output := &Tensor{
		Data: result,
		Op:   "add",
	}

	// 记录依赖关系
	output.Prevs[0] = t
	output.Prevs[1] = other
	output.Num_Prevs = 2

	// 定义反向传播函数
	output.backwardFunc = func() {
		if t.Grad != nil {
			// 对t的梯度等于输出梯度
			grad, _ := t.Grad.Add(output.Grad)
			t.Grad = grad
		}
		if other.Grad != nil {
			// 对other的梯度等于输出梯度
			grad, _ := other.Grad.Add(output.Grad)
			other.Grad = grad
		}
	}

	return output, nil
}

// Subtract Tensor减法
func (t *Tensor) Subtract(other *Tensor) (*Tensor, error) {
	result, err := t.Data.Subtract(other.Data)
	if err != nil {
		return nil, err
	}

	output := &Tensor{
		Data: result,
		Op:   "sub",
	}

	output.Prevs[0] = t
	output.Prevs[1] = other
	output.Num_Prevs = 2

	output.backwardFunc = func() {
		if t.Grad != nil {
			// 对t的梯度等于输出梯度
			grad, _ := t.Grad.Add(output.Grad)
			t.Grad = grad
		}
		if other.Grad != nil {
			// 对other的梯度等于输出梯度的负值
			negGrad := output.Grad.Scale(-1.0)
			grad, _ := other.Grad.Add(negGrad)
			other.Grad = grad
		}
	}

	return output, nil
}

// Multiply Tensor乘法（逐元素相乘）
func (t *Tensor) Multiply(other *Tensor) (*Tensor, error) {
	result, err := t.Data.Multiply(other.Data)
	if err != nil {
		return nil, err
	}

	output := &Tensor{
		Data: result,
		Op:   "mul",
	}

	output.Prevs[0] = t
	output.Prevs[1] = other
	output.Num_Prevs = 2

	output.backwardFunc = func() {
		if t.Grad != nil {
			// 对t的梯度 = output.Grad * other.Data
			gradWithOther, _ := output.Grad.Multiply(other.Data)
			grad, _ := t.Grad.Add(gradWithOther)
			t.Grad = grad
		}
		if other.Grad != nil {
			// 对other的梯度 = output.Grad * t.Data
			gradWithT, _ := output.Grad.Multiply(t.Data)
			grad, _ := other.Grad.Add(gradWithT)
			other.Grad = grad
		}
	}

	return output, nil
}

// MatMul Tensor矩阵乘法
func (t *Tensor) MatMul(other *Tensor) (*Tensor, error) {
	result, err := t.Data.MatMul(other.Data)
	if err != nil {
		return nil, err
	}

	output := &Tensor{
		Data: result,
		Op:   "matmul",
	}

	output.Prevs[0] = t
	output.Prevs[1] = other
	output.Num_Prevs = 2

	output.backwardFunc = func() {
		if t.Grad != nil {
			// 对t的梯度 = output.Grad * other.Data^T
			otherT, _ := other.Data.Transpose()
			gradResult, _ := output.Grad.MatMul(otherT)
			grad, _ := t.Grad.Add(gradResult)
			t.Grad = grad
		}
		if other.Grad != nil {
			// 对other的梯度 = t.Data^T * output.Grad
			tT, _ := t.Data.Transpose()
			gradResult, _ := tT.MatMul(output.Grad)
			grad, _ := other.Grad.Add(gradResult)
			other.Grad = grad
		}
	}

	return output, nil
}

// Scale Tensor数乘
func (t *Tensor) Scale(factor float64) *Tensor {
	result := t.Data.Scale(factor)

	output := &Tensor{
		Data: result,
		Op:   "scale",
	}

	output.Prevs[0] = t
	output.Num_Prevs = 1
	output.Args[0] = factor

	output.backwardFunc = func() {
		if t.Grad != nil {
			// 对t的梯度 = output.Grad * factor
			scaledGrad := output.Grad.Scale(factor)
			grad, _ := t.Grad.Add(scaledGrad)
			t.Grad = grad
		}
	}

	return output
}

// ZeroGrad 将梯度置零
func (t *Tensor) ZeroGrad() {
	if t.Grad != nil {
		shape := t.Grad.Shape()
		zeroGrad, _ := gomatrix.NewZeros(shape)
		t.Grad = zeroGrad
	}
}

// Shape 返回Tensor的形状
func (t *Tensor) Shape() []int {
	return t.Data.Shape()
}

// Size 返回Tensor的大小
func (t *Tensor) Size() int {
	return t.Data.Size()
}

// Get 获取指定位置的值
func (t *Tensor) Get(indices ...int) (float64, error) {
	return t.Data.Get(indices...)
}

// Set 设置指定位置的值
func (t *Tensor) Set(value float64, indices ...int) error {
	return t.Data.Set(value, indices...)
}

// String 返回Tensor的字符串表示
func (t *Tensor) String() string {
	return t.Data.String()
}

// Sigmoid 激活函数
func (t *Tensor) Sigmoid() *Tensor {
	// 计算每个元素的sigmoid值
	inputShape := t.Data.Shape()
	resultData := make([]float64, t.Data.Size())
	
	// 遍历所有元素计算sigmoid
	idx := 0
	if len(inputShape) == 1 {
		for i := 0; i < inputShape[0]; i++ {
			val, _ := t.Data.Get(i)
			resultData[idx] = sigmoid(val)
			idx++
		}
	} else if len(inputShape) == 2 {
		rows, cols := inputShape[0], inputShape[1]
		for i := 0; i < rows; i++ {
			for j := 0; j < cols; j++ {
				val, _ := t.Data.Get(i, j)
				resultData[idx] = sigmoid(val)
				idx++
			}
		}
	} else if len(inputShape) == 4 {
		batch, ch, h, w := inputShape[0], inputShape[1], inputShape[2], inputShape[3]
		for b := 0; b < batch; b++ {
			for c := 0; c < ch; c++ {
				for h_idx := 0; h_idx < h; h_idx++ {
					for w_idx := 0; w_idx < w; w_idx++ {
						val, _ := t.Data.Get(b, c, h_idx, w_idx)
						resultData[idx] = sigmoid(val)
						idx++
					}
				}
			}
		}
	}
	
	result, _ := gomatrix.NewMatrix(resultData, inputShape)

	output := &Tensor{
		Data: result,
		Op:   "sigmoid",
	}

	output.Prevs[0] = t
	output.Num_Prevs = 1

	output.backwardFunc = func() {
		if t.Grad != nil {
			// Sigmoid的导数: sigmoid(x) * (1 - sigmoid(x))
			gradData := make([]float64, t.Data.Size())
			gradIdx := 0
			
			if len(inputShape) == 1 {
				for i := 0; i < inputShape[0]; i++ {
					sigmoidVal, _ := result.Get(i)
					gradData[gradIdx] = sigmoidVal * (1 - sigmoidVal)
					gradIdx++
				}
			} else if len(inputShape) == 2 {
				rows, cols := inputShape[0], inputShape[1]
				for i := 0; i < rows; i++ {
					for j := 0; j < cols; j++ {
						sigmoidVal, _ := result.Get(i, j)
						gradData[gradIdx] = sigmoidVal * (1 - sigmoidVal)
						gradIdx++
					}
				}
			} else if len(inputShape) == 4 {
				batch, ch, h, w := inputShape[0], inputShape[1], inputShape[2], inputShape[3]
				for b := 0; b < batch; b++ {
					for c := 0; c < ch; c++ {
						for h_idx := 0; h_idx < h; h_idx++ {
							for w_idx := 0; w_idx < w; w_idx++ {
								sigmoidVal, _ := result.Get(b, c, h_idx, w_idx)
								gradData[gradIdx] = sigmoidVal * (1 - sigmoidVal)
								gradIdx++
							}
						}
					}
				}
			}
			
			gradMatrix, _ := gomatrix.NewMatrix(gradData, inputShape)
			
			// 与输出梯度相乘
			gradWithOutput, _ := gradMatrix.Multiply(output.Grad)
			grad, _ := t.Grad.Add(gradWithOutput)
			t.Grad = grad
		}
	}

	return output
}

// ReLU 激活函数
func (t *Tensor) ReLU() *Tensor {
	// 计算每个元素的ReLU值
	inputShape := t.Data.Shape()
	resultData := make([]float64, t.Data.Size())
	
	// 遍历所有元素计算ReLU
	idx := 0
	if len(inputShape) == 1 {
		for i := 0; i < inputShape[0]; i++ {
			val, _ := t.Data.Get(i)
			if val > 0 {
				resultData[idx] = val
			} else {
				resultData[idx] = 0
			}
			idx++
		}
	} else if len(inputShape) == 2 {
		rows, cols := inputShape[0], inputShape[1]
		for i := 0; i < rows; i++ {
			for j := 0; j < cols; j++ {
				val, _ := t.Data.Get(i, j)
				if val > 0 {
					resultData[idx] = val
				} else {
					resultData[idx] = 0
				}
				idx++
			}
		}
	} else if len(inputShape) == 4 {
		batch, ch, h, w := inputShape[0], inputShape[1], inputShape[2], inputShape[3]
		for b := 0; b < batch; b++ {
			for c := 0; c < ch; c++ {
				for h_idx := 0; h_idx < h; h_idx++ {
					for w_idx := 0; w_idx < w; w_idx++ {
						val, _ := t.Data.Get(b, c, h_idx, w_idx)
						if val > 0 {
							resultData[idx] = val
						} else {
							resultData[idx] = 0
						}
						idx++
					}
				}
			}
		}
	}
	
	result, _ := gomatrix.NewMatrix(resultData, inputShape)

	output := &Tensor{
		Data: result,
		Op:   "relu",
	}

	output.Prevs[0] = t
	output.Num_Prevs = 1

	output.backwardFunc = func() {
		if t.Grad != nil {
			// ReLU的导数: 输入大于0时为1，否则为0
			gradData := make([]float64, t.Data.Size())
			gradIdx := 0
			
			if len(inputShape) == 1 {
				for i := 0; i < inputShape[0]; i++ {
					val, _ := t.Data.Get(i)
					if val > 0 {
						gradData[gradIdx] = 1
					} else {
						gradData[gradIdx] = 0
					}
					gradIdx++
				}
			} else if len(inputShape) == 2 {
				rows, cols := inputShape[0], inputShape[1]
				for i := 0; i < rows; i++ {
					for j := 0; j < cols; j++ {
						val, _ := t.Data.Get(i, j)
						if val > 0 {
							gradData[gradIdx] = 1
						} else {
							gradData[gradIdx] = 0
						}
						gradIdx++
					}
				}
			} else if len(inputShape) == 4 {
				batch, ch, h, w := inputShape[0], inputShape[1], inputShape[2], inputShape[3]
				for b := 0; b < batch; b++ {
					for c := 0; c < ch; c++ {
						for h_idx := 0; h_idx < h; h_idx++ {
							for w_idx := 0; w_idx < w; w_idx++ {
								val, _ := t.Data.Get(b, c, h_idx, w_idx)
								if val > 0 {
									gradData[gradIdx] = 1
								} else {
									gradData[gradIdx] = 0
								}
								gradIdx++
							}
						}
					}
				}
			}
			
			gradMatrix, _ := gomatrix.NewMatrix(gradData, inputShape)
			
			// 与输出梯度相乘
			gradWithOutput, _ := gradMatrix.Multiply(output.Grad)
			grad, _ := t.Grad.Add(gradWithOutput)
			t.Grad = grad
		}
	}

	return output
}

// sigmoid 计算sigmoid函数值
func sigmoid(x float64) float64 {
	if x > 0 {
		z := math.Exp(-x)
		return 1 / (1 + z)
	} else {
		z := math.Exp(x)
		return z / (1 + z)
	}
}

// Backward 执行反向传播
func (t *Tensor) Backward() {
	if t.backwardFunc != nil {
		t.backwardFunc()
	}
}